From 97791896e1ac60c985294c1c250a703a851b9534 Mon Sep 17 00:00:00 2001
From: Zhi Wang <zwang@cs.fsu.edu>
Date: Thu, 6 Jan 2022 20:21:11 -0500
Subject: [PATCH] add an option for command line

---
 main.go | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/main.go b/main.go
index 47e994d..d7a56d2 100644
--- a/main.go
+++ b/main.go
@@ -155,6 +155,8 @@ func fromPtyStdout(ws *websocket.Conn, ptmx *os.File, done chan struct{}) {
 	time.Sleep(closeGracePeriod)
 }
 
+var cmdToExec string = "bash"
+
 // handle websockets
 func wsHandler(w http.ResponseWriter, r *http.Request) {
 	ws, err := upgrader.Upgrade(w, r, nil)
@@ -168,7 +170,7 @@ func wsHandler(w http.ResponseWriter, r *http.Request) {
 
 	log.Println("\n\nCreated the websocket")
 
-	ptmx, cmd, err := createPty("bash")
+	ptmx, cmd, err := createPty(cmdToExec)
 
 	if err != nil {
 		log.Println("Failed to create PTY: ", err)
@@ -235,6 +237,17 @@ func main() {
 		gin.DefaultWriter = fp
 	}
 
+	// parse the arguments. User can pass the command to execute
+	// by default, we use bash, but macos users might want to use zsh
+	// you can also run single program, such as pstree, htop...
+	// but program might misbehave (htop seems to be fine)
+	args := os.Args
+
+	if len(args) > 1 {
+		cmdToExec = strings.Join(args[1:], " ")
+		log.Println(cmdToExec)
+	}
+
 	rt := gin.Default()
 
 	rt.SetTrustedProxies(nil)