path: root/eows/eows.go

// Package eows is used to Execute commands Over WebSocket
package eows

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
	"unsafe"

	"github.com/Sirupsen/logrus"
	"github.com/googollee/go-socket.io"
	"github.com/kr/pty"
)

// OnInputCB is the function callback used to receive data
type OnInputCB func(e *ExecOverWS, stdin []byte) ([]byte, error)

// EmitOutputCB is the function callback used to emit data
type EmitOutputCB func(e *ExecOverWS, stdout, stderr []byte)

// EmitExitCB is the function callback used to emit exit proc code
type EmitExitCB func(e *ExecOverWS, code int, err error)

// SplitType Type of spliting method to tokenize stdout/stderr
type SplitType uint8

const (
	// SplitLine Split line by line
	SplitLine SplitType = iota
	// SplitChar Split character by character
	SplitChar
	// SplitLineTime Split by line or until a timeout has passed
	SplitLineTime
	// SplitTime Split until a timeout has passed
	SplitTime
)

// Inspired by :
// https://github.com/gorilla/websocket/blob/master/examples/command/main.go

// ExecOverWS .
type ExecOverWS struct {
	Cmd      string           // command name to execute
	Args     []string         // command arguments
	SocketIO *socketio.Socket // websocket
	Sid      string           // websocket ID
	CmdID    string           // command ID

	// Optional fields
	Env            []string                // command environment variables
	CmdExecTimeout int                     // command execution time timeout
	Log            *logrus.Logger          // logger (nil if disabled)
	InputEvent     string                  // websocket input event name
	InputCB        OnInputCB               // stdin callback
	OutputCB       EmitOutputCB            // stdout/stderr callback
	ExitCB         EmitExitCB              // exit proc callback
	UserData       *map[string]interface{} // user data passed to callbacks
	OutSplit       SplitType               // split method to tokenize stdout/stderr
	LineTimeSpan   int64                   // time span (only used with SplitTime or SplitLineTime)
	PtyMode        bool                    // Allocate a pseudo-terminal (allow to execute screen-based program)
	PtyTermEcho    bool                    // Turn on/off terminal echo

	// Private fields

	proc       *os.Process
	command    *exec.Cmd
	ptmx       *os.File
	procExited bool
}

var cmdIDMap = make(map[string]*ExecOverWS)

// New creates a new instace of eows
func New(cmd string, args []string, so *socketio.Socket, soID, cmdID string) *ExecOverWS {

	e := &ExecOverWS{
		Cmd:            cmd,
		Args:           args,
		SocketIO:       so,
		Sid:            soID,
		CmdID:          cmdID,
		CmdExecTimeout: -1,            // default no timeout
		OutSplit:       SplitLineTime, // default split by line with time
		LineTimeSpan:   500 * time.Millisecond.Nanoseconds(),
		PtyMode:        false,
		PtyTermEcho:    true,
	}

	cmdIDMap[cmdID] = e

	return e
}

// GetEows gets ExecOverWS object from command ID
func GetEows(cmdID string) *ExecOverWS {
	if _, ok := cmdIDMap[cmdID]; !ok {
		return nil
	}
	return cmdIDMap[cmdID]
}

