diff --git a/README.md b/README.md index c16a781..ace2c58 100644 --- a/README.md +++ b/README.md @@ -12,9 +12,9 @@ Beelzebub is an advanced honeypot framework designed to provide a highly secure Beelzebub Logo -## OpenAI GPT Integration +## LLM Honeypot -Learn how to integrate Beelzebub with OpenAI GPT-3 by referring to our comprehensive guide on Medium: [Medium Article](https://medium.com/@mario.candela.personal/how-to-build-a-highly-effective-honeypot-with-beelzebub-and-chatgpt-a2f0f05b3e1) +Learn how to integrate Beelzebub with LLM OpenAI by referring to our comprehensive guide on Medium: [Medium Article](https://medium.com/@mario.candela.personal/how-to-build-a-highly-effective-honeypot-with-beelzebub-and-chatgpt-a2f0f05b3e1) ## Telegram Bot for Real-Time Attacks diff --git a/configurations/services/ssh-2222.yaml b/configurations/services/ssh-2222.yaml index 28bb235..58c1af6 100644 --- a/configurations/services/ssh-2222.yaml +++ b/configurations/services/ssh-2222.yaml @@ -4,10 +4,10 @@ address: ":2222" description: "SSH interactive ChatGPT" commands: - regex: "^(.+)$" - plugin: "OpenAIGPTLinuxTerminal" + plugin: "LLMHoneypot" serverVersion: "OpenSSH" serverName: "ubuntu" passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$" deadlineTimeoutSeconds: 60 plugin: - openAPIChatGPTSecretKey: "" \ No newline at end of file + openAISecretKey: "" \ No newline at end of file diff --git a/parser/configurations_parser.go b/parser/configurations_parser.go index 6b28272..a5e0dae 100644 --- a/parser/configurations_parser.go +++ b/parser/configurations_parser.go @@ -49,7 +49,7 @@ type Prometheus struct { } type Plugin struct { - OpenAPIChatGPTSecretKey string `yaml:"openAPIChatGPTSecretKey"` + OpenAISecretKey string `yaml:"openAISecretKey"` } // BeelzebubServiceConfiguration is the struct that contains the configurations of the honeypot service diff --git a/plugins/openai-gpt.go b/plugins/openai-gpt.go index 95019c3..f1feeda 100644 --- a/plugins/openai-gpt.go +++ b/plugins/openai-gpt.go @@ -3,31 +3,28 @@ package plugins import ( "encoding/json" "errors" - "fmt" "github.com/go-resty/resty/v2" - "strings" + "github.com/mariocandela/beelzebub/v3/tracer" log "github.com/sirupsen/logrus" ) const ( - promptVirtualizeLinuxTerminal = "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. Do not provide explanations or type commands unless explicitly instructed by the user. Remember previous commands and consider their effects on subsequent outputs.\n\nA:pwd\n\nQ:/home/user\n\n" - ChatGPTPluginName = "OpenAIGPTLinuxTerminal" - openAIGPTEndpoint = "https://api.openai.com/v1/completions" + systemPromptVirtualizeLinuxTerminal = "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. Do not provide explanations or type commands unless explicitly instructed by the user. Remember previous commands and consider their effects on subsequent outputs." + systemPromptVirtualizeHTTPServer = "You will act as an unsecure HTTP Server with multiple vulnerability like aws and git credentials stored into root http directory. The user will send HTTP requests, and you are to reply with what the server should show. Do not provide explanations or type commands unless explicitly instructed by the user." + ChatGPTPluginName = "LLMHoneypot" + openAIGPTEndpoint = "https://api.openai.com/v1/chat/completions" ) -type History struct { - Input, Output string -} - -type openAIGPTVirtualTerminal struct { - Histories []History +type openAIVirtualHoneypot struct { + Histories []Message openAIKey string client *resty.Client + protocol tracer.Protocol } type Choice struct { - Text string `json:"text"` + Message Message `json:"message"` Index int `json:"index"` Logprobs interface{} `json:"logprobs"` FinishReason string `json:"finish_reason"` @@ -47,61 +44,106 @@ type gptResponse struct { } type gptRequest struct { - Model string `json:"model"` - Prompt string `json:"prompt"` - Temperature int `json:"temperature"` - MaxTokens int `json:"max_tokens"` - TopP int `json:"top_p"` - FrequencyPenalty int `json:"frequency_penalty"` - PresencePenalty int `json:"presence_penalty"` - Stop []string `json:"stop"` + Model string `json:"model"` + Messages []Message `json:"messages"` } -func Init(history []History, openAIKey string) *openAIGPTVirtualTerminal { - return &openAIGPTVirtualTerminal{ +type Message struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type Role int + +const ( + SYSTEM Role = iota + USER + ASSISTANT +) + +func (role Role) String() string { + return [...]string{"system", "user", "assistant"}[role] +} + +func Init(history []Message, openAIKey string, protocol tracer.Protocol) *openAIVirtualHoneypot { + return &openAIVirtualHoneypot{ Histories: history, openAIKey: openAIKey, client: resty.New(), + protocol: protocol, } } -func buildPrompt(histories []History, command string) string { - var sb strings.Builder +func buildPrompt(histories []Message, protocol tracer.Protocol, command string) ([]Message, error) { + var messages []Message - sb.WriteString(promptVirtualizeLinuxTerminal) - - for _, history := range histories { - sb.WriteString(fmt.Sprintf("A:%s\n\nQ:%s\n\n", history.Input, history.Output)) + switch protocol { + case tracer.SSH: + messages = append(messages, Message{ + Role: SYSTEM.String(), + Content: systemPromptVirtualizeLinuxTerminal, + }) + messages = append(messages, Message{ + Role: USER.String(), + Content: "pwd", + }) + messages = append(messages, Message{ + Role: ASSISTANT.String(), + Content: "/home/user", + }) + for _, history := range histories { + messages = append(messages, history) + } + case tracer.HTTP: + messages = append(messages, Message{ + Role: SYSTEM.String(), + Content: systemPromptVirtualizeHTTPServer, + }) + messages = append(messages, Message{ + Role: USER.String(), + Content: "GET /index.html", + }) + messages = append(messages, Message{ + Role: ASSISTANT.String(), + Content: "Hello, World!", + }) + default: + return nil, errors.New("no prompt for protocol selected") } - // Append command to evaluate - sb.WriteString(fmt.Sprintf("A:%s\n\nQ:", command)) + messages = append(messages, Message{ + Role: USER.String(), + Content: command, + }) - return sb.String() + return messages, nil } -func (openAIGPTVirtualTerminal *openAIGPTVirtualTerminal) GetCompletions(command string) (string, error) { +func (openAIVirtualHoneypot *openAIVirtualHoneypot) GetCompletions(command string) (string, error) { + var err error + + prompt, err := buildPrompt(openAIVirtualHoneypot.Histories, openAIVirtualHoneypot.protocol, command) + + if err != nil { + return "", err + } + requestJson, err := json.Marshal(gptRequest{ - Model: "gpt-3.5-turbo-instruct", - Prompt: buildPrompt(openAIGPTVirtualTerminal.Histories, command), - Temperature: 0, - MaxTokens: 100, - TopP: 1, - FrequencyPenalty: 0, - PresencePenalty: 0, - Stop: []string{"\n"}, + Model: "gpt-4", //"gpt-3.5-turbo", + Messages: prompt, }) if err != nil { return "", err } - if openAIGPTVirtualTerminal.openAIKey == "" { + if openAIVirtualHoneypot.openAIKey == "" { return "", errors.New("openAIKey is empty") } - response, err := openAIGPTVirtualTerminal.client.R(). + log.Debug(string(requestJson)) + response, err := openAIVirtualHoneypot.client.R(). SetHeader("Content-Type", "application/json"). SetBody(requestJson). - SetAuthToken(openAIGPTVirtualTerminal.openAIKey). + SetAuthToken(openAIVirtualHoneypot.openAIKey). SetResult(&gptResponse{}). Post(openAIGPTEndpoint) @@ -113,5 +155,5 @@ func (openAIGPTVirtualTerminal *openAIGPTVirtualTerminal) GetCompletions(command return "", errors.New("no choices") } - return response.Result().(*gptResponse).Choices[0].Text, nil + return response.Result().(*gptResponse).Choices[0].Message.Content, nil } diff --git a/plugins/openai-gpt_test.go b/plugins/openai-gpt_test.go index 89352c8..24dedc6 100644 --- a/plugins/openai-gpt_test.go +++ b/plugins/openai-gpt_test.go @@ -3,58 +3,63 @@ package plugins import ( "github.com/go-resty/resty/v2" "github.com/jarcoal/httpmock" + "github.com/mariocandela/beelzebub/v3/tracer" "github.com/stretchr/testify/assert" "net/http" "testing" ) +const SystemPromptLen = 4 + func TestBuildPromptEmptyHistory(t *testing.T) { //Given - var histories []History + var histories []Message command := "pwd" //When - prompt := buildPrompt(histories, command) + prompt, err := buildPrompt(histories, tracer.SSH, command) //Then - assert.Equal(t, - "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. Do not provide explanations or type commands unless explicitly instructed by the user. Remember previous commands and consider their effects on subsequent outputs.\n\nA:pwd\n\nQ:/home/user\n\nA:pwd\n\nQ:", - prompt) + assert.Nil(t, err) + assert.Equal(t, SystemPromptLen, len(prompt)) } func TestBuildPromptWithHistory(t *testing.T) { //Given - var histories = []History{ + var histories = []Message{ { - Input: "cat hello.txt", - Output: "world", - }, - { - Input: "echo 1234", - Output: "1234", + Role: "cat hello.txt", + Content: "world", }, } command := "pwd" //When - prompt := buildPrompt(histories, command) + prompt, err := buildPrompt(histories, tracer.SSH, command) //Then - assert.Equal(t, - "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. Do not provide explanations or type commands unless explicitly instructed by the user. Remember previous commands and consider their effects on subsequent outputs.\n\nA:pwd\n\nQ:/home/user\n\nA:cat hello.txt\n\nQ:world\n\nA:echo 1234\n\nQ:1234\n\nA:pwd\n\nQ:", - prompt) + assert.Nil(t, err) + assert.Equal(t, SystemPromptLen+1, len(prompt)) } func TestBuildGetCompletionsFailValidation(t *testing.T) { - openAIGPTVirtualTerminal := Init(make([]History, 0), "") + openAIGPTVirtualTerminal := Init(make([]Message, 0), "", tracer.SSH) _, err := openAIGPTVirtualTerminal.GetCompletions("test") assert.Equal(t, "openAIKey is empty", err.Error()) } -func TestBuildGetCompletionsWithResults(t *testing.T) { +func TestBuildGetCompletionsFailValidationStrategyType(t *testing.T) { + openAIGPTVirtualTerminal := Init(make([]Message, 0), "", tracer.TCP) + + _, err := openAIGPTVirtualTerminal.GetCompletions("test") + + assert.Equal(t, "no prompt for protocol selected", err.Error()) +} + +func TestBuildGetCompletionsSSHWithResults(t *testing.T) { client := resty.New() httpmock.ActivateNonDefault(client.GetClient()) defer httpmock.DeactivateAndReset() @@ -65,7 +70,10 @@ func TestBuildGetCompletionsWithResults(t *testing.T) { resp, err := httpmock.NewJsonResponse(200, &gptResponse{ Choices: []Choice{ { - Text: "prova.txt", + Message: Message{ + Role: SYSTEM.String(), + Content: "prova.txt", + }, }, }, }) @@ -76,7 +84,7 @@ func TestBuildGetCompletionsWithResults(t *testing.T) { }, ) - openAIGPTVirtualTerminal := Init(make([]History, 0), "sdjdnklfjndslkjanfk") + openAIGPTVirtualTerminal := Init(make([]Message, 0), "sdjdnklfjndslkjanfk", tracer.SSH) openAIGPTVirtualTerminal.client = client //When @@ -87,7 +95,7 @@ func TestBuildGetCompletionsWithResults(t *testing.T) { assert.Equal(t, "prova.txt", str) } -func TestBuildGetCompletionsWithoutResults(t *testing.T) { +func TestBuildGetCompletionsSSHWithoutResults(t *testing.T) { client := resty.New() httpmock.ActivateNonDefault(client.GetClient()) defer httpmock.DeactivateAndReset() @@ -105,7 +113,7 @@ func TestBuildGetCompletionsWithoutResults(t *testing.T) { }, ) - openAIGPTVirtualTerminal := Init(make([]History, 0), "sdjdnklfjndslkjanfk") + openAIGPTVirtualTerminal := Init(make([]Message, 0), "sdjdnklfjndslkjanfk", tracer.SSH) openAIGPTVirtualTerminal.client = client //When @@ -114,3 +122,67 @@ func TestBuildGetCompletionsWithoutResults(t *testing.T) { //Then assert.Equal(t, "no choices", err.Error()) } + +func TestBuildGetCompletionsHTTPWithResults(t *testing.T) { + client := resty.New() + httpmock.ActivateNonDefault(client.GetClient()) + defer httpmock.DeactivateAndReset() + + // Given + httpmock.RegisterResponder("POST", openAIGPTEndpoint, + func(req *http.Request) (*http.Response, error) { + resp, err := httpmock.NewJsonResponse(200, &gptResponse{ + 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 + }, + ) + + openAIGPTVirtualTerminal := Init(make([]Message, 0), "sdjdnklfjndslkjanfk", tracer.HTTP) + openAIGPTVirtualTerminal.client = client + + //When + str, err := openAIGPTVirtualTerminal.GetCompletions("GET /.aws/credentials") + + //Then + assert.Nil(t, err) + assert.Equal(t, "[default]\nregion = us-west-2\noutput = json", str) +} + +func TestBuildGetCompletionsHTTPWithoutResults(t *testing.T) { + client := resty.New() + httpmock.ActivateNonDefault(client.GetClient()) + defer httpmock.DeactivateAndReset() + + // Given + httpmock.RegisterResponder("POST", openAIGPTEndpoint, + func(req *http.Request) (*http.Response, error) { + resp, err := httpmock.NewJsonResponse(200, &gptResponse{ + Choices: []Choice{}, + }) + if err != nil { + return httpmock.NewStringResponse(500, ""), nil + } + return resp, nil + }, + ) + + openAIGPTVirtualTerminal := Init(make([]Message, 0), "sdjdnklfjndslkjanfk", tracer.HTTP) + openAIGPTVirtualTerminal.client = client + + //When + _, err := openAIGPTVirtualTerminal.GetCompletions("GET /.aws/credentials") + + //Then + assert.Equal(t, "no choices", err.Error()) +} diff --git a/protocols/strategies/http.go b/protocols/strategies/http.go index fcaa714..8f38f3e 100644 --- a/protocols/strategies/http.go +++ b/protocols/strategies/http.go @@ -3,6 +3,7 @@ package strategies import ( "fmt" "github.com/mariocandela/beelzebub/v3/parser" + "github.com/mariocandela/beelzebub/v3/plugins" "github.com/mariocandela/beelzebub/v3/tracer" "io" "net/http" @@ -31,8 +32,24 @@ func (httpStrategy HTTPStrategy) Init(beelzebubServiceConfiguration parser.Beelz } if matched { + responseHTTPBody := command.Handler + + if command.Plugin == plugins.ChatGPTPluginName { + openAIGPTVirtualTerminal := plugins.Init(make([]plugins.Message, 0), beelzebubServiceConfiguration.Plugin.OpenAISecretKey, tracer.HTTP) + + command := fmt.Sprintf("%s %s", request.Method, request.RequestURI) + + if completions, err := openAIGPTVirtualTerminal.GetCompletions(command); err != nil { + log.Errorf("Error GetCompletions: %s, %s", command, err.Error()) + responseHTTPBody = "404 Not Found!" + } else { + responseHTTPBody = completions + } + + } + setResponseHeaders(responseWriter, command.Headers, command.StatusCode) - fmt.Fprintf(responseWriter, command.Handler) + fmt.Fprintf(responseWriter, responseHTTPBody) break } } diff --git a/protocols/strategies/ssh.go b/protocols/strategies/ssh.go index 779d43d..80003ae 100644 --- a/protocols/strategies/ssh.go +++ b/protocols/strategies/ssh.go @@ -42,7 +42,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze }) term := terminal.NewTerminal(sess, buildPrompt(sess.User(), beelzebubServiceConfiguration.ServerName)) - var histories []plugins.History + var histories []plugins.Message for { commandInput, err := term.ReadLine() if err != nil { @@ -63,7 +63,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze commandOutput := command.Handler if command.Plugin == plugins.ChatGPTPluginName { - openAIGPTVirtualTerminal := plugins.Init(histories, beelzebubServiceConfiguration.Plugin.OpenAPIChatGPTSecretKey) + openAIGPTVirtualTerminal := plugins.Init(histories, beelzebubServiceConfiguration.Plugin.OpenAISecretKey, tracer.SSH) if commandOutput, err = openAIGPTVirtualTerminal.GetCompletions(commandInput); err != nil { log.Errorf("Error GetCompletions: %s, %s", commandInput, err.Error()) @@ -71,7 +71,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze } } - histories = append(histories, plugins.History{Input: commandInput, Output: commandOutput}) + histories = append(histories, plugins.Message{Role: plugins.USER.String(), Content: commandOutput}) term.Write(append([]byte(commandOutput), '\n'))