From aca580237943f44b917bc0cb327c68b70f38a5ad Mon Sep 17 00:00:00 2001
From: Vasile Popescu <elisescu@elisescu.com>
Date: Fri, 18 Nov 2022 11:59:56 +0100
Subject: [PATCH] Add an option to run headless (no interactive input)

Allow `tty-share` to run without having an interactive input terminal.
---
 main.go       | 30 +++++++++++++++++++-----------
 pty_master.go | 49 ++++++++++++++++++++++++++++++++++---------------
 2 files changed, 53 insertions(+), 26 deletions(-)

diff --git a/main.go b/main.go
index a6925dd..bbaf4f3 100644
--- a/main.go
+++ b/main.go
@@ -75,6 +75,9 @@ Flags:
 	publicSession := flag.Bool("public", false, "Create a public session")
 	noTLS := flag.Bool("no-tls", false, "Don't use TLS to connect to the tty-proxy server. Useful for local debugging")
 	noWaitEnter := flag.Bool("no-wait", false, "Don't wait for the Enter press before starting the session")
+	headless := flag.Bool("headless", false, "Don't expect an interactive terminal at stdin")
+	headlessCols := flag.Int("headless-cols", 80, "Number of cols for the allocated pty when running headless")
+	headlessRows := flag.Int("headless-rows", 25, "Number of rows for the allocated pty when running headless")
 	detachKeys := flag.String("detach-keys", "ctrl-o,ctrl-c", "Sequence of keys to press for closing the connection. Supported: https://godoc.org/github.com/moby/term#pkg-variables.")
 	verbose := flag.Bool("verbose", false, "Verbose logging")
 	flag.Usage = func() {
@@ -121,7 +124,7 @@ Flags:
 	}
 
 	// tty-share works as a server, from here on
-	if !isStdinTerminal() {
+	if !isStdinTerminal() && !*headless {
 		fmt.Printf("Input not a tty\n")
 		os.Exit(1)
 	}
@@ -153,7 +156,7 @@ Flags:
 		)
 	}
 
-	ptyMaster := ptyMasterNew()
+	ptyMaster := ptyMasterNew(*headless, *headlessCols, *headlessRows)
 	err := ptyMaster.Start(*commandName, strings.Fields(*commandArgs), envVars)
 	if err != nil {
 		log.Errorf("Cannot start the %s command: %s", *commandName, err.Error())
@@ -168,7 +171,7 @@ Flags:
 
 	fmt.Printf("local session: http://%s/s/local/\n", *listenAddress)
 
-	if !*noWaitEnter {
+	if !*noWaitEnter && !*headless {
 		fmt.Printf("Press Enter to continue!\n")
 		bufio.NewReader(os.Stdin).ReadString('\n')
 	}
@@ -195,8 +198,11 @@ Flags:
 		server.WindowSize(cols, rows)
 	})
 
-	mw := io.MultiWriter(os.Stdout, server)
-
+	var mw io.Writer
+	mw = server
+	if !*headless {
+		mw = io.MultiWriter(os.Stdout, server)
+	}
 
 	go func() {
 		err := server.Run()
@@ -213,12 +219,14 @@ Flags:
 		}
 	}()
 
-	go func() {
-		_, err := io.Copy(ptyMaster, os.Stdin)
-		if err != nil {
-			stopPtyAndRestore()
-		}
-	}()
+	if !*headless {
+		go func() {
+			_, err := io.Copy(ptyMaster, os.Stdin)
+			if err != nil {
+				stopPtyAndRestore()
+			}
+		}()
+	}
 
 	ptyMaster.Wait()
 	fmt.Printf("tty-share finished\n\n\r")
diff --git a/pty_master.go b/pty_master.go
index 002c389..487e3c9 100644
--- a/pty_master.go
+++ b/pty_master.go
@@ -20,10 +20,13 @@ type ptyMaster struct {
 	ptyFile           *os.File
 	command           *exec.Cmd
 	terminalInitState *terminal.State
+	headless          bool
+	headlessCols      int
+	headlessRows      int
 }
 
-func ptyMasterNew() *ptyMaster {
-	return &ptyMaster{}
+func ptyMasterNew(headless bool, headlessCols, headlessRows int) *ptyMaster {
+	return &ptyMaster{headless: headless, headlessCols: headlessCols, headlessRows: headlessRows}
 }
 
 func isStdinTerminal() bool {
@@ -40,12 +43,21 @@ func (pty *ptyMaster) Start(command string, args []string, envVars []string) (er
 	}
 
 	// Set the initial window size
-	cols, rows, err := terminal.GetSize(0)
+	cols, rows := pty.headlessCols, pty.headlessRows
+
+	if !pty.headless {
+		cols, rows, err = terminal.GetSize(0)
+	}
+
 	pty.SetWinSize(rows, cols)
 	return
 }
 
 func (pty *ptyMaster) MakeRaw() (err error) {
+	// don't do anything if running headless
+	if pty.headless {
+		return nil
+	}
 
 	// Save the initial state of the terminal, before making it RAW. Note that this terminal is the
 	// terminal under which the tty-share command has been started, and it's identified via the
@@ -59,20 +71,25 @@ func (pty *ptyMaster) MakeRaw() (err error) {
 }
 
 func (pty *ptyMaster) SetWinChangeCB(winChangedCB onWindowChangedCB) {
-	// Start listening for window changes
-	go onWindowChanges(func(cols, rows int) {
-		// TODO:policy: should the server decide here if we care about the size and set it
-		// right here?
-		pty.SetWinSize(rows, cols)
-
-		// Notify the ptyMaster user of the window changes, to be sent to the remote side
-		winChangedCB(cols, rows)
-	})
+	// Start listening for window changes if not running headless
+	if !pty.headless {
+		go onWindowChanges(func(cols, rows int) {
+			// TODO:policy: should the server decide here if we care about the size and set it
+			// right here?
+			pty.SetWinSize(rows, cols)
+
+			// Notify the ptyMaster user of the window changes, to be sent to the remote side
+			winChangedCB(cols, rows)
+		})
+	}
 }
 
 func (pty *ptyMaster) GetWinSize() (int, int, error) {
-	cols, rows, err := terminal.GetSize(0)
-	return cols, rows, err
+	if pty.headless {
+		return pty.headlessCols, pty.headlessRows, nil
+	} else {
+		return terminal.GetSize(0)
+	}
 }
 
 func (pty *ptyMaster) Write(b []byte) (int, error) {
@@ -110,7 +127,9 @@ func (pty *ptyMaster) Wait() (err error) {
 }
 
 func (pty *ptyMaster) Restore() {
-	terminal.Restore(0, pty.terminalInitState)
+	if !pty.headless {
+		terminal.Restore(0, pty.terminalInitState)
+	}
 	return
 }