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: Refactoring plugin:LLM honeypot custom prompt (#154)

Browse Source 
refactoring LLM honeypot custom prompt
This commit is contained in:
Mario Candela
2025-01-16 08:46:13 +01:00
committed by GitHub
parent c3d2ff885d
commit 99c7287c02
2 changed files with 53 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
plugins/llm-integration.go
View File

@ -96,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(),
@ -113,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(),
@ -214,18 +223,7 @@ func (llmHoneypot *LLMHoneypot) ExecuteModel(command string) (string, error) {
var err error
var prompt []Message
if llmHoneypot.CustomPrompt != "" {
prompt = append(prompt, Message{
Role: SYSTEM.String(),
Content: llmHoneypot.CustomPrompt,
})
prompt = append(prompt, Message{
Role: USER.String(),
Content: command,
})
} else {
prompt, err = buildPrompt(llmHoneypot.Histories, llmHoneypot.Protocol, command)
}
prompt, err = llmHoneypot.buildPrompt(command)
if err != nil {
return "", err

40
plugins/llm-integration_test.go
View File

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