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

Feature: add support for llama, refactor yaml service interface (#115)

Browse Source 
* refactor and add llama LMM support

* update readme

* improve code coverage
This commit is contained in:
Mario Candela
2024-07-21 20:11:18 +02:00
committed by GitHub
parent 0af1a05ae9
commit 2088163b54
11 changed files with 511 additions and 275 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
README.md
View File

@ -104,7 +104,8 @@ $ make test.dependencies.down
Beelzebub offers a wide range of features to enhance your honeypot environment: Beelzebub offers a wide range of features to enhance your honeypot environment:
- OpenAI Generative Pre-trained Transformer act as Linux virtualization - Support for Ollama
- Support for OpenAI
- SSH Honeypot - SSH Honeypot
- HTTP Honeypot - HTTP Honeypot
- TCP Honeypot - TCP Honeypot
@ -210,7 +211,9 @@ commands:
#### Example SSH Honeypot #### Example SSH Honeypot
###### Honeypot with GPT-3 on Port 2222 ###### Honeypot LLM Honeypots
Example with OpenAI GPT-4:
```yaml ```yaml
apiVersion: "v1" apiVersion: "v1"
@ -219,13 +222,54 @@ address: ":2222"
description: "SSH interactive ChatGPT" description: "SSH interactive ChatGPT"
commands: commands:
- regex: "^(.+)$" - regex: "^(.+)$"
plugin: "OpenAIGPTLinuxTerminal" plugin: "LLMHoneypot"
serverVersion: "OpenSSH" serverVersion: "OpenSSH"
serverName: "ubuntu" serverName: "ubuntu"
passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$" passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
deadlineTimeoutSeconds: 60 deadlineTimeoutSeconds: 60
plugin: plugin:
openAISecretKey: "Your OpenAI Secret Key" llmModel: "gpt4-o"
openAISecretKey: "sk-proj-123456"
```
###### Honeypot LLM Honeypots
Example with OpenAI GPT-4:
```yaml
apiVersion: "v1"
protocol: "ssh"
address: ":2222"
description: "SSH interactive ChatGPT"
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"
```
Example with Ollama Llama3:
```yaml
apiVersion: "v1"
protocol: "ssh"
address: ":2222"
description: "SSH interactive ChatGPT"
commands:
- regex: "^(.+)$"
plugin: "LLMHoneypot"
serverVersion: "OpenSSH"
serverName: "ubuntu"
passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
deadlineTimeoutSeconds: 60
plugin:
llmModel: "llama3"
host: "http://example.com/api/chat" #default http://localhost:11434/api/chat
``` ```
###### SSH Honeypot on Port 22 ###### SSH Honeypot on Port 22

42
configurations/services/http-80.yaml
View File

@ -18,41 +18,9 @@ commands:
- "Server: Apache/2.4.53 (Debian)" - "Server: Apache/2.4.53 (Debian)"
- "X-Powered-By: PHP/7.4.29" - "X-Powered-By: PHP/7.4.29"
statusCode: 200 statusCode: 200
- regex: "^(/wp-login.php|/wp-admin)$"
handler:
<html>
<header>
<title>Wordpress 6 test page</title>
</header>
<body>
<form action="" method="post">
<label for="uname"><b>Username</b></label>
<input type="text" placeholder="Enter Username" name="uname" required>
<label for="psw"><b>Password</b></label>
<input type="password" placeholder="Enter Password" name="psw" required>
<button type="submit">Login</button>
</form>
</body>
</html>
headers:
- "Content-Type: text/html"
- "Server: Apache/2.4.53 (Debian)"
- "X-Powered-By: PHP/7.4.29"
statusCode: 200
- regex: "^.*$" - regex: "^.*$"
handler: plugin: "LLMHoneypot"
<html> statusCode: 200
<header> plugin:
<title>404</title> llmModel: "gpt4-o"
</header> openAISecretKey: "sk-proj-123456"
<body>
<h1>Not found!</h1>
</body>
</html>
headers:
- "Content-Type: text/html"
- "Server: Apache/2.4.53 (Debian)"
- "X-Powered-By: PHP/7.4.29"
statusCode: 404

4
configurations/services/ssh-2222.yaml
View File

@ -8,6 +8,6 @@ commands:
serverVersion: "OpenSSH" serverVersion: "OpenSSH"
serverName: "ubuntu" serverName: "ubuntu"
passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$" passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
deadlineTimeoutSeconds: 60 deadlineTimeoutSeconds: 6000
plugin: plugin:
openAISecretKey: "" llmModel: "llama3"

14
parser/configurations_parser.go
View File

@ -3,6 +3,7 @@ package parser
import ( import (
"fmt" "fmt"
"github.com/mariocandela/beelzebub/v3/plugins"
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
@ -50,6 +51,19 @@ type Prometheus struct {
type Plugin struct { type Plugin struct {
OpenAISecretKey string `yaml:"openAISecretKey"` OpenAISecretKey string `yaml:"openAISecretKey"`
Host string `yaml:"host"`
LLMModel string `yaml:"llmModel"`
}
func FromString(llmModel string) (plugins.LLMModel, error) {
switch llmModel {
case "llama3":
return plugins.LLAMA3, nil
case "gpt4-o":
return plugins.GPT4O, nil
default:
return -1, fmt.Errorf("model %s not found", llmModel)
}
} }
// BeelzebubServiceConfiguration is the struct that contains the configurations of the honeypot service // BeelzebubServiceConfiguration is the struct that contains the configurations of the honeypot service

26
parser/configurations_parser_test.go
View File

@ -2,6 +2,7 @@ package parser
import ( import (
"errors" "errors"
"github.com/mariocandela/beelzebub/v3/plugins"
"os" "os"
"testing" "testing"
@ -53,7 +54,12 @@ commands:
- regex: "wp-admin" - regex: "wp-admin"
handler: "login" handler: "login"
headers: headers:
- "Content-Type: text/html"`) - "Content-Type: text/html"
plugin:
openAISecretKey: "qwerty"
llmModel: "llama3"
host: "localhost:1563"
`)
return beelzebubServiceConfiguration, nil return beelzebubServiceConfiguration, nil
} }
@ -112,10 +118,10 @@ func TestReadConfigurationsServicesValid(t *testing.T) {
configurationsParser.gelAllFilesNameByDirNameDependency = mockReadDirValid configurationsParser.gelAllFilesNameByDirNameDependency = mockReadDirValid
beelzebubServicesConfiguration, err := configurationsParser.ReadConfigurationsServices() beelzebubServicesConfiguration, err := configurationsParser.ReadConfigurationsServices()
assert.Nil(t, err)
firstBeelzebubServiceConfiguration := beelzebubServicesConfiguration[0] firstBeelzebubServiceConfiguration := beelzebubServicesConfiguration[0]
assert.Nil(t, err)
assert.Equal(t, firstBeelzebubServiceConfiguration.Protocol, "http") assert.Equal(t, firstBeelzebubServiceConfiguration.Protocol, "http")
assert.Equal(t, firstBeelzebubServiceConfiguration.ApiVersion, "v1") assert.Equal(t, firstBeelzebubServiceConfiguration.ApiVersion, "v1")
assert.Equal(t, firstBeelzebubServiceConfiguration.Address, ":8080") assert.Equal(t, firstBeelzebubServiceConfiguration.Address, ":8080")
@ -125,6 +131,9 @@ func TestReadConfigurationsServicesValid(t *testing.T) {
assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].Handler, "login") assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].Handler, "login")
assert.Equal(t, len(firstBeelzebubServiceConfiguration.Commands[0].Headers), 1) assert.Equal(t, len(firstBeelzebubServiceConfiguration.Commands[0].Headers), 1)
assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].Headers[0], "Content-Type: text/html") assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].Headers[0], "Content-Type: text/html")
assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.OpenAISecretKey, "qwerty")
assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.LLMModel, "llama3")
assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.Host, "localhost:1563")
} }
func TestGelAllFilesNameByDirName(t *testing.T) { func TestGelAllFilesNameByDirName(t *testing.T) {
@ -177,3 +186,16 @@ func TestReadFileBytesByFilePath(t *testing.T) {
assert.Equal(t, "", string(bytes)) assert.Equal(t, "", string(bytes))
} }
func TestFromString(t *testing.T) {
model, err := FromString("llama3")
assert.Nil(t, err)
assert.Equal(t, plugins.LLAMA3, model)
model, err = FromString("gpt4-o")
assert.Nil(t, err)
assert.Equal(t, plugins.GPT4O, model)
model, err = FromString("beelzebub-model")
assert.Errorf(t, err, "model beelzebub-model not found")
}

1
plugins/beelzebub-cloud.go
View File

@ -36,7 +36,6 @@ func (beelzebubCloud *beelzebubCloud) SendEvent(event tracer.Event) (bool, error
SetHeader("Content-Type", "application/json"). SetHeader("Content-Type", "application/json").
SetBody(requestJson). SetBody(requestJson).
SetHeader("Authorization", beelzebubCloud.AuthToken). SetHeader("Authorization", beelzebubCloud.AuthToken).
SetResult(&gptResponse{}).
Post(beelzebubCloud.URI) Post(beelzebubCloud.URI)
log.Debug(response) log.Debug(response)

124
plugins/openai-gpt.go → plugins/llm-integration.go
View File

@ -12,30 +12,33 @@ import (
const ( const (
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. Your entire response/output is going to consist of a simple text with \n for new line, and you will NOT wrap it within string md markers" 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. Your entire response/output is going to consist of a simple text with \n for new line, and you will NOT wrap it within string md markers"
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." 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" LLMPluginName = "LLMHoneypot"
openAIGPTEndpoint = "https://api.openai.com/v1/chat/completions" openAIGPTEndpoint = "https://api.openai.com/v1/chat/completions"
ollamaEndpoint = "http://localhost:11434/api/chat"
) )
type openAIVirtualHoneypot struct { type LLMHoneypot struct {
Histories []Message Histories []Message
openAIKey string OpenAIKey string
client *resty.Client client *resty.Client
protocol tracer.Protocol Protocol tracer.Protocol
Model LLMModel
Host string
} }
type Choice struct { type Choice struct {
Message Message `json:"message"` Message Message `json:"message"`
Index int `json:"index"` Index int `json:"index"`
Logprobs interface{} `json:"logprobs"` FinishReason string `json:"finish_reason"`
FinishReason string `json:"finish_reason"`
} }
type gptResponse struct { type Response struct {
ID string `json:"id"` ID string `json:"id"`
Object string `json:"object"` Object string `json:"object"`
Created int `json:"created"` Created int `json:"created"`
Model string `json:"model"` Model string `json:"model"`
Choices []Choice `json:"choices"` Choices []Choice `json:"choices"`
Message Message `json:"message"`
Usage struct { Usage struct {
PromptTokens int `json:"prompt_tokens"` PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"` CompletionTokens int `json:"completion_tokens"`
@ -43,9 +46,10 @@ type gptResponse struct {
} `json:"usage"` } `json:"usage"`
} }
type gptRequest struct { type Request struct {
Model string `json:"model"` Model string `json:"model"`
Messages []Message `json:"messages"` Messages []Message `json:"messages"`
Stream bool `json:"stream"`
} }
type Message struct { type Message struct {
@ -65,13 +69,18 @@ func (role Role) String() string {
return [...]string{"system", "user", "assistant"}[role] return [...]string{"system", "user", "assistant"}[role]
} }
func Init(history []Message, openAIKey string, protocol tracer.Protocol) *openAIVirtualHoneypot { type LLMModel int
return &openAIVirtualHoneypot{
Histories: history, const (
openAIKey: openAIKey, LLAMA3 LLMModel = iota
client: resty.New(), GPT4O
protocol: protocol, )
}
func InitLLMHoneypot(config LLMHoneypot) *LLMHoneypot {
// Inject the dependencies
config.client = resty.New()
return &config
} }
func buildPrompt(histories []Message, protocol tracer.Protocol, command string) ([]Message, error) { func buildPrompt(histories []Message, protocol tracer.Protocol, command string) ([]Message, error) {
@ -118,42 +127,91 @@ func buildPrompt(histories []Message, protocol tracer.Protocol, command string)
return messages, nil return messages, nil
} }
func (openAIVirtualHoneypot *openAIVirtualHoneypot) GetCompletions(command string) (string, error) { func (llmHoneypot *LLMHoneypot) openAICaller(messages []Message) (string, error) {
var err error var err error
prompt, err := buildPrompt(openAIVirtualHoneypot.Histories, openAIVirtualHoneypot.protocol, command) requestJson, err := json.Marshal(Request{
if err != nil {
return "", err
}
requestJson, err := json.Marshal(gptRequest{
Model: "gpt-4o", Model: "gpt-4o",
Messages: prompt, Messages: messages,
Stream: false,
}) })
if err != nil { if err != nil {
return "", err return "", err
} }
if openAIVirtualHoneypot.openAIKey == "" { if llmHoneypot.OpenAIKey == "" {
return "", errors.New("openAIKey is empty") return "", errors.New("openAIKey is empty")
} }
if llmHoneypot.Host == "" {
llmHoneypot.Host = openAIGPTEndpoint
}
log.Debug(string(requestJson)) log.Debug(string(requestJson))
response, err := openAIVirtualHoneypot.client.R(). response, err := llmHoneypot.client.R().
SetHeader("Content-Type", "application/json"). SetHeader("Content-Type", "application/json").
SetBody(requestJson). SetBody(requestJson).
SetAuthToken(openAIVirtualHoneypot.openAIKey). SetAuthToken(llmHoneypot.OpenAIKey).
SetResult(&gptResponse{}). SetResult(&Response{}).
Post(openAIGPTEndpoint) Post(llmHoneypot.Host)
if err != nil { if err != nil {
return "", err return "", err
} }
log.Debug(response) log.Debug(response)
if len(response.Result().(*gptResponse).Choices) == 0 { if len(response.Result().(*Response).Choices) == 0 {
return "", errors.New("no choices") return "", errors.New("no choices")
} }
return response.Result().(*gptResponse).Choices[0].Message.Content, nil return response.Result().(*Response).Choices[0].Message.Content, nil
}
func (llmHoneypot *LLMHoneypot) ollamaCaller(messages []Message) (string, error) {
var err error
requestJson, err := json.Marshal(Request{
Model: "llama3",
Messages: messages,
Stream: false,
})
if err != nil {
return "", err
}
if llmHoneypot.Host == "" {
llmHoneypot.Host = ollamaEndpoint
}
log.Debug(string(requestJson))
response, err := llmHoneypot.client.R().
SetHeader("Content-Type", "application/json").
SetBody(requestJson).
SetResult(&Response{}).
Post(llmHoneypot.Host)
if err != nil {
return "", err
}
log.Debug(response)
return response.Result().(*Response).Message.Content, nil
}
func (llmHoneypot *LLMHoneypot) ExecuteModel(command string) (string, error) {
var err error
prompt, err := buildPrompt(llmHoneypot.Histories, llmHoneypot.Protocol, command)
if err != nil {
return "", err
}
switch llmHoneypot.Model {
case LLAMA3:
return llmHoneypot.ollamaCaller(prompt)
case GPT4O:
return llmHoneypot.openAICaller(prompt)
default:
return "", errors.New("no model selected")
}
} }

287
plugins/llm-integration_test.go Normal file
View File

@ -0,0 +1,287 @@
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 []Message
command := "pwd"
//When
prompt, err := buildPrompt(histories, tracer.SSH, command)
//Then
assert.Nil(t, err)
assert.Equal(t, SystemPromptLen, len(prompt))
}
func TestBuildPromptWithHistory(t *testing.T) {
//Given
var histories = []Message{
{
Role: "cat hello.txt",
Content: "world",
},
}
command := "pwd"
//When
prompt, err := buildPrompt(histories, tracer.SSH, command)
//Then
assert.Nil(t, err)
assert.Equal(t, SystemPromptLen+1, len(prompt))
}
func TestBuildExecuteModelFailValidation(t *testing.T) {
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
OpenAIKey: "",
Protocol: tracer.SSH,
Model: GPT4O,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
_, err := openAIGPTVirtualTerminal.ExecuteModel("test")
assert.Equal(t, "openAIKey is empty", err.Error())
}
func TestBuildExecuteModelFailValidationStrategyType(t *testing.T) {
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
OpenAIKey: "",
Protocol: tracer.TCP,
Model: GPT4O,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
_, err := openAIGPTVirtualTerminal.ExecuteModel("test")
assert.Equal(t, "no prompt for protocol selected", err.Error())
}
func TestBuildExecuteModelFailValidationModelType(t *testing.T) {
// Given
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
Protocol: tracer.SSH,
Model: 5,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
//When
_, err := openAIGPTVirtualTerminal.ExecuteModel("ls")
//Then
assert.Errorf(t, err, "no model selected")
}
func TestBuildExecuteModelSSHWithResultsOpenAI(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, &Response{
Choices: []Choice{
{
Message: Message{
Role: SYSTEM.String(),
Content: "prova.txt",
},
},
},
})
if err != nil {
return httpmock.NewStringResponse(500, ""), nil
}
return resp, nil
},
)
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
OpenAIKey: "sdjdnklfjndslkjanfk",
Protocol: tracer.SSH,
Model: GPT4O,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
openAIGPTVirtualTerminal.client = client
//When
str, err := openAIGPTVirtualTerminal.ExecuteModel("ls")
//Then
assert.Nil(t, err)
assert.Equal(t, "prova.txt", str)
}
func TestBuildExecuteModelSSHWithResultsLLama(t *testing.T) {
client := resty.New()
httpmock.ActivateNonDefault(client.GetClient())
defer httpmock.DeactivateAndReset()
// Given
httpmock.RegisterResponder("POST", ollamaEndpoint,
func(req *http.Request) (*http.Response, error) {
resp, err := httpmock.NewJsonResponse(200, &Response{
Message: Message{
Role: SYSTEM.String(),
Content: "prova.txt",
},
})
if err != nil {
return httpmock.NewStringResponse(500, ""), nil
}
return resp, nil
},
)
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
Protocol: tracer.SSH,
Model: LLAMA3,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
openAIGPTVirtualTerminal.client = client
//When
str, err := openAIGPTVirtualTerminal.ExecuteModel("ls")
//Then
assert.Nil(t, err)
assert.Equal(t, "prova.txt", str)
}
func TestBuildExecuteModelSSHWithoutResults(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, &Response{
Choices: []Choice{},
})
if err != nil {
return httpmock.NewStringResponse(500, ""), nil
}
return resp, nil
},
)
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
OpenAIKey: "sdjdnklfjndslkjanfk",
Protocol: tracer.SSH,
Model: GPT4O,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
openAIGPTVirtualTerminal.client = client
//When
_, err := openAIGPTVirtualTerminal.ExecuteModel("ls")
//Then
assert.Equal(t, "no choices", err.Error())
}
func TestBuildExecuteModelHTTPWithResults(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, &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,
}
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 TestBuildExecuteModelHTTPWithoutResults(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, &Response{
Choices: []Choice{},
})
if err != nil {
return httpmock.NewStringResponse(500, ""), nil
}
return resp, nil
},
)
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
OpenAIKey: "sdjdnklfjndslkjanfk",
Protocol: tracer.HTTP,
Model: GPT4O,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
openAIGPTVirtualTerminal.client = client
//When
_, err := openAIGPTVirtualTerminal.ExecuteModel("GET /.aws/credentials")
//Then
assert.Equal(t, "no choices", err.Error())
}

