From bcd37285c1398b08199a190843dd9d5527deecfd Mon Sep 17 00:00:00 2001
From: Shulhan <ms@kilabit.info>
Date: Sun, 15 Aug 2021 03:38:35 +0700
Subject: lib/ssh: add method to set session output and error

Previously, all of the SSH output and error goes to os.Stdout and
os.Stderr.

This changes add method SetSessionOutputError to change the output and
error for future remote execution.
---
 lib/ssh/client.go | 34 ++++++++++++++++++++++++++--------
 1 file changed, 26 insertions(+), 8 deletions(-)

(limited to 'lib/ssh')

diff --git a/lib/ssh/client.go b/lib/ssh/client.go
index 77eb99d3..7deaefee 100644
--- a/lib/ssh/client.go
+++ b/lib/ssh/client.go
@@ -6,6 +6,7 @@ package ssh
 
 import (
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"os"
@@ -23,7 +24,9 @@ import (
 type Client struct {
 	*ssh.Client
 
-	cfg *config.Section
+	cfg    *config.Section
+	stdout io.Writer
+	stderr io.Writer
 }
 
 //
@@ -66,7 +69,9 @@ func NewClientFromConfig(cfg *config.Section) (cl *Client, err error) {
 	}
 
 	cl = &Client{
-		cfg: cfg,
+		cfg:    cfg,
+		stdout: os.Stdout,
+		stderr: os.Stderr,
 	}
 
 	remoteAddr := fmt.Sprintf("%s:%s", cfg.Hostname, cfg.Port)
@@ -88,8 +93,8 @@ func (cl *Client) Execute(cmd string) (err error) {
 		return fmt.Errorf("ssh: NewSession: " + err.Error())
 	}
 
-	sess.Stdout = os.Stdout
-	sess.Stderr = os.Stderr
+	sess.Stdout = cl.stdout
+	sess.Stderr = cl.stderr
 
 	for k, v := range cl.cfg.Environments {
 		err = sess.Setenv(k, v)
@@ -137,8 +142,8 @@ func (cl *Client) ScpGet(remote, local string) (err error) {
 	cmd := exec.Command("scp", args...)
 
 	cmd.Dir = cl.cfg.WorkingDir
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
+	cmd.Stdout = cl.stdout
+	cmd.Stderr = cl.stderr
 
 	err = cmd.Run()
 	if err != nil {
@@ -177,8 +182,8 @@ func (cl *Client) ScpPut(local, remote string) (err error) {
 	cmd := exec.Command("scp", args...)
 
 	cmd.Dir = cl.cfg.WorkingDir
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
+	cmd.Stdout = cl.stdout
+	cmd.Stderr = cl.stderr
 
 	err = cmd.Run()
 	if err != nil {
@@ -188,6 +193,19 @@ func (cl *Client) ScpPut(local, remote string) (err error) {
 	return nil
 }
 
+//
+// SetSessionOutputError set the standard output and error for future remote
+// execution.
+//
+func (cl *Client) SetSessionOutputError(stdout, stderr io.Writer) {
+	if stdout != nil {
+		cl.stdout = stdout
+	}
+	if stderr != nil {
+		cl.stderr = stderr
+	}
+}
+
 func (cl *Client) String() string {
 	return cl.cfg.User + "@" + cl.cfg.Hostname + ":" + cl.cfg.Port
 }
-- 
cgit v1.3-5-g9baa