Skip to content
Permalink

Comparing changes

Choose two branches to see what’s changed or to start a new pull request. If you need to, you can also or learn more about diff comparisons.

Open a pull request

Create a new pull request by comparing changes across two branches. If you need to, you can also . Learn more about diff comparisons here.
base repository: mariocandela/beelzebub
Failed to load repositories. Confirm that selected base ref is valid, then try again.
Loading
base: v3.2.9
Choose a base ref
...
head repository: mariocandela/beelzebub
Failed to load repositories. Confirm that selected head ref is valid, then try again.
Loading
compare: main
Choose a head ref
  • 2 commits
  • 7 files changed
  • 1 contributor

Commits on Jan 14, 2025

  1. Feat: LLM Honeypot allow specifying the custom prompt #152 (#153)

    * implement new feature, custom prompt
    
    * Add doc for custom prompt
    mariocandela authored Jan 14, 2025
    Copy the full SHA
    c3d2ff8 View commit details

Commits on Jan 16, 2025

  1. Feat: Refactoring plugin:LLM honeypot custom prompt (#154)

    refactoring LLM honeypot custom prompt
    mariocandela authored Jan 16, 2025
    Copy the full SHA
    99c7287 View commit details
Showing with 146 additions and 29 deletions.
  1. +19 −0 README.md
  2. +1 −0 parser/configurations_parser.go
  3. +2 −0 parser/configurations_parser_test.go
  4. +23 −12 plugins/llm-integration.go
  5. +83 −2 plugins/llm-integration_test.go
  6. +6 −5 protocols/strategies/http.go
  7. +12 −10 protocols/strategies/ssh.go
19 changes: 19 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -250,6 +250,25 @@ plugin:
llmModel: "llama3"
host: "http://example.com/api/chat" #default http://localhost:11434/api/chat
```
Example with custom prompt:
```yaml
apiVersion: "v1"
protocol: "ssh"
address: ":2222"
description: "SSH interactive OpenAI GPT-4"
commands:
- regex: "^(.+)$"
plugin: "LLMHoneypot"
serverVersion: "OpenSSH"
serverName: "ubuntu"
passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
deadlineTimeoutSeconds: 60
plugin:
llmModel: "gpt4-o"
openAISecretKey: "sk-proj-123456"
prompt: "You will act as an Ubuntu Linux terminal. The user will type commands, and you are to reply with what the terminal should show. Your responses must be contained within a single code block."
```
###### SSH Honeypot on Port 22
1 change: 1 addition & 0 deletions parser/configurations_parser.go
Original file line number Diff line number Diff line change
@@ -52,6 +52,7 @@ type Plugin struct {
OpenAISecretKey string `yaml:"openAISecretKey"`
Host string `yaml:"host"`
LLMModel string `yaml:"llmModel"`
Prompt string `yaml:"prompt"`
}

// BeelzebubServiceConfiguration is the struct that contains the configurations of the honeypot service
2 changes: 2 additions & 0 deletions parser/configurations_parser_test.go
Original file line number Diff line number Diff line change
@@ -58,6 +58,7 @@ plugin:
openAISecretKey: "qwerty"
llmModel: "llama3"
host: "localhost:1563"
prompt: "hello world"
`)
return beelzebubServiceConfiguration, nil
}
@@ -133,6 +134,7 @@ func TestReadConfigurationsServicesValid(t *testing.T) {
assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.OpenAISecretKey, "qwerty")
assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.LLMModel, "llama3")
assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.Host, "localhost:1563")
assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.Prompt, "hello world")
}

