chrissewell/beelzebub
1
0
Fork 0
mirror of https://github.com/mariocandela/beelzebub.git synced 2025-07-01 18:47:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
This commit is contained in:
Mario Candela
2025-03-09 13:17:04 +01:00
committed by GitHub
parent ef07ca1203
commit 933f02911b
5 changed files with 57 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
configurations/services/http-80.yaml
View File

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

1
docker-compose.yml
View File

@ -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"

5
plugins/llm-integration.go
View File

@ -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
}

33
plugins/llm-integration_test.go
View File

@ -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,
}

29
protocols/strategies/ssh.go
View File

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