tramp-blog/commands/shell.go

package commands

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")

type State struct {
	Variables map[string]string
	Commands  map[string]CommandExecutor
}

type Statement interface {
	Words() []Word
	Statements() []Statement
	AddStatement(Statement)
	Evaluate(state State, output io.Writer) (exitcode uint8)
}

type BaseStatement struct {
	words      []Word
	statements []Statement
}

func (s *BaseStatement) Words() []Word {
	return s.words
}

func (s *BaseStatement) Statements() []Statement {
	return s.statements
}

func (s *BaseStatement) AddStatement(statement Statement) {
	s.statements = append(s.statements, statement)
}

func (s *BaseStatement) AddWord(word Word) {
	s.words = append(s.words, word)
}

type Word string

func (w Word) Statements() []Statement {
	return []Statement{}
}

func (w Word) AddStatement(statement Statement) {}

func (w Word) AddWord(word Word) {}

type ParserParameters interface {
	MakeStatement() Statement
	ShouldLeave(Word, rune) bool
	Supports(rune) bool
	SubParsers() []Parser
}

type Parser interface {
	Parse(*bufio.Reader) (Statement, error)
	Parameters() ParserParameters
}

type BaseParser struct {
	parameters ParserParameters
}

func (p *BaseParser) Parameters() ParserParameters {
	return p.parameters
}

type BaseParameters struct {
	EntryMarkers []rune
	subParsers   []Parser
}

func (p *BaseParameters) SubParsers() []Parser {
	return p.subParsers
}

func (p *BaseParameters) Supports(r rune) bool {
	for _, char := range p.EntryMarkers {
		if char == r {
			return true
		}
	}

	return false
}

func (p *BaseParameters) MakeStatement() Statement {
	return &BaseStatement{}
}

func (p *BaseParameters) ShouldLeave(w Word, r rune) bool {
	if r == '\n' || r == ';' {
		return true
	}
}

func (p *BaseParser) Parse(r *bufio.Reader) (Statement, error) {
	var currentWord Word = ""
	statement := p.Parameters().MakeStatement()
	for {
		char, _, err := r.ReadRune()
		if err != nil {
			return statement, err
		}

		if char == ' ' || char == '\t' {
			statement.AddWord(currentWord)
			currentWord = ""
		}

		if p.Parameters().ShouldLeave(currentWord, char) {
			return statement, nil
		}

		var matchedParser bool

		for _, parser := range p.Parameters().SubParsers() {
			if parser.Parameters().Supports(char) {
				matchedParser = true
				nestedStatement, err := parser.Parse(r)

				if word, ok := nestedStatement.(Word); ok {
					statement.AddWord(word)
				} else {
					statement.AddStatement(nestedStatement)
				}

				if err != nil {
					return statement, err
				}
			}
		}

		if !matchedParser {
			currentWord += Word(char)
		}
	}
}

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)),
	)
}