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 LLM Plugin, update docs. (#165)

Browse Source 
Refactoring LLM Plugin, update docs.
This commit is contained in:
Mario Candela
2025-02-16 22:48:59 +01:00
committed by GitHub
parent 8703d1afda
commit 38297faed2
8 changed files with 82 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
plugins/llm-integration.go
View File

@ -8,13 +8,14 @@ import (
"github.com/mariocandela/beelzebub/v3/tracer"
log "github.com/sirupsen/logrus"
"regexp"
"strings"
)
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 note. 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."
LLMPluginName = "LLMHoneypot"
openAIGPTEndpoint = "https://api.openai.com/v1/chat/completions"
openAIEndpoint = "https://api.openai.com/v1/chat/completions"
ollamaEndpoint = "http://localhost:11434/api/chat"
)
@ -23,7 +24,8 @@ type LLMHoneypot struct {
OpenAIKey string
client *resty.Client
Protocol tracer.Protocol
Model LLMModel
Provider LLMProvider
Model string
Host string
CustomPrompt string
}
@ -71,21 +73,21 @@ func (role Role) String() string {
return [...]string{"system", "user", "assistant"}[role]
}
type LLMModel int
type LLMProvider int
const (
LLAMA3 LLMModel = iota
GPT4O
Ollama LLMProvider = iota
OpenAI
)
func FromStringToLLMModel(llmModel string) (LLMModel, error) {
switch llmModel {
case "llama3":
return LLAMA3, nil
case "gpt4-o":
return GPT4O, nil
func FromStringToLLMProvider(llmProvider string) (LLMProvider, error) {
switch strings.ToLower(llmProvider) {
case "ollama":
return Ollama, nil
case "openai":
return OpenAI, nil
default:
return -1, fmt.Errorf("model %s not found", llmModel)
return -1, fmt.Errorf("provider %s not found, valid providers: ollama, openai", llmProvider)
}
}
@ -153,7 +155,7 @@ func (llmHoneypot *LLMHoneypot) openAICaller(messages []Message) (string, error)
var err error
requestJson, err := json.Marshal(Request{
Model: "gpt-4o",
Model: llmHoneypot.Model,
Messages: messages,
Stream: false,
})
@ -166,7 +168,7 @@ func (llmHoneypot *LLMHoneypot) openAICaller(messages []Message) (string, error)
}
if llmHoneypot.Host == "" {
llmHoneypot.Host = openAIGPTEndpoint
llmHoneypot.Host = openAIEndpoint
}
log.Debug(string(requestJson))
@ -192,7 +194,7 @@ func (llmHoneypot *LLMHoneypot) ollamaCaller(messages []Message) (string, error)
var err error
requestJson, err := json.Marshal(Request{
Model: "llama3",
Model: llmHoneypot.Model,
Messages: messages,
Stream: false,
})
@ -229,13 +231,13 @@ func (llmHoneypot *LLMHoneypot) ExecuteModel(command string) (string, error) {
return "", err
}
switch llmHoneypot.Model {
case LLAMA3:
switch llmHoneypot.Provider {
case Ollama:
return llmHoneypot.ollamaCaller(prompt)
case GPT4O:
case OpenAI:
return llmHoneypot.openAICaller(prompt)
default:
return "", errors.New("no model selected")
return "", fmt.Errorf("provider %d not found, valid providers: ollama, openai", llmHoneypot.Provider)
}
}

54
plugins/llm-integration_test.go
View File

@ -85,7 +85,8 @@ func TestBuildExecuteModelFailValidation(t *testing.T) {
Histories: make([]Message, 0),
OpenAIKey: "",
Protocol: tracer.SSH,
Model: GPT4O,
Model: "gpt4-o",
Provider: OpenAI,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@ -101,7 +102,7 @@ func TestBuildExecuteModelWithCustomPrompt(t *testing.T) {
defer httpmock.DeactivateAndReset()
// Given
httpmock.RegisterMatcherResponder("POST", openAIGPTEndpoint,
httpmock.RegisterMatcherResponder("POST", openAIEndpoint,
httpmock.BodyContainsString("hello world"),
func(req *http.Request) (*http.Response, error) {
resp, err := httpmock.NewJsonResponse(200, &Response{
@ -125,7 +126,8 @@ func TestBuildExecuteModelWithCustomPrompt(t *testing.T) {
Histories: make([]Message, 0),
OpenAIKey: "sdjdnklfjndslkjanfk",
Protocol: tracer.HTTP,
Model: GPT4O,
Model: "gpt4-o",
Provider: OpenAI,
CustomPrompt: "hello world",
}
@ -146,7 +148,8 @@ func TestBuildExecuteModelFailValidationStrategyType(t *testing.T) {
Histories: make([]Message, 0),
OpenAIKey: "",
Protocol: tracer.TCP,
Model: GPT4O,
Model: "gpt4-o",
Provider: OpenAI,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@ -161,7 +164,8 @@ func TestBuildExecuteModelFailValidationModelType(t *testing.T) {
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
Protocol: tracer.SSH,
Model: 5,
Model: "llama3",
Provider: 5,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@ -179,7 +183,7 @@ func TestBuildExecuteModelSSHWithResultsOpenAI(t *testing.T) {
defer httpmock.DeactivateAndReset()
// Given
httpmock.RegisterResponder("POST", openAIGPTEndpoint,
httpmock.RegisterResponder("POST", openAIEndpoint,
func(req *http.Request) (*http.Response, error) {
resp, err := httpmock.NewJsonResponse(200, &Response{
Choices: []Choice{
@ -202,7 +206,8 @@ func TestBuildExecuteModelSSHWithResultsOpenAI(t *testing.T) {
Histories: make([]Message, 0),
OpenAIKey: "sdjdnklfjndslkjanfk",
Protocol: tracer.SSH,
Model: GPT4O,
Model: "gpt4-o",
Provider: OpenAI,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@ -240,7 +245,8 @@ func TestBuildExecuteModelSSHWithResultsLLama(t *testing.T) {
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
Protocol: tracer.SSH,
Model: LLAMA3,
Model: "llama3",
Provider: Ollama,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@ -260,7 +266,7 @@ func TestBuildExecuteModelSSHWithoutResults(t *testing.T) {
defer httpmock.DeactivateAndReset()
// Given
httpmock.RegisterResponder("POST", openAIGPTEndpoint,
httpmock.RegisterResponder("POST", openAIEndpoint,
func(req *http.Request) (*http.Response, error) {
resp, err := httpmock.NewJsonResponse(200, &Response{
Choices: []Choice{},
@ -276,7 +282,8 @@ func TestBuildExecuteModelSSHWithoutResults(t *testing.T) {
Histories: make([]Message, 0),
OpenAIKey: "sdjdnklfjndslkjanfk",
Protocol: tracer.SSH,
Model: GPT4O,
Model: "gpt4-o",
Provider: OpenAI,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@ -295,7 +302,7 @@ func TestBuildExecuteModelHTTPWithResults(t *testing.T) {
defer httpmock.DeactivateAndReset()
// Given
httpmock.RegisterResponder("POST", openAIGPTEndpoint,
httpmock.RegisterResponder("POST", openAIEndpoint,
func(req *http.Request) (*http.Response, error) {
resp, err := httpmock.NewJsonResponse(200, &Response{
Choices: []Choice{
@ -318,7 +325,8 @@ func TestBuildExecuteModelHTTPWithResults(t *testing.T) {
Histories: make([]Message, 0),
OpenAIKey: "sdjdnklfjndslkjanfk",
Protocol: tracer.HTTP,
Model: GPT4O,
Model: "gpt4-o",
Provider: OpenAI,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@ -338,7 +346,7 @@ func TestBuildExecuteModelHTTPWithoutResults(t *testing.T) {
defer httpmock.DeactivateAndReset()
// Given
httpmock.RegisterResponder("POST", openAIGPTEndpoint,
httpmock.RegisterResponder("POST", openAIEndpoint,
func(req *http.Request) (*http.Response, error) {
resp, err := httpmock.NewJsonResponse(200, &Response{
Choices: []Choice{},
@ -354,7 +362,8 @@ func TestBuildExecuteModelHTTPWithoutResults(t *testing.T) {
Histories: make([]Message, 0),
OpenAIKey: "sdjdnklfjndslkjanfk",
Protocol: tracer.HTTP,
Model: GPT4O,
Model: "gpt4-o",
Provider: OpenAI,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@ -368,16 +377,16 @@ func TestBuildExecuteModelHTTPWithoutResults(t *testing.T) {
}
func TestFromString(t *testing.T) {
model, err := FromStringToLLMModel("llama3")
model, err := FromStringToLLMProvider("openai")
assert.Nil(t, err)
assert.Equal(t, LLAMA3, model)
assert.Equal(t, OpenAI, model)
model, err = FromStringToLLMModel("gpt4-o")
model, err = FromStringToLLMProvider("ollama")
assert.Nil(t, err)
assert.Equal(t, GPT4O, model)
assert.Equal(t, Ollama, model)
model, err = FromStringToLLMModel("beelzebub-model")
assert.Errorf(t, err, "model beelzebub-model not found")
model, err = FromStringToLLMProvider("beelzebub-model")
assert.Errorf(t, err, "provider beelzebub-model not found")
}
func TestBuildExecuteModelSSHWithoutPlaintextSection(t *testing.T) {
@ -404,7 +413,7 @@ func TestBuildExecuteModelSSHWithoutPlaintextSection(t *testing.T) {
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
Protocol: tracer.SSH,
Model: LLAMA3,
Model: "llama3",
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@ -442,7 +451,8 @@ func TestBuildExecuteModelSSHWithoutQuotesSection(t *testing.T) {
llmHoneypot := LLMHoneypot{
Histories: make([]Message, 0),
Protocol: tracer.SSH,
Model: LLAMA3,
Model: "llama3",
Provider: Ollama,
}
openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)