// Start executes the command and redirect stdout/stderr into a WebSocket
func (e *ExecOverWS) Start() error {
	var err error
	var outr, outw, errr, errw, inr, inw *os.File

	bashArgs := []string{"/bin/bash", "-c", e.Cmd + " " + strings.Join(e.Args, " ")}

	// no timeout == 1 year
	if e.CmdExecTimeout == -1 {
		e.CmdExecTimeout = 365 * 24 * 60 * 60
	}

	e.procExited = false

	if e.PtyMode {

		e.command = exec.Command(bashArgs[0], bashArgs[1:]...)
		e.command.Env = append(os.Environ(), e.Env...)
		e.ptmx, err = pty.Start(e.command)
		if err != nil {
			err = fmt.Errorf("Process start error: " + err.Error())
			goto exitErr
		}
		e.proc = e.command.Process

		// Turn off terminal echo
		if !e.PtyTermEcho {
			e.terminalEcho(e.ptmx, false)
		}

	} else {

		// Create pipes
		outr, outw, err = os.Pipe()
		if err != nil {
			err = fmt.Errorf("Pipe stdout error: " + err.Error())
			goto exitErr
		}

		errr, errw, err = os.Pipe()
		if err != nil {
			err = fmt.Errorf("Pipe stderr error: " + err.Error())
			goto exitErr
		}

		inr, inw, err = os.Pipe()
		if err != nil {
			err = fmt.Errorf("Pipe stdin error: " + err.Error())
			goto exitErr
		}

		e.proc, err = os.StartProcess(bashArgs[0], bashArgs, &os.ProcAttr{
			Files: []*os.File{inr, outw, errw},
			Env:   append(os.Environ(), e.Env...),
		})
		if err != nil {
			err = fmt.Errorf("Process start error: " + err.Error())
			goto exitErr
		}
	}

	go func() {
		stdoutDone := make(chan struct{})

		if e.PtyMode {
			// Make sure to close the pty at the end.
			defer e.ptmx.Close()

			// Handle both stdout mixed with stderr
			go e.ptsPumpStdout(e.ptmx, stdoutDone)

			// Blocking function that poll input or wait for end of process
			e.pumpStdin(e.ptmx)

		} else {
			// Make sure to close all pipes
			defer outr.Close()
			defer outw.Close()
			defer errr.Close()
			defer errw.Close()
			defer inr.Close()
			defer inw.Close()

			// Handle stdout + stderr
			go e.pipePumpStdout(outr, stdoutDone)
			go e.pipePumpStderr(errr)

			// Blocking function that poll input or wait for end of process
			e.pumpStdin(inw)
		}

		if status, err := e.proc.Wait(); err == nil {
			// Other commands need a bonk on the head.
			if !status.Exited() {
				if err := e.proc.Signal(os.Interrupt); err != nil {
					e.logError("Proc interrupt:", err)
				}

				select {
				case <-stdoutDone:
				case <-time.After(time.Second):
					// A bigger bonk on the head.
					if err := e.proc.Signal(os.Kill); err != nil {
						e.logError("Proc term:", err)
					}
					<-stdoutDone
				}
			}
		}

		delete(cmdIDMap, e.CmdID)
	}()

	return nil

exitErr:
	for _, pf := range []*os.File{outr, outw, errr, errw, inr, inw} {
		pf.Close()
	}
	return err
}

// TerminalSetSize Set terminal size
func (e *ExecOverWS) TerminalSetSize(rows, cols uint16) error {
	if !e.PtyMode || e.ptmx == nil {
		return fmt.Errorf("PtyMode not set")
	}
	w, err := pty.GetsizeFull(e.ptmx)
	if err != nil {
		return err
	}
	return e.TerminalSetSizePos(rows, cols, w.X, w.Y)
}

// TerminalSetSizePos Set terminal size and position
func (e *ExecOverWS) TerminalSetSizePos(rows, cols, x, y uint16) error {
	if !e.PtyMode || e.ptmx == nil {
		return fmt.Errorf("PtyMode not set")
	}
	winSz := pty.Winsize{Rows: rows, Cols: cols, X: x, Y: y}
	return pty.Setsize(e.ptmx, &winSz)
}

/**
 * Private functions
 **/

// terminalEcho Enable or disable echoing terminal input.
// This is useful specifically for when users enter passwords.
func (e *ExecOverWS) terminalEcho(ff *os.File, show bool) {
	var termios = &syscall.Termios{}

	fd := ff.Fd()

	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd,
		syscall.TCGETS, uintptr(unsafe.Pointer(termios))); err != 0 {
		return
	}

	if show {
		termios.Lflag |= syscall.ECHO
	} else {
		termios.Lflag &^= syscall.ECHO
	}

	if _, _, err := syscall.Syscall(syscall.SYS_IOCTL, fd,
		uintptr(syscall.TCSETS),
		uintptr(unsafe.Pointer(termios))); err != 0 {
		return
	}
}

func (e *ExecOverWS) logDebug(format string, a ...interface{}) {
	if e.Log != nil {
		e.Log.Debugf(format, a...)
	}
}

func (e *ExecOverWS) logError(format string, a ...interface{}) {
	if e.Log != nil {
		e.Log.Errorf(format, a...)
	}
}