From 99c7287c0239b48f2ed03ae45f4efd2dbd08112a Mon Sep 17 00:00:00 2001 From: Mario Candela Date: Thu, 16 Jan 2025 08:46:13 +0100 Subject: [PATCH] Feat: Refactoring plugin:LLM honeypot custom prompt (#154) refactoring LLM honeypot custom prompt --- plugins/llm-integration.go | 32 +++++++++++++------------- plugins/llm-integration_test.go | 40 +++++++++++++++++++++++++++++++-- 2 files changed, 53 insertions(+), 19 deletions(-) diff --git a/plugins/llm-integration.go b/plugins/llm-integration.go index 9e5888d..99ee29b 100644 --- a/plugins/llm-integration.go +++ b/plugins/llm-integration.go @@ -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 diff --git a/plugins/llm-integration_test.go b/plugins/llm-integration_test.go index 572cab3..15e2516 100644 --- a/plugins/llm-integration_test.go +++ b/plugins/llm-integration_test.go @@ -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{