fix typo README.md

Feature: add support for llama, refactor yaml service interface (#115 )
* refactor and add llama LMM support * update readme * improve code coverage
2025-07-01 18:47:26 -04:00 · 2024-07-21 20:15:09 +02:00 · 2024-07-21 20:11:18 +02:00 · 2024-07-08 09:01:31 +02:00 · 2024-06-26 23:22:56 +02:00 · 2024-06-23 20:04:43 +02:00
13 changed files with 503 additions and 289 deletions
--- a/README.md
+++ b/README.md
@ -8,13 +8,14 @@

 ## Overview

-Beelzebub is an advanced honeypot framework designed to provide a highly secure environment for detecting and analyzing cyber attacks. It offers a low code approach for easy implementation and utilizes virtualization techniques powered by OpenAI Generative Pre-trained Transformer.
+Beelzebub is an advanced honeypot framework designed to provide a highly secure environment for detecting and analyzing cyber attacks. It offers a low code approach for easy implementation and uses AI to mimic the behavior of a high-interaction honeypot.

 <img src="https://beelzebub.netlify.app/go-beelzebub.png" alt="Beelzebub Logo" width="200"/>

 ## LLM Honeypot

-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)
+[![asciicast](https://asciinema.org/a/665295.svg)](https://asciinema.org/a/665295)
+

 ## Telegram Bot for Real-Time Attacks

@ -103,7 +104,8 @@ $ make test.dependencies.down

 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
 - HTTP Honeypot
 - TCP Honeypot
@ -209,22 +211,44 @@ commands:

 #### Example SSH Honeypot

-###### Honeypot with GPT-3 on Port 2222
+###### Honeypot LLM Honeypots
+
+Example with OpenAI GPT-4:

 ```yaml
 apiVersion: "v1"
 protocol: "ssh"
 address: ":2222"
-description: "SSH interactive ChatGPT"
+description: "SSH interactive OpenAI  GPT-4"
 commands:
  - regex: "^(.+)$"
-    plugin: "OpenAIGPTLinuxTerminal"
+    plugin: "LLMHoneypot"
 serverVersion: "OpenSSH"
 serverName: "ubuntu"
 passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
 deadlineTimeoutSeconds: 60
 plugin:
-  openAISecretKey: "Your OpenAI Secret Key"
+   llmModel: "gpt4-o"
+   openAISecretKey: "sk-proj-123456"
+```
+
+Example with Ollama Llama3:
+
+```yaml
+apiVersion: "v1"
+protocol: "ssh"
+address: ":2222"
+description: "SSH Ollama Llama3"
+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
@ -261,8 +285,6 @@ passwordRegex: "^(root|qwerty|Smoker666)$"
 deadlineTimeoutSeconds: 60
 ```

-[![asciicast](https://asciinema.org/a/604522.svg)](https://asciinema.org/a/604522)
-
 ## Roadmap

 Our future plans for Beelzebub include developing it into a robust PaaS platform.
--- a/configurations/services/http-80.yaml
+++ b/configurations/services/http-80.yaml
@ -18,41 +18,9 @@ commands:
      - "Server: Apache/2.4.53 (Debian)"
      - "X-Powered-By: PHP/7.4.29"
    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: "^.*$"
-    handler:
-      <html>
-        <header>
-          <title>404</title>
-        </header>
-        <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
+    plugin: "LLMHoneypot"
+    statusCode: 200
+plugin:
+  llmModel: "gpt4-o"
+  openAISecretKey: "sk-proj-123456"
--- a/configurations/services/ssh-2222.yaml
+++ b/configurations/services/ssh-2222.yaml
@ -8,6 +8,6 @@ commands:
 serverVersion: "OpenSSH"
 serverName: "ubuntu"
 passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
-deadlineTimeoutSeconds: 60
+deadlineTimeoutSeconds: 6000
 plugin:
-  openAISecretKey: ""
+  llmModel: "llama3"
--- a/go.mod
+++ b/go.mod
@ -12,7 +12,7 @@ require (
 	github.com/rabbitmq/amqp091-go v1.10.0
 	github.com/sirupsen/logrus v1.9.3
 	github.com/stretchr/testify v1.9.0
-	golang.org/x/crypto v0.24.0
+	golang.org/x/crypto v0.25.0
 	gopkg.in/yaml.v3 v3.0.1
 )

@ -30,7 +30,7 @@ require (
 	github.com/prometheus/common v0.48.0 // indirect
 	github.com/prometheus/procfs v0.12.0 // indirect
 	golang.org/x/net v0.25.0 // indirect
-	golang.org/x/sys v0.21.0 // indirect
-	golang.org/x/term v0.21.0 // indirect
+	golang.org/x/sys v0.22.0 // indirect
+	golang.org/x/term v0.22.0 // indirect
 	google.golang.org/protobuf v1.33.0 // indirect
 )
--- a/go.sum
+++ b/go.sum
@ -56,8 +56,8 @@ golang.org/x/crypto v0.0.0-20211215153901-e495a2d5b3d3/go.mod h1:IxCIyHEi3zRg3s0
 golang.org/x/crypto v0.6.0/go.mod h1:OFC/31mSvZgRz0V1QTNCzfAI1aIRzbiufJtkMIlEp58=
 golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
 golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
-golang.org/x/crypto v0.24.0 h1:mnl8DM0o513X8fdIkmyFE/5hTYxbwYOjDS/+rK6qpRI=
-golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM=
+golang.org/x/crypto v0.25.0 h1:ypSNr+bnYL2YhwoMt2zPxHFmbAN1KZs/njMG3hxUp30=
+golang.org/x/crypto v0.25.0/go.mod h1:T+wALwcMOSE0kXgUAnPAHqTLW+XHgcELELW8VaDgm/M=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
@ -84,16 +84,16 @@ golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
-golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI=
+golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
 golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo=
 golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
 golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
-golang.org/x/term v0.21.0 h1:WVXCp+/EBEHOj53Rvu+7KiT/iElMrO8ACK16SMZ3jaA=
-golang.org/x/term v0.21.0/go.mod h1:ooXLefLobQVslOqselCNF4SxFAaoS6KujMbsGzSDmX0=
+golang.org/x/term v0.22.0 h1:BbsgPEJULsl2fV/AT3v15Mjva5yXKQDyKf+TbDz7QJk=
+golang.org/x/term v0.22.0/go.mod h1:F3qCibpT5AMpCRfhfT53vVJwhLtIVHhB9XDjfFvnMI4=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
--- a/parser/configurations_parser.go
+++ b/parser/configurations_parser.go
@ -3,6 +3,7 @@ package parser

 import (
 	"fmt"
+	"github.com/mariocandela/beelzebub/v3/plugins"
 	"os"
 	"path/filepath"
 	"strings"
@ -50,6 +51,19 @@ type Prometheus struct {

 type Plugin struct {
 	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
--- a/parser/configurations_parser_test.go
+++ b/parser/configurations_parser_test.go
@ -2,6 +2,7 @@ package parser

 import (
 	"errors"
+	"github.com/mariocandela/beelzebub/v3/plugins"
 	"os"
 	"testing"

@ -53,7 +54,12 @@ commands:
  - regex: "wp-admin"
    handler: "login"
    headers:
-      - "Content-Type: text/html"`)
+      - "Content-Type: text/html"
+plugin:
+  openAISecretKey: "qwerty"
+  llmModel: "llama3"
+  host: "localhost:1563"
+`)
 	return beelzebubServiceConfiguration, nil
 }

@ -112,10 +118,10 @@ func TestReadConfigurationsServicesValid(t *testing.T) {
 	configurationsParser.gelAllFilesNameByDirNameDependency = mockReadDirValid

 	beelzebubServicesConfiguration, err := configurationsParser.ReadConfigurationsServices()
+	assert.Nil(t, err)

 	firstBeelzebubServiceConfiguration := beelzebubServicesConfiguration[0]

-	assert.Nil(t, err)
 	assert.Equal(t, firstBeelzebubServiceConfiguration.Protocol, "http")
 	assert.Equal(t, firstBeelzebubServiceConfiguration.ApiVersion, "v1")
 	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, len(firstBeelzebubServiceConfiguration.Commands[0].Headers), 1)
 	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) {
@ -177,3 +186,16 @@ func TestReadFileBytesByFilePath(t *testing.T) {

 	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")
+}
--- a/plugins/beelzebub-cloud.go
+++ b/plugins/beelzebub-cloud.go
@ -36,7 +36,6 @@ func (beelzebubCloud *beelzebubCloud) SendEvent(event tracer.Event) (bool, error
 		SetHeader("Content-Type", "application/json").
 		SetBody(requestJson).
 		SetHeader("Authorization", beelzebubCloud.AuthToken).
-		SetResult(&gptResponse{}).
 		Post(beelzebubCloud.URI)

 	log.Debug(response)
--- a/plugins/llm-integration.go
+++ b/plugins/llm-integration.go
@ -12,30 +12,33 @@ import (
 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"
 	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"
+	ollamaEndpoint                      = "http://localhost:11434/api/chat"
 )

-type openAIVirtualHoneypot struct {
+type LLMHoneypot struct {
 	Histories []Message
-	openAIKey string
+	OpenAIKey string
 	client    *resty.Client
-	protocol  tracer.Protocol
+	Protocol  tracer.Protocol
+	Model     LLMModel
+	Host      string
 }

 type Choice struct {
-	Message      Message     `json:"message"`
-	Index        int         `json:"index"`
-	Logprobs     interface{} `json:"logprobs"`
-	FinishReason string      `json:"finish_reason"`
+	Message      Message `json:"message"`
+	Index        int     `json:"index"`
+	FinishReason string  `json:"finish_reason"`
 }

-type gptResponse struct {
+type Response struct {
 	ID      string   `json:"id"`
 	Object  string   `json:"object"`
 	Created int      `json:"created"`
 	Model   string   `json:"model"`
 	Choices []Choice `json:"choices"`
+	Message Message  `json:"message"`
 	Usage   struct {
 		PromptTokens     int `json:"prompt_tokens"`
 		CompletionTokens int `json:"completion_tokens"`
@ -43,9 +46,10 @@ type gptResponse struct {
 	} `json:"usage"`
 }

-type gptRequest struct {
+type Request struct {
 	Model    string    `json:"model"`
 	Messages []Message `json:"messages"`
+	Stream   bool      `json:"stream"`
 }

 type Message struct {
@ -65,13 +69,18 @@ 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,
-	}
+type LLMModel int
+
+const (
+	LLAMA3 LLMModel = iota
+	GPT4O
+)
+
+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) {
@ -118,42 +127,91 @@ func buildPrompt(histories []Message, protocol tracer.Protocol, command string)
 	return messages, nil
 }

-func (openAIVirtualHoneypot *openAIVirtualHoneypot) GetCompletions(command string) (string, error) {
+func (llmHoneypot *LLMHoneypot) openAICaller(messages []Message) (string, error) {
 	var err error

-	prompt, err := buildPrompt(openAIVirtualHoneypot.Histories, openAIVirtualHoneypot.protocol, command)
-
-	if err != nil {
-		return "", err
-	}
-
-	requestJson, err := json.Marshal(gptRequest{
+	requestJson, err := json.Marshal(Request{
 		Model:    "gpt-4o",
-		Messages: prompt,
+		Messages: messages,
+		Stream:   false,
 	})
 	if err != nil {
 		return "", err
 	}

-	if openAIVirtualHoneypot.openAIKey == "" {
+	if llmHoneypot.OpenAIKey == "" {
 		return "", errors.New("openAIKey is empty")
 	}

+	if llmHoneypot.Host == "" {
+		llmHoneypot.Host = openAIGPTEndpoint
+	}
+
 	log.Debug(string(requestJson))
-	response, err := openAIVirtualHoneypot.client.R().
+	response, err := llmHoneypot.client.R().
 		SetHeader("Content-Type", "application/json").
 		SetBody(requestJson).
-		SetAuthToken(openAIVirtualHoneypot.openAIKey).
-		SetResult(&gptResponse{}).
-		Post(openAIGPTEndpoint)
+		SetAuthToken(llmHoneypot.OpenAIKey).
+		SetResult(&Response{}).
+		Post(llmHoneypot.Host)

 	if err != nil {
 		return "", err
 	}
 	log.Debug(response)
-	if len(response.Result().(*gptResponse).Choices) == 0 {
+	if len(response.Result().(*Response).Choices) == 0 {
 		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")
+	}
 }
--- a/plugins/llm-integration_test.go
+++ b/plugins/llm-integration_test.go
@ -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())
+}
--- a/plugins/openai-gpt_test.go
+++ b/plugins/openai-gpt_test.go
@ -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())
-}
--- a/protocols/strategies/http.go
+++ b/protocols/strategies/http.go
@ -34,13 +34,29 @@ 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)
+				if command.Plugin == plugins.LLMPluginName {
+
+					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)

-					if completions, err := openAIGPTVirtualTerminal.GetCompletions(command); err != nil {
-						log.Errorf("Error GetCompletions: %s, %s", command, err.Error())
+					if completions, err := llmHoneypotInstance.ExecuteModel(command); err != nil {
+						log.Errorf("Error ExecuteModel: %s, %s", command, err.Error())
 						responseHTTPBody = "404 Not Found!"
 					} else {
 						responseHTTPBody = completions
--- a/protocols/strategies/ssh.go
+++ b/protocols/strategies/ssh.go
@ -62,11 +62,27 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 						if matched {
 							commandOutput := command.Handler

-							if command.Plugin == plugins.ChatGPTPluginName {
-								openAIGPTVirtualTerminal := plugins.Init(histories, beelzebubServiceConfiguration.Plugin.OpenAISecretKey, tracer.SSH)
+							if command.Plugin == plugins.LLMPluginName {

-								if commandOutput, err = openAIGPTVirtualTerminal.GetCompletions(commandInput); err != nil {
-									log.Errorf("Error GetCompletions: %s, %s", commandInput, err.Error())
+								llmModel, err := parser.FromString(beelzebubServiceConfiguration.Plugin.LLMModel)
+
+								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"
 								}
 							}
Author	SHA1	Message	Date
Mario	cd284877cf	fix typo README.md	2024-07-21 20:15:09 +02:00
Mario Candela	2088163b54	Feature: add support for llama, refactor yaml service interface (#115 ) * refactor and add llama LMM support * update readme * improve code coverage	2024-07-21 20:11:18 +02:00
dependabot[bot]	0af1a05ae9	Bump golang.org/x/crypto from 0.24.0 to 0.25.0 (#113 ) Bumps [golang.org/x/crypto](https://github.com/golang/crypto) from 0.24.0 to 0.25.0. - [Commits](https://github.com/golang/crypto/compare/v0.24.0...v0.25.0) --- updated-dependencies: - dependency-name: golang.org/x/crypto dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>	2024-07-08 09:01:31 +02:00
Mario Candela	5086f5ba08	Update README.md Signed-off-by: Mario Candela <mario.candela.personal@gmail.com>	2024-06-26 23:22:56 +02:00
Mario Candela	e1f82db2d0	Update README.md Signed-off-by: Mario Candela <mario.candela.personal@gmail.com>	2024-06-23 20:04:43 +02:00