telnetroulette/terminal_handler.go

package main

import (
	"github.com/BRNSystems/go-telnet"
	"github.com/reiver/go-oi"
	"strings"
)

type internalTerminalHandler struct {
	ShellCharacter  string
	WelcomeMessage  string
	CommandRegistry *commandRegistry
}

func newInternalTerminalHandler(shellCharacter, welcomeMessage string, commandRegistry *commandRegistry) *internalTerminalHandler {
	return &internalTerminalHandler{
		ShellCharacter:  shellCharacter,
		WelcomeMessage:  welcomeMessage,
		CommandRegistry: commandRegistry,
	}
}

func clearLine(w telnet.Writer) error {
	_, err := w.RawWrite([]byte{0x1B, 0x5B, 0x4B})
	return err
}

func (handler *internalTerminalHandler) ServeTELNET(ctx telnet.Context, w telnet.Writer, r telnet.Reader) {
	clientID := newClient(&ctx, &w, &r)

	defer removeClient(clientID)
	_, err := w.RawWrite([]byte{0xff, 0xfb, 0x03})
	if err != nil {
		return
	} //send char by char
	_, err = w.RawWrite([]byte{0xff, 0xfb, 0x01})
	if err != nil {
		return
	} //echo from server
	_, err = w.Write([]byte(strings.ReplaceAll(handler.WelcomeMessage, "\n", "\r\n") + "\r\n"))
	if err != nil {
		return
	}
	//replace every \n with \r\n in WelcomeMessage
	_, err = w.Write([]byte(handler.ShellCharacter))
	if err != nil {
		return
	}

	var lineBuffer []byte

	CSISequence := 0

	for {
		var bt [1]byte
		n, err := r.Read(bt[:])
		if err != nil {
			break
		}

		if n > 0 {
			b := bt[0]

			if b == 0x03 || b == 0x04 {
				return
			} else if b == 0x1B { //escape
				CSISequence = 1
			} else if CSISequence == 1 && b == 0x5B {
				CSISequence = 2
			} else if CSISequence == 2 && b == 'A' {
				if len(lineBuffer) == 0 && len(clients[clientID].History) > 0 {
					if clients[clientID].HistoryIndex < len(clients[clientID].History)-1 {
						clients[clientID].HistoryIndex++
					}
					lineBuffer = []byte(clients[clientID].History[clients[clientID].HistoryIndex])
					if clearLine(w) != nil {
						return
					}
					_, err := w.Write(lineBuffer)
					if err != nil {
						return
					}
				}
			} else if CSISequence == 2 && b == 'B' {
				if len(lineBuffer) == 0 && len(clients[clientID].History) > 0 {
					if clients[clientID].HistoryIndex > 0 {
						clients[clientID].HistoryIndex--
					}
					lineBuffer = []byte(clients[clientID].History[clients[clientID].HistoryIndex])
					if clearLine(w) != nil {
						return
					}
					_, err := w.Write(lineBuffer)
					if err != nil {
						return
					}
				}
			} else if b == 0x7F {
				if len(lineBuffer) > 0 {
					_, err := w.Write([]byte{0x08, 0x20, 0x08})
					if err != nil {
						return
					}
					lineBuffer = lineBuffer[:len(lineBuffer)-1]
				}
			} else if b != '\r' && b != '\n' && b != 0 {
				CSISequence = 0
				lineBuffer = append(lineBuffer, b)
				_, err := w.RawWrite([]byte{b})
				if err != nil {
					return
				}
			} else if b == '\r' {
				_, err := w.Write([]byte("\r\n"))
				if err != nil {
					return
				}
				command := string(lineBuffer)
				args := strings.Fields(command)
				var message strings.Builder

				if len(args) > 0 {
					cmd := args[0]
					internalCommand := false
					if cmd == "exit" {
						internalCommand = true
						return
					} else if cmd == "help" {
						internalCommand = true
						var commands strings.Builder
						for command, commandEntry := range handler.CommandRegistry.commands {
							commands.WriteString(command + " - " + commandEntry.description + "\r\n")
						}
						// Handle help command
						if len(clients) > 0 {
							clients[clientID].send("Just help yourself, " + clients[clientID].Username + "\r\nCommands:\r\n" + commands.String())
						} else {
							clients[clientID].send("No clients available" + "\r\nCommands:\r\n" + commands.String())
						}
					}
					handlerFunc, ok, clientID := handler.CommandRegistry.getCommandHandler(cmd, clientID)
					clients[clientID].HistoryIndex = 0
					if ok || internalCommand {
						if ok {
							message.WriteString(handlerFunc(args[1:], clientID))
						}
						clients[clientID].History = append(clients[clientID].History, command)
					} else {
						message.WriteString("Command not found")
					}
					message.WriteString("\n" + handler.ShellCharacter)
				} else {
					message.WriteString(handler.ShellCharacter)
				}
				_, err = oi.LongWrite(w, []byte(strings.ReplaceAll(message.String(), "\n", "\r\n")))
				if err != nil {
					return
				}
				lineBuffer = nil
			}
		}
	}
}