188
plugins/openai-gpt_test.go
View File

@ -1,188 +0,0 @@
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 []Message
command := "pwd"
//When
prompt, err := buildPrompt(histories, tracer.SSH, command)
//Then
assert.Nil(t, err)
assert.Equal(t, SystemPromptLen, len(prompt))
}
func TestBuildPromptWithHistory(t *testing.T) {
//Given
var histories = []Message{
{
Role: "cat hello.txt",
Content: "world",
},
}
command := "pwd"
//When
prompt, err := buildPrompt(histories, tracer.SSH, command)
//Then
assert.Nil(t, err)
assert.Equal(t, SystemPromptLen+1, len(prompt))
}
func TestBuildGetCompletionsFailValidation(t *testing.T) {
openAIGPTVirtualTerminal := Init(make([]Message, 0), "", tracer.SSH)
_, err := openAIGPTVirtualTerminal.GetCompletions("test")
assert.Equal(t, "openAIKey is empty", err.Error())
}
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()
// 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: "prova.txt",
},
},
},
})
if err != nil {
return httpmock.NewStringResponse(500, ""), nil
}
return resp, nil
},
)
openAIGPTVirtualTerminal := Init(make([]Message, 0), "sdjdnklfjndslkjanfk", tracer.SSH)
openAIGPTVirtualTerminal.client = client
//When
str, err := openAIGPTVirtualTerminal.GetCompletions("ls")
//Then
assert.Nil(t, err)
assert.Equal(t, "prova.txt", str)
}
func TestBuildGetCompletionsSSHWithoutResults(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.SSH)
openAIGPTVirtualTerminal.client = client
//When
_, err := openAIGPTVirtualTerminal.GetCompletions("ls")
//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())
}