func TestGelAllFilesNameByDirName(t *testing.T) {
35 changes: 23 additions & 12 deletions plugins/llm-integration.go
Original file line number Diff line number Diff line change
@@ -19,12 +19,13 @@ const (
)

type LLMHoneypot struct {
Histories []Message
OpenAIKey string
client *resty.Client
Protocol tracer.Protocol
Model LLMModel
Host string
Histories []Message
OpenAIKey string
client *resty.Client
Protocol tracer.Protocol
Model LLMModel
Host string
CustomPrompt string
}

type Choice struct {
@@ -95,14 +96,19 @@ func InitLLMHoneypot(config LLMHoneypot) *LLMHoneypot {
return &config
}

func buildPrompt(histories []Message, protocol tracer.Protocol, command string) ([]Message, error) {
func (llmHoneypot *LLMHoneypot) buildPrompt(command string) ([]Message, error) {
var messages []Message
var prompt string

switch protocol {
switch llmHoneypot.Protocol {
case tracer.SSH:
prompt = systemPromptVirtualizeLinuxTerminal
if llmHoneypot.CustomPrompt != "" {
prompt = llmHoneypot.CustomPrompt
}
messages = append(messages, Message{
Role: SYSTEM.String(),
Content: systemPromptVirtualizeLinuxTerminal,
Content: prompt,
})
messages = append(messages, Message{
Role: USER.String(),
@@ -112,13 +118,17 @@ func buildPrompt(histories []Message, protocol tracer.Protocol, command string)
Role: ASSISTANT.String(),
Content: "/home/user",
})
for _, history := range histories {
for _, history := range llmHoneypot.Histories {
messages = append(messages, history)
}
case tracer.HTTP:
prompt = systemPromptVirtualizeHTTPServer
if llmHoneypot.CustomPrompt != "" {
prompt = llmHoneypot.CustomPrompt
}
messages = append(messages, Message{
Role: SYSTEM.String(),
Content: systemPromptVirtualizeHTTPServer,
Content: prompt,
})
messages = append(messages, Message{
Role: USER.String(),
@@ -211,8 +221,9 @@ func (llmHoneypot *LLMHoneypot) ollamaCaller(messages []Message) (string, error)

func (llmHoneypot *LLMHoneypot) ExecuteModel(command string) (string, error) {
var err error
var prompt []Message

prompt, err := buildPrompt(llmHoneypot.Histories, llmHoneypot.Protocol, command)
prompt, err = llmHoneypot.buildPrompt(command)

if err != nil {
return "", err
85 changes: 83 additions & 2 deletions plugins/llm-integration_test.go
Original file line number Diff line number Diff line change
@@ -16,8 +16,13 @@ func TestBuildPromptEmptyHistory(t *testing.T) {
var histories []Message
command := "pwd"

honeypot := LLMHoneypot{
Histories: histories,
Protocol: tracer.SSH,
}

//When
prompt, err := buildPrompt(histories, tracer.SSH, command)
prompt, err := honeypot.buildPrompt(command)

//Then
assert.Nil(t, err)
@@ -35,14 +40,45 @@ func TestBuildPromptWithHistory(t *testing.T) {

command := "pwd"

honeypot := LLMHoneypot{
Histories: histories,
Protocol: tracer.SSH,
}

//When
prompt, err := buildPrompt(histories, tracer.SSH, command)
prompt, err := honeypot.buildPrompt(command)

//Then
assert.Nil(t, err)
assert.Equal(t, SystemPromptLen+1, len(prompt))
}

func TestBuildPromptWithCustomPrompt(t *testing.T) {
//Given
var histories = []Message{
{
Role: "cat hello.txt",
Content: "world",
},
}

command := "pwd"

honeypot := LLMHoneypot{
Histories: histories,
Protocol: tracer.SSH,
CustomPrompt: "act as calculator",
}

//When
prompt, err := honeypot.buildPrompt(command)

//Then
assert.Nil(t, err)
assert.Equal(t, prompt[0].Content, "act as calculator")
assert.Equal(t, prompt[0].Role, SYSTEM.String())
}

func TestBuildExecuteModelFailValidation(t *testing.T) {

llmHoneypot := LLMHoneypot{
@@ -59,6 +95,51 @@ func TestBuildExecuteModelFailValidation(t *testing.T) {
assert.Equal(t, "openAIKey is empty", err.Error())
}

func TestBuildExecuteModelWithCustomPrompt(t *testing.T) {
client := resty.New()
httpmock.ActivateNonDefault(client.GetClient())
defer httpmock.DeactivateAndReset()

// Given
httpmock.RegisterMatcherResponder("POST", openAIGPTEndpoint,
httpmock.BodyContainsString("hello world"),
func(req *http.Request) (*http.Response, error) {
resp, err := httpmock.NewJsonResponse(200, &Response{
Choices: []Choice{
{
Message: Message{
Role: SYSTEM.String(),
Content: "[default]\nregion = us-west-2\noutput = json",
},
},
},
})
if err != nil {
return httpmock.NewStringResponse(500, ""), nil
}
return resp, nil
},
)

llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
OpenAIKey: "sdjdnklfjndslkjanfk",
Protocol: tracer.HTTP,
Model: GPT4O,
CustomPrompt: "hello world",
}

openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
openAIGPTVirtualTerminal.client = client

//When
str, err := openAIGPTVirtualTerminal.ExecuteModel("GET /.aws/credentials")

//Then
assert.Nil(t, err)
assert.Equal(t, "[default]\nregion = us-west-2\noutput = json", str)
}

func TestBuildExecuteModelFailValidationStrategyType(t *testing.T) {

llmHoneypot := LLMHoneypot{
11 changes: 6 additions & 5 deletions protocols/strategies/http.go
Original file line number Diff line number Diff line change
@@ -45,11 +45,12 @@ func (httpStrategy HTTPStrategy) Init(beelzebubServiceConfiguration parser.Beelz
}

llmHoneypot := plugins.LLMHoneypot{
Histories: make([]plugins.Message, 0),
OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
Protocol: tracer.HTTP,
Host: beelzebubServiceConfiguration.Plugin.Host,
Model: llmModel,
Histories: make([]plugins.Message, 0),
OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
Protocol: tracer.HTTP,
Host: beelzebubServiceConfiguration.Plugin.Host,
Model: llmModel,
CustomPrompt: beelzebubServiceConfiguration.Plugin.Prompt,
}

llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot)
22 changes: 12 additions & 10 deletions protocols/strategies/ssh.go
Original file line number Diff line number Diff line change
@@ -52,11 +52,12 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
}

llmHoneypot := plugins.LLMHoneypot{
Histories: make([]plugins.Message, 0),
OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
Protocol: tracer.SSH,
Host: beelzebubServiceConfiguration.Plugin.Host,
Model: llmModel,
Histories: make([]plugins.Message, 0),
OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
Protocol: tracer.SSH,
Host: beelzebubServiceConfiguration.Plugin.Host,
Model: llmModel,
CustomPrompt: beelzebubServiceConfiguration.Plugin.Prompt,
}

llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot)
@@ -137,11 +138,12 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
}

llmHoneypot := plugins.LLMHoneypot{
Histories: histories,
OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
Protocol: tracer.SSH,
Host: beelzebubServiceConfiguration.Plugin.Host,
Model: llmModel,
Histories: histories,
OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
Protocol: tracer.SSH,
Host: beelzebubServiceConfiguration.Plugin.Host,
Model: llmModel,
CustomPrompt: beelzebubServiceConfiguration.Plugin.Prompt,
}

llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot)