From 933f02911b51e202fe46e7b7a6489df7ecdc62d9 Mon Sep 17 00:00:00 2001 From: Mario Candela Date: Sun, 9 Mar 2025 13:17:04 +0100 Subject: [PATCH] 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 | 2 +- docker-compose.yml | 1 + plugins/llm-integration.go | 5 +++++ plugins/llm-integration_test.go | 33 ++++++++++++++++++++++------ protocols/strategies/ssh.go | 29 +++++++++++++++++++----- 5 files changed, 57 insertions(+), 13 deletions(-) diff --git a/configurations/services/http-80.yaml b/configurations/services/http-80.yaml index c9e1110..faed94b 100644 --- a/configurations/services/http-80.yaml +++ b/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" \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml index 30b234b..0ad1743 100644 --- a/docker-compose.yml +++ b/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" \ No newline at end of file diff --git a/plugins/llm-integration.go b/plugins/llm-integration.go index b109c7c..30b17a2 100644 --- a/plugins/llm-integration.go +++ b/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 } diff --git a/plugins/llm-integration_test.go b/plugins/llm-integration_test.go index 7b3096b..3810ffa 100644 --- a/plugins/llm-integration_test.go +++ b/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, } diff --git a/protocols/strategies/ssh.go b/protocols/strategies/ssh.go index 6cb7af2..6c1ecf5 100644 --- a/protocols/strategies/ssh.go +++ b/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(),