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 plugin:LLM honeypot custom prompt (#154)

Browse Source 
refactoring LLM honeypot custom prompt
This commit is contained in:
Mario Candela
2025-01-16 08:46:13 +01:00
committed by GitHub
parent c3d2ff885d
commit 99c7287c02
2 changed files with 53 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
plugins/llm-integration.go
View File

@ -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