feat: Improve SSH LLM honeypot, preserve session after attacker logout (#179)

* Migrate from deprecated library "golang.org/x/crypto/ssh/terminal" to "golang.org/x/term"

* Feat: Inject OpenAI secret key from environment variable

* Feat: Add test for OpenAI secret key injection from environment variable

* Fix: Correct llmModel value in http-80.yaml configuration

* Feat: Add OPEN_AI_SECRET_KEY environment variable to docker-compose.yml

* Feat: Implement session management for SSHStrategy with command history

configurations/services/http-80.yaml

@@ -22,5 +22,5 @@ commands:
     plugin: "LLMHoneypot"
     statusCode: 200
 plugin:
-  llmModel: "gpt4-o"
+  llmModel: "gpt-4o"
   openAISecretKey: "sk-proj-123456"

docker-compose.yml

@@ -15,5 +15,6 @@ services:
       - "2112:2112" #Prometheus Open Metrics
     environment:
       RABBITMQ_URI: ${RABBITMQ_URI}
+      OPEN_AI_SECRET_KEY: ${OPEN_AI_SECRET_KEY}
     volumes:
       - "./configurations:/configurations"

plugins/llm-integration.go

@@ -7,6 +7,7 @@ import (
 	"github.com/go-resty/resty/v2"
 	"github.com/mariocandela/beelzebub/v3/tracer"
 	log "github.com/sirupsen/logrus"
+	"os"
 	"regexp"
 	"strings"
 )
@@ -95,6 +96,10 @@ func InitLLMHoneypot(config LLMHoneypot) *LLMHoneypot {
 	// Inject the dependencies
 	config.client = resty.New()
 
+	if os.Getenv("OPEN_AI_SECRET_KEY") != "" {
+		config.OpenAIKey = os.Getenv("OPEN_AI_SECRET_KEY")
+	}
+
 	return &config
 }
 

plugins/llm-integration_test.go

@@ -6,6 +6,7 @@ import (
 	"github.com/mariocandela/beelzebub/v3/tracer"
 	"github.com/stretchr/testify/assert"
 	"net/http"
+	"os"
 	"testing"
 )
 
@@ -85,7 +86,7 @@ func TestBuildExecuteModelFailValidation(t *testing.T) {
 		Histories: make([]Message, 0),
 		OpenAIKey: "",
 		Protocol:  tracer.SSH,
-		Model:     "gpt4-o",
+		Model:     "gpt-4o",
 		Provider:  OpenAI,
 	}
 
@@ -96,6 +97,24 @@ func TestBuildExecuteModelFailValidation(t *testing.T) {
 	assert.Equal(t, "openAIKey is empty", err.Error())
 }
 
+func TestBuildExecuteModelOpenAISecretKeyFromEnv(t *testing.T) {
+
+	llmHoneypot := LLMHoneypot{
+		Histories: make([]Message, 0),
+		OpenAIKey: "",
+		Protocol:  tracer.SSH,
+		Model:     "gpt-4o",
+		Provider:  OpenAI,
+	}
+
+	os.Setenv("OPEN_AI_SECRET_KEY", "sdjdnklfjndslkjanfk")
+
+	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
+
+	assert.Equal(t, "sdjdnklfjndslkjanfk", openAIGPTVirtualTerminal.OpenAIKey)
+
+}
+
 func TestBuildExecuteModelWithCustomPrompt(t *testing.T) {
 	client := resty.New()
 	httpmock.ActivateNonDefault(client.GetClient())
@@ -126,7 +145,7 @@ func TestBuildExecuteModelWithCustomPrompt(t *testing.T) {
 		Histories:    make([]Message, 0),
 		OpenAIKey:    "sdjdnklfjndslkjanfk",
 		Protocol:     tracer.HTTP,
-		Model:        "gpt4-o",
+		Model:        "gpt-4o",
 		Provider:     OpenAI,
 		CustomPrompt: "hello world",
 	}
@@ -148,7 +167,7 @@ func TestBuildExecuteModelFailValidationStrategyType(t *testing.T) {
 		Histories: make([]Message, 0),
 		OpenAIKey: "",
 		Protocol:  tracer.TCP,
-		Model:     "gpt4-o",
+		Model:     "gpt-4o",
 		Provider:  OpenAI,
 	}
 
@@ -206,7 +225,7 @@ func TestBuildExecuteModelSSHWithResultsOpenAI(t *testing.T) {
 		Histories: make([]Message, 0),
 		OpenAIKey: "sdjdnklfjndslkjanfk",
 		Protocol:  tracer.SSH,
-		Model:     "gpt4-o",
+		Model:     "gpt-4o",
 		Provider:  OpenAI,
 	}
 
@@ -282,7 +301,7 @@ func TestBuildExecuteModelSSHWithoutResults(t *testing.T) {
 		Histories: make([]Message, 0),
 		OpenAIKey: "sdjdnklfjndslkjanfk",
 		Protocol:  tracer.SSH,
-		Model:     "gpt4-o",
+		Model:     "gpt-4o",
 		Provider:  OpenAI,
 	}
 
@@ -325,7 +344,7 @@ func TestBuildExecuteModelHTTPWithResults(t *testing.T) {
 		Histories: make([]Message, 0),
 		OpenAIKey: "sdjdnklfjndslkjanfk",
 		Protocol:  tracer.HTTP,
-		Model:     "gpt4-o",
+		Model:     "gpt-4o",
 		Provider:  OpenAI,
 	}
 
@@ -362,7 +381,7 @@ func TestBuildExecuteModelHTTPWithoutResults(t *testing.T) {
 		Histories: make([]Message, 0),
 		OpenAIKey: "sdjdnklfjndslkjanfk",
 		Protocol:  tracer.HTTP,
-		Model:     "gpt4-o",
+		Model:     "gpt-4o",
 		Provider:  OpenAI,
 	}
 

protocols/strategies/ssh.go

@@ -13,13 +13,15 @@ import (
 	"github.com/gliderlabs/ssh"
 	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
-	"golang.org/x/crypto/ssh/terminal"
+	"golang.org/x/term"
 )
 
 type SSHStrategy struct {
+	Sessions map[string][]plugins.Message
 }
 
 func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error {
+	sshStrategy.Sessions = make(map[string][]plugins.Message)
 	go func() {
 		server := &ssh.Server{
 			Addr:        beelzebubServiceConfiguration.Address,
@@ -30,7 +32,9 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 				uuidSession := uuid.New()
 
 				host, port, _ := net.SplitHostPort(sess.RemoteAddr().String())
+				sessionKey := host + sess.User()
 
+				// Inline SSH command
 				if sess.RawCommand() != "" {
 					for _, command := range beelzebubServiceConfiguration.Commands {
 						matched, err := regexp.MatchString(command.Regex, sess.RawCommand())
@@ -52,8 +56,14 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 									llmProvider = plugins.OpenAI
 								}
 
+								histories := make([]plugins.Message, 0)
+
+								if sshStrategy.Sessions[sessionKey] != nil {
+									histories = sshStrategy.Sessions[sessionKey]
+								}
+
 								llmHoneypot := plugins.LLMHoneypot{
-									Histories:    make([]plugins.Message, 0),
+									Histories:    histories,
 									OpenAIKey:    beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
 									Protocol:     tracer.SSH,
 									Host:         beelzebubServiceConfiguration.Plugin.Host,
@@ -86,6 +96,10 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 								Command:       sess.RawCommand(),
 								CommandOutput: commandOutput,
 							})
+							var histories []plugins.Message
+							histories = append(histories, plugins.Message{Role: plugins.USER.String(), Content: sess.RawCommand()})
+							histories = append(histories, plugins.Message{Role: plugins.ASSISTANT.String(), Content: commandOutput})
+							sshStrategy.Sessions[sessionKey] = histories
 							tr.TraceEvent(tracer.Event{
 								Msg:    "End SSH Session",
 								Status: tracer.End.String(),
@@ -109,10 +123,14 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 					Description: beelzebubServiceConfiguration.Description,
 				})
 
-				term := terminal.NewTerminal(sess, buildPrompt(sess.User(), beelzebubServiceConfiguration.ServerName))
+				terminal := term.NewTerminal(sess, buildPrompt(sess.User(), beelzebubServiceConfiguration.ServerName))
 				var histories []plugins.Message
+				if sshStrategy.Sessions[sessionKey] != nil {
+					histories = sshStrategy.Sessions[sessionKey]
+				}
+
 				for {
-					commandInput, err := term.ReadLine()
+					commandInput, err := terminal.ReadLine()
 					if err != nil {
 						break
 					}
@@ -160,7 +178,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 							histories = append(histories, plugins.Message{Role: plugins.USER.String(), Content: commandInput})
 							histories = append(histories, plugins.Message{Role: plugins.ASSISTANT.String(), Content: commandOutput})
 
-							term.Write(append([]byte(commandOutput), '\n'))
+							terminal.Write(append([]byte(commandOutput), '\n'))
 
 							tr.TraceEvent(tracer.Event{
 								Msg:           "New SSH Terminal Session",
@@ -178,6 +196,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 						}
 					}
 				}
+				sshStrategy.Sessions[sessionKey] = histories
 				tr.TraceEvent(tracer.Event{
 					Msg:    "End SSH Session",
 					Status: tracer.End.String(),