Low-level shell wrapper package

wrapper.go

package wrapper

import (
	"os/exec"
	"os"
	"sync"
	"github.com/kr/pty"
	"github.com/creack/termios/raw"
	"github.com/creack/termios/win"
	"io"
)

const outputLimit = 10 * 1024 * 1024

type ShellWrapper struct {
	output  []byte
	size    int

	mutex   sync.Mutex

	command *exec.Cmd
	pty     *os.File
}

func New(command *exec.Cmd) *ShellWrapper {
	return &ShellWrapper{
		output: make([]byte, outputLimit),
		size: 0,
		mutex: sync.Mutex{},
		command: command,
	}
}

func (w *ShellWrapper) Start() error {
	f, err := pty.Start(w.command)
	if err != nil {
		return err
	}

	w.command.Stdin = os.Stdin
	w.pty = f

	_, err = raw.MakeRaw(os.Stdin.Fd())
	if err != nil {
		return err
	}
	ws, err := win.GetWinsize(os.Stdin.Fd())
	if err != nil {
		return err
	}

	if err := win.SetWinsize(f.Fd(), ws); err != nil {
		return err
	}

	go w.handleStdout()
	go func() { _, _ = io.Copy(f, os.Stdin) }()

	return nil
}

func (w *ShellWrapper) Reset() {
	w.mutex.Lock()
	w.size = 0
	w.mutex.Unlock()
}

func (w *ShellWrapper) GetOutput() string {
	return string(w.output[:w.size])
}

func (w *ShellWrapper) handleStdout() {
	buf := make([]byte, 1)
	for {
		n, err := w.pty.Read(buf)
		if err != nil {
			panic(err)
		}

		if n == 0 {
			os.Exit(0)
		}

		os.Stdout.Write(buf)

		if (w.size < outputLimit) {
			w.mutex.Lock()
			w.output[w.size] = buf[0]
			w.size += 1
			w.mutex.Unlock()
		}
	}
}