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

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