From c3d2ff885d82fdff616f5f289d57b6309320ab48 Mon Sep 17 00:00:00 2001 From: Mario Candela Date: Tue, 14 Jan 2025 08:45:30 +0100 Subject: [PATCH] Feat: LLM Honeypot allow specifying the custom prompt #152 (#153) * implement new feature, custom prompt * Add doc for custom prompt --- README.md | 19 ++++++++++++ parser/configurations_parser.go | 1 + parser/configurations_parser_test.go | 2 ++ plugins/llm-integration.go | 27 ++++++++++++----- plugins/llm-integration_test.go | 45 ++++++++++++++++++++++++++++ protocols/strategies/http.go | 11 +++---- protocols/strategies/ssh.go | 22 +++++++------- 7 files changed, 105 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index eafa85a..90c1f1c 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/parser/configurations_parser.go b/parser/configurations_parser.go index f96fbc8..fd0131a 100644 --- a/parser/configurations_parser.go +++ b/parser/configurations_parser.go @@ -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 diff --git a/parser/configurations_parser_test.go b/parser/configurations_parser_test.go index 1b90922..fff8271 100644 --- a/parser/configurations_parser_test.go +++ b/parser/configurations_parser_test.go @@ -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) { diff --git a/plugins/llm-integration.go b/plugins/llm-integration.go index f63f44e..9e5888d 100644 --- a/plugins/llm-integration.go +++ b/plugins/llm-integration.go @@ -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 { @@ -211,8 +212,20 @@ 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) + 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) + } if err != nil { return "", err diff --git a/plugins/llm-integration_test.go b/plugins/llm-integration_test.go index a2aaa2e..572cab3 100644 --- a/plugins/llm-integration_test.go +++ b/plugins/llm-integration_test.go @@ -59,6 +59,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{ diff --git a/protocols/strategies/http.go b/protocols/strategies/http.go index e0a9d8c..e37d049 100644 --- a/protocols/strategies/http.go +++ b/protocols/strategies/http.go @@ -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) diff --git a/protocols/strategies/ssh.go b/protocols/strategies/ssh.go index aa1caff..22eb471 100644 --- a/protocols/strategies/ssh.go +++ b/protocols/strategies/ssh.go @@ -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)