From ffcd15b8acb425c7905d0fddeffe0162270946f8 Mon Sep 17 00:00:00 2001
From: Claudia Beresford <claudiaberesford@gmail.com>
Date: Fri, 23 Jun 2023 14:30:55 +0100
Subject: [PATCH] feat: Add ExecStream for long-running commands

Attaches a given writer to session StdoutPipe and StderrPipe so that
long-running commands can be seen in real time, rather than waiting for
a dump of info at the end.
---
 README.md    |  2 +-
 simplessh.go | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 086b0c8..a20bdb0 100644
--- a/README.md
+++ b/README.md
@@ -44,4 +44,4 @@ func main() {
 ```
 
 ## License
-SimpleSSH is licensed under the MIT license.
\ No newline at end of file
+SimpleSSH is licensed under the MIT license.
diff --git a/simplessh.go b/simplessh.go
index c4d4fb3..cc8ab65 100644
--- a/simplessh.go
+++ b/simplessh.go
@@ -1,6 +1,7 @@
 package simplessh
 
 import (
+	"bufio"
 	"bytes"
 	"fmt"
 	"io"
@@ -232,6 +233,31 @@ func (c *Client) Exec(cmd string) ([]byte, error) {
 	return session.CombinedOutput(cmd)
 }
 
+// Execute cmd on the remote host and return combined stderr and stdout in
+// real time
+func (c *Client) ExecStream(cmd string, stdout io.Writer) error {
+	session, err := c.SSHClient.NewSession()
+	if err != nil {
+		return err
+	}
+	defer session.Close()
+
+	sessOut, err := session.StdoutPipe()
+	if err != nil {
+		return err
+	}
+
+	sessErr, err := session.StderrPipe()
+	if err != nil {
+		return err
+	}
+
+	go streamOutput(sessOut, stdout)
+	go streamOutput(sessErr, stdout)
+
+	return session.Run(cmd)
+}
+
 // Execute cmd via sudo. Do not include the sudo command in
 // the cmd string. For example: Client.ExecSudo("uptime", "password").
 // If you are using passwordless sudo you can use the regular Exec()
@@ -375,3 +401,26 @@ func addPortToHost(host string) string {
 
 	return host
 }
+
+func streamOutput(read io.Reader, write io.Writer) {
+	var (
+		r    = bufio.NewReader(read)
+		line = ""
+	)
+
+	for {
+		b, err := r.ReadByte()
+		if err != nil {
+			break
+		}
+
+		if b == byte('\n') {
+			fmt.Fprintln(write, line)
+			line = ""
+
+			continue
+		}
+
+		line += string(b)
+	}
+}