From 49a4022f569ecd19f91da3d5ff753e3f441a3ccd Mon Sep 17 00:00:00 2001
From: Mario Candela <m4r10.php@gmail.com>
Date: Fri, 14 Apr 2023 22:35:11 +0200
Subject: [PATCH] feat: Trace command output for SSH Strategy (#35)

manage command output SSH strategy
---
 protocols/strategies/secure_shell.go | 28 ++++++++++++++++------------
 tracer/tracer.go                     |  1 +
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/protocols/strategies/secure_shell.go b/protocols/strategies/secure_shell.go
index 7a41167..5bef19c 100644
--- a/protocols/strategies/secure_shell.go
+++ b/protocols/strategies/secure_shell.go
@@ -5,13 +5,14 @@ import (
 	"beelzebub/plugins"
 	"beelzebub/tracer"
 	"fmt"
+	"regexp"
+	"strings"
+	"time"
+
 	"github.com/gliderlabs/ssh"
 	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/crypto/ssh/terminal"
-	"regexp"
-	"strings"
-	"time"
 )
 
 type SecureShellStrategy struct {
@@ -46,15 +47,7 @@ func (SSHStrategy *SecureShellStrategy) Init(beelzebubServiceConfiguration parse
 					if err != nil {
 						break
 					}
-					tr.TraceEvent(tracer.Event{
-						Msg:         "New SSH Terminal Session",
-						RemoteAddr:  sess.RemoteAddr().String(),
-						Status:      tracer.Interaction.String(),
-						Command:     commandInput,
-						ID:          uuidSession.String(),
-						Protocol:    tracer.SSH.String(),
-						Description: beelzebubServiceConfiguration.Description,
-					})
+
 					if commandInput == "exit" {
 						break
 					}
@@ -81,6 +74,17 @@ func (SSHStrategy *SecureShellStrategy) Init(beelzebubServiceConfiguration parse
 							histories = append(histories, plugins.History{Input: commandInput, Output: commandOutput})
 
 							term.Write(append([]byte(commandOutput), '\n'))
+
+							tr.TraceEvent(tracer.Event{
+								Msg:           "New SSH Terminal Session",
+								RemoteAddr:    sess.RemoteAddr().String(),
+								Status:        tracer.Interaction.String(),
+								Command:       commandInput,
+								CommandOutput: commandOutput,
+								ID:            uuidSession.String(),
+								Protocol:      tracer.SSH.String(),
+								Description:   beelzebubServiceConfiguration.Description,
+							})
 							break
 						}
 					}
diff --git a/tracer/tracer.go b/tracer/tracer.go
index ea42464..658cc90 100644
--- a/tracer/tracer.go
+++ b/tracer/tracer.go
@@ -12,6 +12,7 @@ type Event struct {
 	RemoteAddr      string
 	Protocol        string
 	Command         string
+	CommandOutput   string
 	Status          string
 	Msg             string
 	ID              string