24
protocols/strategies/http.go
View File

@ -34,13 +34,29 @@ func (httpStrategy HTTPStrategy) Init(beelzebubServiceConfiguration parser.Beelz
if matched { if matched {
responseHTTPBody := command.Handler responseHTTPBody := command.Handler
if command.Plugin == plugins.ChatGPTPluginName { if command.Plugin == plugins.LLMPluginName {
openAIGPTVirtualTerminal := plugins.Init(make([]plugins.Message, 0), beelzebubServiceConfiguration.Plugin.OpenAISecretKey, tracer.HTTP)
llmModel, err := parser.FromString(beelzebubServiceConfiguration.Plugin.LLMModel)
if err != nil {
log.Errorf("Error fromString: %s", err.Error())
responseHTTPBody = "404 Not Found!"
}
llmHoneypot := plugins.LLMHoneypot{
Histories: make([]plugins.Message, 0),
OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
Protocol: tracer.HTTP,
Host: beelzebubServiceConfiguration.Plugin.Host,
Model: llmModel,
}
llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot)
command := fmt.Sprintf("%s %s", request.Method, request.RequestURI) command := fmt.Sprintf("%s %s", request.Method, request.RequestURI)
if completions, err := openAIGPTVirtualTerminal.GetCompletions(command); err != nil { if completions, err := llmHoneypotInstance.ExecuteModel(command); err != nil {
log.Errorf("Error GetCompletions: %s, %s", command, err.Error()) log.Errorf("Error ExecuteModel: %s, %s", command, err.Error())
responseHTTPBody = "404 Not Found!" responseHTTPBody = "404 Not Found!"
} else { } else {
responseHTTPBody = completions responseHTTPBody = completions

24
protocols/strategies/ssh.go
View File

@ -62,11 +62,27 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
if matched { if matched {
commandOutput := command.Handler commandOutput := command.Handler
if command.Plugin == plugins.ChatGPTPluginName { if command.Plugin == plugins.LLMPluginName {
openAIGPTVirtualTerminal := plugins.Init(histories, beelzebubServiceConfiguration.Plugin.OpenAISecretKey, tracer.SSH)
if commandOutput, err = openAIGPTVirtualTerminal.GetCompletions(commandInput); err != nil { llmModel, err := parser.FromString(beelzebubServiceConfiguration.Plugin.LLMModel)
log.Errorf("Error GetCompletions: %s, %s", commandInput, err.Error())
if err != nil {
log.Errorf("Error fromString: %s", err.Error())
commandOutput = "command not found"
}
llmHoneypot := plugins.LLMHoneypot{
Histories: histories,
OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
Protocol: tracer.SSH,
Host: beelzebubServiceConfiguration.Plugin.Host,
Model: llmModel,
}
llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot)
if commandOutput, err = llmHoneypotInstance.ExecuteModel(commandInput); err != nil {
log.Errorf("Error ExecuteModel: %s, %s", commandInput, err.Error())
commandOutput = "command not found" commandOutput = "command not found"
} }
} }