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

.github/workflows/codeql.yml

@@ -27,7 +27,7 @@ jobs:
       uses: actions/checkout@v3
 
     - name: Initialize CodeQL
-      uses: github/codeql-action/init@v2
+      uses: github/codeql-action/init@v3
       with:
         languages: ${{ matrix.language }}
 
@@ -35,6 +35,6 @@ jobs:
       run: go build ./...
 
     - name: Perform CodeQL Analysis
-      uses: github/codeql-action/analyze@v2
+      uses: github/codeql-action/analyze@v3
       with:
         category: "/language:${{matrix.language}}"

.github/workflows/docker-image.yml

@@ -1,31 +1,32 @@
+---
 name: Docker Hub Image
 
 on:
   push:
     tags:
       - 'v*.*.*'
-
 jobs:
   CD:
     runs-on: ubuntu-latest
     steps:
-      -
+      - name: Checkout
-        name: Checkout
         uses: actions/checkout@v3
-      -
+      - name: Login to Docker Hub
-        name: Login to Docker Hub
         uses: docker/login-action@v2
         with:
           username: ${{ secrets.DOCKER_USER }}
           password: ${{ secrets.DOCKER_ACCESS_TOKEN }}
-      -
+      - name: Set up QEMU
-        name: Set up Docker Buildx
+        uses: docker/setup-qemu-action@v3
+      - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v2
-      -
+      - name: Build and push
-        name: Build and push
         uses: docker/build-push-action@v4
         with:
           context: .
           file: ./Dockerfile
           push: true
-          tags: m4r10/beelzebub:${{  github.ref_name }}
+          tags: |
+            m4r10/beelzebub:${{  github.ref_name }}
+            m4r10/beelzebub:latest
+          platforms: linux/amd64,linux/arm64

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"

integration_test/integration_test.go

@@ -67,8 +67,11 @@ func (suite *IntegrationTestSuite) TestInvokeHTTPHoneypot() {
 	response, err := resty.New().R().
 		Get(suite.httpHoneypotHost + "/index.php")
 
+	response.Header().Del("Date")
+
 	suite.Require().NoError(err)
 	suite.Equal(http.StatusOK, response.StatusCode())
+	suite.Equal(http.Header{"Content-Length": []string{"15"}, "Content-Type": []string{"text/html"}, "Server": []string{"Apache/2.4.53 (Debian)"}, "X-Powered-By": []string{"PHP/7.4.29"}}, response.Header())
 	suite.Equal("mocked response", string(response.Body()))
 
 	response, err = resty.New().R().

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/http.go

@@ -117,7 +117,7 @@ func traceRequest(request *http.Request, tr tracer.Tracer, HoneypotDescription s
 		HostHTTPRequest: request.Host,
 		UserAgent:       request.UserAgent(),
 		Cookies:         mapCookiesToString(request.Cookies()),
-		Headers:         mapHeaderToString(request.Header),
+		Headers:         request.Header,
 		Status:          tracer.Stateless.String(),
 		RemoteAddr:      request.RemoteAddr,
 		SourceIp:        host,
@@ -133,18 +133,6 @@ func traceRequest(request *http.Request, tr tracer.Tracer, HoneypotDescription s
 	tr.TraceEvent(event)
 }
 
-func mapHeaderToString(headers http.Header) string {
-	headersString := ""
-
-	for key := range headers {
-		for _, values := range headers[key] {
-			headersString += fmt.Sprintf("[Key: %s, values: %s],", key, values)
-		}
-	}
-
-	return headersString
-}
-
 func mapCookiesToString(cookies []*http.Cookie) string {
 	cookiesString := ""
 

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(),

tracer/tracer.go

@@ -27,7 +27,7 @@ type Event struct {
 	User            string
 	Password        string
 	Client          string
-	Headers         string
+	Headers         map[string][]string
 	Cookies         string
 	UserAgent       string
 	HostHTTPRequest string