mirror of
https://github.com/mariocandela/beelzebub.git
synced 2025-07-01 18:47:26 -04:00
Compare commits
3 Commits
Author | SHA1 | Date | |
---|---|---|---|
4813685834 | |||
6f6acb212b | |||
99c7287c02 |
2
go.mod
2
go.mod
@ -4,7 +4,7 @@ go 1.20
|
||||
|
||||
require (
|
||||
github.com/gliderlabs/ssh v0.3.8
|
||||
github.com/go-resty/resty/v2 v2.16.3
|
||||
github.com/go-resty/resty/v2 v2.16.5
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/jarcoal/httpmock v1.3.1
|
||||
github.com/melbahja/goph v1.4.0
|
||||
|
4
go.sum
4
go.sum
@ -10,8 +10,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/gliderlabs/ssh v0.3.8 h1:a4YXD1V7xMF9g5nTkdfnja3Sxy1PVDCj1Zg4Wb8vY6c=
|
||||
github.com/gliderlabs/ssh v0.3.8/go.mod h1:xYoytBv1sV0aL3CavoDuJIQNURXkkfPA/wxQ1pL1fAU=
|
||||
github.com/go-resty/resty/v2 v2.16.3 h1:zacNT7lt4b8M/io2Ahj6yPypL7bqx9n1iprfQuodV+E=
|
||||
github.com/go-resty/resty/v2 v2.16.3/go.mod h1:hkJtXbA2iKHzJheXYvQ8snQES5ZLGKMwQ07xAwp/fiA=
|
||||
github.com/go-resty/resty/v2 v2.16.5 h1:hBKqmWrr7uRc3euHVqmh1HTHcKn99Smr7o5spptdhTM=
|
||||
github.com/go-resty/resty/v2 v2.16.5/go.mod h1:hkJtXbA2iKHzJheXYvQ8snQES5ZLGKMwQ07xAwp/fiA=
|
||||
github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
|
@ -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
|
||||
|
@ -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{
|
||||
|
Reference in New Issue
Block a user