summaryrefslogtreecommitdiffstats
path: root/eows/eows.go
diff options
context:
space:
mode:
Diffstat (limited to 'eows/eows.go')
-rw-r--r--eows/eows.go287
1 files changed, 287 insertions, 0 deletions
diff --git a/eows/eows.go b/eows/eows.go
new file mode 100644
index 0000000..9d0b520
--- /dev/null
+++ b/eows/eows.go
@@ -0,0 +1,287 @@
+// 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...)
+ }
+}