tramp-blog/shell.go

package main

import (
	"bufio"
	"errors"
	"io"
	"log"
	"os"
	"strings"
	"unsafe"

	"github.com/google/shlex"
	"github.com/pkg/term/termios"
	"golang.org/x/sys/unix"
)

type Shell struct {
	pty    *os.File
	tty    *os.File
	prompt []byte

	readData  []byte
	readIndex int
	winsize   *Winsize
}

type Command struct {
	Name      string
	Arguments []string
}

func (s *Shell) GetPty() *os.File {
	return s.pty
}

func (s *Shell) Attach(conn io.ReadWriteCloser) error {
	pty, tty, err := termios.Pty()
	if err != nil {
		return err
	}

	s.pty = pty
	s.tty = tty

	if s.winsize != nil {
		SetWinsize(s.pty.Fd(), s.winsize)
	}

	go func() {
		defer s.pty.Close()
		io.Copy(s.pty, conn)
	}()

	go func() {
		defer conn.Close()
		io.Copy(conn, s.pty)
	}()

	_, err = s.tty.Write(s.prompt)
	return err
}

func NewShell(prompt string) *Shell {
	return &Shell{
		prompt: []byte(prompt),
	}
}

func (s *Shell) SetWinsize(w, h uint32) {
	s.winsize = &Winsize{Width: uint16(w), Height: uint16(h)}

	if s.pty != nil {
		SetWinsize(s.pty.Fd(), s.winsize)
	}
}

func (s *Shell) Close() (error, error) {
	return s.pty.Close(), s.tty.Close()
}

var ErrEmptyCommand error = errors.New("Empty command line")

func (s *Shell) ReadCommand() (*Command, error) {
	var currentLine string
	var backslash bool
	var insideDoubleQuote bool
	var insideSingleQuote bool
	var maybeHeredoc bool
	var insideHeredoc bool
	var heredocMarker string
	var heredocMarkerComplete bool
	var heredocCurrentWord string

	reader := bufio.NewReader(s.tty)

GatherCommand:
	for {
		char, _, err := reader.ReadRune()
		if err != nil {
			return nil, err
		}

		if insideHeredoc {
			currentLine += string(char)

			if !heredocMarkerComplete && char == '<' {
				insideHeredoc = false
				continue
			}

			if char == '\n' {
				log.Printf("Heredoc current word: \"%s\"", heredocCurrentWord)
				if heredocCurrentWord == heredocMarker {
					log.Println("exiting heredoc")
					insideHeredoc = false
					heredocMarkerComplete = false
					break GatherCommand
				}

				if !heredocMarkerComplete {
					log.Println("Encountered heredoc marker: " + heredocCurrentWord)
					heredocMarker = strings.Trim(heredocCurrentWord, "'")
					heredocMarkerComplete = true
				}

				heredocCurrentWord = ""
				continue
			}

			heredocCurrentWord += string(char)
			continue
		}

		if backslash {
			if !insideDoubleQuote && char == '\n' {
				currentLine += string(" ")
				continue
			}

			currentLine += "\\" + string(char)
			backslash = false
			continue
		}

		switch rune(char) {
		case '<':
			currentLine += string(char)
			if maybeHeredoc {
				maybeHeredoc = false
				insideHeredoc = true
				continue
			}

			maybeHeredoc = true
			continue
		case '\\':
			backslash = true
			continue
		case '"':
			insideDoubleQuote = !insideDoubleQuote
			currentLine += string(char)
		case '\'':
			insideSingleQuote = !insideSingleQuote
			currentLine += string(char)
		case '\n':
			if insideSingleQuote || insideDoubleQuote {
				currentLine += string(char)
				continue
			}

			break GatherCommand
		default:
			currentLine += string(char)
		}
	}

	commandLine, err := shlex.Split(currentLine)
	if err != nil {
		return nil, err
	}

	if len(commandLine) == 0 {
		return nil, ErrEmptyCommand
	}

	comm := &Command{Name: commandLine[0]}
	if len(commandLine) > 1 {
		comm.Arguments = commandLine[1:]
	}

	return comm, nil
}

func (s *Shell) SetPrompt(prompt string) {
	log.Printf("Changing prompt to \"%s\"", prompt)
	s.prompt = []byte(prompt)
}

func (s *Shell) WritePrompt() error {
	_, err := s.tty.Write(s.prompt)
	return err
}

func (s *Shell) WriteOutput(output []byte) (int, error) {
	return s.tty.Write(output)
}

func (s *Shell) WriteOutputString(output string) (int, error) {
	return s.WriteOutput([]byte(output))
}

// Winsize stores the Height and Width of a terminal.
type Winsize struct {
	Height uint16
	Width  uint16
	x      uint16 // unused
	y      uint16 // unused
}

// SetWinsize sets the size of the given pty.
func SetWinsize(fd uintptr, winsize *Winsize) {
	log.Printf("window resize %dx%d", winsize.Width, winsize.Height)
	unix.Syscall(
		unix.SYS_IOCTL,
		fd,
		uintptr(unix.TIOCSWINSZ), uintptr(unsafe.Pointer(winsize)),
	)
}