From c64797f2d9621a95f2145d60ac3ff8d5e3ef8f3c Mon Sep 17 00:00:00 2001
From: Russ Magee <rmagee@gmail.com>
Date: Fri, 29 Jun 2018 16:54:20 -0700
Subject: [PATCH] Basic server-side recording of exitStatus of pty(cmd). TODO:
 sending of exitStatus to client and client handling of said packet via   a
 WritePacket() with unique existStatus op.

---
 hkexshd/hkexshd.go | 103 ++++++++++++++++++++++++++++-----------------
 1 file changed, 65 insertions(+), 38 deletions(-)

diff --git a/hkexshd/hkexshd.go b/hkexshd/hkexshd.go
index 5fbb77e..e70e0d4 100644
--- a/hkexshd/hkexshd.go
+++ b/hkexshd/hkexshd.go
@@ -80,7 +80,7 @@ func runCmdAs(who string, cmd string, conn hkex.Conn) (err error) {
 // Run a command (via default shell) as a specific user
 //
 // Uses ptys to support commands which expect a terminal.
-func runShellAs(who string, cmd string, interactive bool, conn hkexsh.Conn, chaffing bool) (err error) {
+func runShellAs(who string, cmd string, interactive bool, conn hkexsh.Conn, chaffing bool) (err error, exitStatus int) {
 	u, _ := user.Lookup(who)
 	var uid, gid uint32
 	fmt.Sscanf(u.Uid, "%d", &uid)
@@ -117,46 +117,65 @@ func runShellAs(who string, cmd string, interactive bool, conn hkexsh.Conn, chaf
 	// Start the command with a pty.
 	ptmx, err := pty.Start(c) // returns immediately with ptmx file
 	if err != nil {
-		return err
+		return err, 0
 	}
 	// Make sure to close the pty at the end.
 	defer func() { _ = ptmx.Close() }() // Best effort.
 
-	// Watch for term resizes
-	go func() {
-		for sz := range conn.WinCh {
-			log.Printf("[Setting term size to: %v %v]\n", sz.Rows, sz.Cols)
-			pty.Setsize(ptmx, &pty.Winsize{Rows: sz.Rows, Cols: sz.Cols})
-		}
-	}()
-
-	// Copy stdin to the pty.. (bgnd goroutine)
-	go func() {
-		_, e := io.Copy(ptmx, conn)
-		if e != nil {
-			log.Printf("** std->pty ended **\n")
-			return
-		}
-	}()
-
-	if chaffing {
-		conn.EnableChaff()
-	}
-	defer conn.DisableChaff()
-	defer conn.ShutdownChaff()
-
-	// ..and the pty to stdout.
-	_, e := io.Copy(conn, ptmx)
-	if e != nil {
-		log.Printf("** pty->stdout ended **\n")
-		return
-	}
-
-	//err = c.Run()  // returns when c finishes.
-
 	log.Printf("[%s]\n", cmd)
 	if err != nil {
 		log.Printf("Command finished with error: %v", err)
+	} else {
+
+		// Watch for term resizes
+		go func() {
+			for sz := range conn.WinCh {
+				log.Printf("[Setting term size to: %v %v]\n", sz.Rows, sz.Cols)
+				pty.Setsize(ptmx, &pty.Winsize{Rows: sz.Rows, Cols: sz.Cols})
+			}
+		}()
+
+		// Copy stdin to the pty.. (bgnd goroutine)
+		go func() {
+			_, e := io.Copy(ptmx, conn)
+			if e != nil {
+				log.Printf("** std->pty ended **\n")
+				return
+			}
+		}()
+
+		if chaffing {
+			conn.EnableChaff()
+		}
+		defer conn.DisableChaff()
+		defer conn.ShutdownChaff()
+
+		// ..and the pty to stdout.
+		go func() {
+			_, e := io.Copy(conn, ptmx)
+			if e != nil {
+				log.Printf("** pty->stdout ended **\n")
+				return
+			}
+		}()
+
+		// The above io.Copy() will exit when the command attached
+		// to the pty exits
+
+		if err := c.Wait(); err != nil {
+			if exiterr, ok := err.(*exec.ExitError); ok {
+				// The program has exited with an exit code != 0
+
+				// This works on both Unix and Windows. Although package
+				// syscall is generally platform dependent, WaitStatus is
+				// defined for both Unix and Windows and in both cases has
+				// an ExitStatus() method with the same signature.
+				if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
+					exitStatus = status.ExitStatus()
+					log.Printf("Exit Status: %d", exitStatus)
+				}
+			}
+		}
 	}
 	return
 }
@@ -303,11 +322,15 @@ func main() {
 					hname := goutmp.GetHost(addr.String())
 
 					log.Printf("[Running command for [%s@%s]]\n", rec.who, hname)
-					runShellAs(string(rec.who), string(rec.cmd), false, conn, chaffEnabled)
+					runErr, cmdStatus := runShellAs(string(rec.who), string(rec.cmd), false, conn, chaffEnabled)
 					// Returned hopefully via an EOF or exit/logout;
 					// Clear current op so user can enter next, or EOF
 					rec.op[0] = 0
-					log.Printf("[Command completed for [%s@%s]\n", rec.who, hname)
+					if runErr != nil {
+						log.Printf("[Error spawning cmd for %s@%s]\n", rec.who, hname)
+					} else {
+						log.Printf("[Command completed for %s@%s, status %d]\n", rec.who, hname, cmdStatus)
+					}
 				} else if rec.op[0] == 's' {
 					// Interactive session
 					addr := c.RemoteAddr()
@@ -317,11 +340,15 @@ func main() {
 					utmpx := goutmp.Put_utmp(string(rec.who), hname)
 					defer func() { goutmp.Unput_utmp(utmpx) }()
 					goutmp.Put_lastlog_entry("hkexsh", string(rec.who), hname)
-					runShellAs(string(rec.who), string(rec.cmd), true, conn, chaffEnabled)
+					runErr, cmdStatus := runShellAs(string(rec.who), string(rec.cmd), true, conn, chaffEnabled)
 					// Returned hopefully via an EOF or exit/logout;
 					// Clear current op so user can enter next, or EOF
 					rec.op[0] = 0
-					log.Printf("[Exiting shell for [%s@%s]]\n", rec.who, hname)
+					if runErr != nil {
+						log.Printf("[Error spawning shell for %s@%s]\n", rec.who, hname)
+					} else {
+						log.Printf("[Shell completed for %s@%s, status %d]\n", rec.who, hname, cmdStatus)
+					}
 				} else {
 					log.Println("[Bad cmdSpec]")
 				}