feat: Improve SSH LLM honeypot, preserve session after attacker logout (#179)

* Migrate from deprecated library "golang.org/x/crypto/ssh/terminal" to "golang.org/x/term"

* Feat: Inject OpenAI secret key from environment variable

* Feat: Add test for OpenAI secret key injection from environment variable

* Fix: Correct llmModel value in http-80.yaml configuration

* Feat: Add OPEN_AI_SECRET_KEY environment variable to docker-compose.yml

* Feat: Implement session management for SSHStrategy with command history

.github/workflows/codeql.yml

@@ -27,7 +27,7 @@ jobs:
       uses: actions/checkout@v3
 
     - name: Initialize CodeQL
-      uses: github/codeql-action/init@v2
+      uses: github/codeql-action/init@v3
       with:
         languages: ${{ matrix.language }}
 
@@ -35,6 +35,6 @@ jobs:
       run: go build ./...
 
     - name: Perform CodeQL Analysis
-      uses: github/codeql-action/analyze@v2
+      uses: github/codeql-action/analyze@v3
       with:
         category: "/language:${{matrix.language}}"

.github/workflows/docker-image.yml

@@ -1,31 +1,32 @@
+---
 name: Docker Hub Image
 
 on:
   push:
     tags:
       - 'v*.*.*'
-
 jobs:
   CD:
     runs-on: ubuntu-latest
     steps:
-      -
+      - name: Checkout
-        name: Checkout
         uses: actions/checkout@v3
-      -
+      - name: Login to Docker Hub
-        name: Login to Docker Hub
         uses: docker/login-action@v2
         with:
           username: ${{ secrets.DOCKER_USER }}
           password: ${{ secrets.DOCKER_ACCESS_TOKEN }}
-      -
+      - name: Set up QEMU
-        name: Set up Docker Buildx
+        uses: docker/setup-qemu-action@v3
+      - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v2
-      -
+      - name: Build and push
-        name: Build and push
         uses: docker/build-push-action@v4
         with:
           context: .
           file: ./Dockerfile
           push: true
-          tags: m4r10/beelzebub:${{  github.ref_name }}
+          tags: |
+            m4r10/beelzebub:${{  github.ref_name }}
+            m4r10/beelzebub:latest
+          platforms: linux/amd64,linux/arm64

README.md

@@ -211,9 +211,9 @@ commands:
 
 #### Example SSH Honeypot
 
-###### Honeypot LLM Honeypots
+###### LLM Honeypots
 
-Example with OpenAI GPT-4:
+Follow a SSH LLM Honeypot using OpenAI as provider LLM:
 
 ```yaml
 apiVersion: "v1"
@@ -228,11 +228,12 @@ serverName: "ubuntu"
 passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
 deadlineTimeoutSeconds: 60
 plugin:
-   llmModel: "gpt4-o"
+   llmProvider: "openai"
+   llmModel: "gpt-4o" #Models https://platform.openai.com/docs/models
    openAISecretKey: "sk-proj-123456"
 ```
 
-Example with Ollama Llama3:
+Examples with local Ollama instance using model codellama:7b:
 
 ```yaml
 apiVersion: "v1"
@@ -247,9 +248,30 @@ serverName: "ubuntu"
 passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
 deadlineTimeoutSeconds: 60
 plugin:
-   llmModel: "llama3"
+   llmProvider: "ollama"
+   llmModel: "codellama:7b" #Models https://ollama.com/search
    host: "http://example.com/api/chat" #default http://localhost:11434/api/chat
 ```
+Example with custom prompt:
+
+```yaml
+apiVersion: "v1"
+protocol: "ssh"
+address: ":2222"
+description: "SSH interactive OpenAI  GPT-4"
+commands:
+  - regex: "^(.+)$"
+    plugin: "LLMHoneypot"
+serverVersion: "OpenSSH"
+serverName: "ubuntu"
+passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
+deadlineTimeoutSeconds: 60
+plugin:
+   llmProvider: "openai"
+   llmModel: "gpt-4o"
+   openAISecretKey: "sk-proj-123456"
+   prompt: "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."
+```
 
 ###### SSH Honeypot on Port 22
 

configurations/services/http-80.yaml

@@ -22,5 +22,5 @@ commands:
     plugin: "LLMHoneypot"
     statusCode: 200
 plugin:
-  llmModel: "gpt4-o"
+  llmModel: "gpt-4o"
   openAISecretKey: "sk-proj-123456"

configurations/services/ssh-2222.yaml

@@ -7,7 +7,9 @@ commands:
     plugin: "LLMHoneypot"
 serverVersion: "OpenSSH"
 serverName: "ubuntu"
-passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
+passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456|1234)$"
 deadlineTimeoutSeconds: 6000
 plugin:
-  llmModel: "llama3"
+  llmProvider: "openai"
+  llmModel: "gpt-4o"
+  openAISecretKey: "sk-proj-12345"

docker-compose.yml

@@ -3,18 +3,18 @@ version: "3.9"
 services:
   beelzebub:
     build: .
-    #network_mode: host # Not work on Mac OS
     container_name: beelzebub
     restart: always
-    ports: # Remove me, if you use configuration network_mode: host
+    ports:
       - "22:22"
       - "2222:2222"
       - "8080:8080"
       - "8081:8081"
       - "80:80"
       - "3306:3306"
-      - "2112:2112" # Prometheus openmetrics
+      - "2112:2112" #Prometheus Open Metrics
     environment:
       RABBITMQ_URI: ${RABBITMQ_URI}
+      OPEN_AI_SECRET_KEY: ${OPEN_AI_SECRET_KEY}
     volumes:
       - "./configurations:/configurations"

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
@@ -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.10.0
-	golang.org/x/crypto v0.32.0
+	golang.org/x/crypto v0.33.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 
@@ -32,7 +32,7 @@ require (
 	github.com/prometheus/common v0.55.0 // indirect
 	github.com/prometheus/procfs v0.15.1 // indirect
 	golang.org/x/net v0.33.0 // indirect
-	golang.org/x/sys v0.29.0 // indirect
+	golang.org/x/sys v0.30.0 // indirect
-	golang.org/x/term v0.28.0 // indirect
+	golang.org/x/term v0.29.0 // indirect
 	google.golang.org/protobuf v1.34.2 // indirect
 )

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.5 h1:hBKqmWrr7uRc3euHVqmh1HTHcKn99Smr7o5spptdhTM=
-github.com/go-resty/resty/v2 v2.16.3/go.mod h1:hkJtXbA2iKHzJheXYvQ8snQES5ZLGKMwQ07xAwp/fiA=
+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=
@@ -59,8 +59,8 @@ golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACk
 golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/crypto v0.0.0-20211215153901-e495a2d5b3d3/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
 golang.org/x/crypto v0.6.0/go.mod h1:OFC/31mSvZgRz0V1QTNCzfAI1aIRzbiufJtkMIlEp58=
-golang.org/x/crypto v0.32.0 h1:euUpcYgM8WcP71gNpTqQCn6rC2t6ULUPiOzfWaXVVfc=
+golang.org/x/crypto v0.33.0 h1:IOBPskki6Lysi0lo9qQvbxiQ+FvsCC/YWOecCHAixus=
-golang.org/x/crypto v0.32.0/go.mod h1:ZnnJkOaASj8g0AjIduWNlq2NRxL0PlBrbKVyZ6V/Ugc=
+golang.org/x/crypto v0.33.0/go.mod h1:bVdXmD7IV/4GdElGPozy6U7lWdRXA4qyRVGJV57uQ5M=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
@@ -80,13 +80,13 @@ golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU=
+golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
-golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.30.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.28.0 h1:/Ts8HFuMR2E6IP/jlo7QVLZHggjKQbhu/7H0LJFr3Gg=
+golang.org/x/term v0.29.0 h1:L6pJp37ocefwRRtYPKSWOWzOtWSxVajvz2ldH/xi3iU=
-golang.org/x/term v0.28.0/go.mod h1:Sw/lC2IAUZ92udQNf3WodGtn4k/XoLyZoh8v/8uiwek=
+golang.org/x/term v0.29.0/go.mod h1:6bl4lRlvVuDgSf3179VpIxBF0o10JUpXWOnI7nErv7s=
 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=

integration_test/integration_test.go

@@ -67,8 +67,11 @@ func (suite *IntegrationTestSuite) TestInvokeHTTPHoneypot() {
 	response, err := resty.New().R().
 		Get(suite.httpHoneypotHost + "/index.php")
 
+	response.Header().Del("Date")
+
 	suite.Require().NoError(err)
 	suite.Equal(http.StatusOK, response.StatusCode())
+	suite.Equal(http.Header{"Content-Length": []string{"15"}, "Content-Type": []string{"text/html"}, "Server": []string{"Apache/2.4.53 (Debian)"}, "X-Powered-By": []string{"PHP/7.4.29"}}, response.Header())
 	suite.Equal("mocked response", string(response.Body()))
 
 	response, err = resty.New().R().

parser/configurations_parser.go

@@ -52,6 +52,8 @@ type Plugin struct {
 	OpenAISecretKey string `yaml:"openAISecretKey"`
 	Host            string `yaml:"host"`
 	LLMModel        string `yaml:"llmModel"`
+	LLMProvider     string `yaml:"llmProvider"`
+	Prompt          string `yaml:"prompt"`
 }
 
 // BeelzebubServiceConfiguration is the struct that contains the configurations of the honeypot service
@@ -67,6 +69,8 @@ type BeelzebubServiceConfiguration struct {
 	Description            string    `yaml:"description"`
 	Banner                 string    `yaml:"banner"`
 	Plugin                 Plugin    `yaml:"plugin"`
+	TLSCertPath            string    `yaml:"tlsCertPath"`
+	TLSKeyPath             string    `yaml:"tlsKeyPath"`
 }
 
 // Command is the struct that contains the configurations of the commands

parser/configurations_parser_test.go

@@ -49,6 +49,8 @@ func mockReadfilebytesBeelzebubServiceConfiguration(filePath string) ([]byte, er
 apiVersion: "v1"
 protocol: "http"
 address: ":8080"
+tlsCertPath: "/tmp/cert.crt"
+tlsKeyPath: "/tmp/cert.key"
 commands:
   - regex: "wp-admin"
     handler: "login"
@@ -57,7 +59,9 @@ commands:
 plugin:
   openAISecretKey: "qwerty"
   llmModel: "llama3"
+  llmProvider: "ollama"
   host: "localhost:1563"
+  prompt: "hello world"
 `)
 	return beelzebubServiceConfiguration, nil
 }
@@ -132,7 +136,11 @@ func TestReadConfigurationsServicesValid(t *testing.T) {
 	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.LLMProvider, "ollama")
 	assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.Host, "localhost:1563")
+	assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.Prompt, "hello world")
+	assert.Equal(t, firstBeelzebubServiceConfiguration.TLSCertPath, "/tmp/cert.crt")
+	assert.Equal(t, firstBeelzebubServiceConfiguration.TLSKeyPath, "/tmp/cert.key")
 }
 
 func TestGelAllFilesNameByDirName(t *testing.T) {

plugins/beelzebub-cloud_test.go

@@ -85,7 +85,7 @@ func TestGetHoneypotsConfigurationsWithResults(t *testing.T) {
 			resp, err := httpmock.NewJsonResponse(200, &[]HoneypotConfigResponseDTO{
 				{
 					ID:      "123456",
-					Config:  "apiVersion: \"v1\"\nprotocol: \"ssh\"\naddress: \":2222\"\ndescription: \"SSH interactive ChatGPT\"\ncommands:\n  - regex: \"^(.+)$\"\n    plugin: \"LLMHoneypot\"\nserverVersion: \"OpenSSH\"\nserverName: \"ubuntu\"\npasswordRegex: \"^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$\"\ndeadlineTimeoutSeconds: 60\nplugin:\n  llmModel: \"gpt4-o\"\n  openAISecretKey: \"1234\"\n",
+					Config:  "apiVersion: \"v1\"\nprotocol: \"ssh\"\naddress: \":2222\"\ndescription: \"SSH interactive ChatGPT\"\ncommands:\n  - regex: \"^(.+)$\"\n    plugin: \"LLMHoneypot\"\nserverVersion: \"OpenSSH\"\nserverName: \"ubuntu\"\npasswordRegex: \"^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$\"\ndeadlineTimeoutSeconds: 60\nplugin:\n  llmModel: \"gpt-4o\"\n  openAISecretKey: \"1234\"\n",
 					TokenID: "1234567",
 				},
 			})
@@ -120,7 +120,7 @@ func TestGetHoneypotsConfigurationsWithResults(t *testing.T) {
 			PasswordRegex:          "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$",
 			DeadlineTimeoutSeconds: 60,
 			Plugin: parser.Plugin{
-				LLMModel:        "gpt4-o",
+				LLMModel:        "gpt-4o",
 				OpenAISecretKey: "1234",
 			},
 		},

plugins/llm-integration.go

@@ -6,25 +6,29 @@ import (
 	"fmt"
 	"github.com/go-resty/resty/v2"
 	"github.com/mariocandela/beelzebub/v3/tracer"
-
 	log "github.com/sirupsen/logrus"
+	"os"
+	"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 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 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"
 )
 
 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
+	Provider     LLMProvider
-	Host      string
+	Model        string
+	Host         string
+	CustomPrompt string
 }
 
 type Choice struct {
@@ -70,21 +74,21 @@ func (role Role) String() string {
 	return [...]string{"system", "user", "assistant"}[role]
 }
 
-type LLMModel int
+type LLMProvider int
 
 const (
-	LLAMA3 LLMModel = iota
+	Ollama LLMProvider = iota
-	GPT4O
+	OpenAI
 )
 
-func FromStringToLLMModel(llmModel string) (LLMModel, error) {
+func FromStringToLLMProvider(llmProvider string) (LLMProvider, error) {
-	switch llmModel {
+	switch strings.ToLower(llmProvider) {
-	case "llama3":
+	case "ollama":
-		return LLAMA3, nil
+		return Ollama, nil
-	case "gpt4-o":
+	case "openai":
-		return GPT4O, nil
+		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)
 	}
 }
 
@@ -92,17 +96,26 @@ func InitLLMHoneypot(config LLMHoneypot) *LLMHoneypot {
 	// Inject the dependencies
 	config.client = resty.New()
 
+	if os.Getenv("OPEN_AI_SECRET_KEY") != "" {
+		config.OpenAIKey = os.Getenv("OPEN_AI_SECRET_KEY")
+	}
+
 	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(),
@@ -112,13 +125,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(),
@@ -143,7 +160,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,
 	})
@@ -156,7 +173,7 @@ func (llmHoneypot *LLMHoneypot) openAICaller(messages []Message) (string, error)
 	}
 
 	if llmHoneypot.Host == "" {
-		llmHoneypot.Host = openAIGPTEndpoint
+		llmHoneypot.Host = openAIEndpoint
 	}
 
 	log.Debug(string(requestJson))
@@ -175,14 +192,14 @@ func (llmHoneypot *LLMHoneypot) openAICaller(messages []Message) (string, error)
 		return "", errors.New("no choices")
 	}
 
-	return response.Result().(*Response).Choices[0].Message.Content, nil
+	return removeQuotes(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",
+		Model:    llmHoneypot.Model,
 		Messages: messages,
 		Stream:   false,
 	})
@@ -206,24 +223,30 @@ func (llmHoneypot *LLMHoneypot) ollamaCaller(messages []Message) (string, error)
 	}
 	log.Debug(response)
 
-	return response.Result().(*Response).Message.Content, nil
+	return removeQuotes(response.Result().(*Response).Message.Content), nil
 }
 
 func (llmHoneypot *LLMHoneypot) ExecuteModel(command string) (string, error) {
 	var err error
+	var prompt []Message
 
-	prompt, err := buildPrompt(llmHoneypot.Histories, llmHoneypot.Protocol, command)
+	prompt, err = llmHoneypot.buildPrompt(command)
 
 	if err != nil {
 		return "", err
 	}
 
-	switch llmHoneypot.Model {
+	switch llmHoneypot.Provider {
-	case LLAMA3:
+	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)
 	}
 }
+
+func removeQuotes(content string) string {
+	regex := regexp.MustCompile("(```( *)?([a-z]*)?(\\n)?)")
+	return regex.ReplaceAllString(content, "")
+}

plugins/llm-integration_test.go

@@ -6,6 +6,7 @@ import (
 	"github.com/mariocandela/beelzebub/v3/tracer"
 	"github.com/stretchr/testify/assert"
 	"net/http"
+	"os"
 	"testing"
 )
 
@@ -16,8 +17,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,21 +41,53 @@ 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{
 		Histories: make([]Message, 0),
 		OpenAIKey: "",
 		Protocol:  tracer.SSH,
-		Model:     GPT4O,
+		Model:     "gpt-4o",
+		Provider:  OpenAI,
 	}
 
 	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@@ -59,13 +97,78 @@ func TestBuildExecuteModelFailValidation(t *testing.T) {
 	assert.Equal(t, "openAIKey is empty", err.Error())
 }
 
+func TestBuildExecuteModelOpenAISecretKeyFromEnv(t *testing.T) {
+
+	llmHoneypot := LLMHoneypot{
+		Histories: make([]Message, 0),
+		OpenAIKey: "",
+		Protocol:  tracer.SSH,
+		Model:     "gpt-4o",
+		Provider:  OpenAI,
+	}
+
+	os.Setenv("OPEN_AI_SECRET_KEY", "sdjdnklfjndslkjanfk")
+
+	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
+
+	assert.Equal(t, "sdjdnklfjndslkjanfk", openAIGPTVirtualTerminal.OpenAIKey)
+
+}
+
+func TestBuildExecuteModelWithCustomPrompt(t *testing.T) {
+	client := resty.New()
+	httpmock.ActivateNonDefault(client.GetClient())
+	defer httpmock.DeactivateAndReset()
+
+	// Given
+	httpmock.RegisterMatcherResponder("POST", openAIEndpoint,
+		httpmock.BodyContainsString("hello world"),
+		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:        "gpt-4o",
+		Provider:     OpenAI,
+		CustomPrompt: "hello world",
+	}
+
+	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 TestBuildExecuteModelFailValidationStrategyType(t *testing.T) {
 
 	llmHoneypot := LLMHoneypot{
 		Histories: make([]Message, 0),
 		OpenAIKey: "",
 		Protocol:  tracer.TCP,
-		Model:     GPT4O,
+		Model:     "gpt-4o",
+		Provider:  OpenAI,
 	}
 
 	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@@ -80,7 +183,8 @@ func TestBuildExecuteModelFailValidationModelType(t *testing.T) {
 	llmHoneypot := LLMHoneypot{
 		Histories: make([]Message, 0),
 		Protocol:  tracer.SSH,
-		Model:     5,
+		Model:     "llama3",
+		Provider:  5,
 	}
 
 	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@@ -98,7 +202,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{
@@ -121,7 +225,8 @@ func TestBuildExecuteModelSSHWithResultsOpenAI(t *testing.T) {
 		Histories: make([]Message, 0),
 		OpenAIKey: "sdjdnklfjndslkjanfk",
 		Protocol:  tracer.SSH,
-		Model:     GPT4O,
+		Model:     "gpt-4o",
+		Provider:  OpenAI,
 	}
 
 	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@@ -159,7 +264,8 @@ func TestBuildExecuteModelSSHWithResultsLLama(t *testing.T) {
 	llmHoneypot := LLMHoneypot{
 		Histories: make([]Message, 0),
 		Protocol:  tracer.SSH,
-		Model:     LLAMA3,
+		Model:     "llama3",
+		Provider:  Ollama,
 	}
 
 	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@@ -179,7 +285,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{},
@@ -195,7 +301,8 @@ func TestBuildExecuteModelSSHWithoutResults(t *testing.T) {
 		Histories: make([]Message, 0),
 		OpenAIKey: "sdjdnklfjndslkjanfk",
 		Protocol:  tracer.SSH,
-		Model:     GPT4O,
+		Model:     "gpt-4o",
+		Provider:  OpenAI,
 	}
 
 	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@@ -214,7 +321,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{
@@ -237,7 +344,8 @@ func TestBuildExecuteModelHTTPWithResults(t *testing.T) {
 		Histories: make([]Message, 0),
 		OpenAIKey: "sdjdnklfjndslkjanfk",
 		Protocol:  tracer.HTTP,
-		Model:     GPT4O,
+		Model:     "gpt-4o",
+		Provider:  OpenAI,
 	}
 
 	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@@ -257,7 +365,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{},
@@ -273,7 +381,8 @@ func TestBuildExecuteModelHTTPWithoutResults(t *testing.T) {
 		Histories: make([]Message, 0),
 		OpenAIKey: "sdjdnklfjndslkjanfk",
 		Protocol:  tracer.HTTP,
-		Model:     GPT4O,
+		Model:     "gpt-4o",
+		Provider:  OpenAI,
 	}
 
 	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
@@ -287,14 +396,105 @@ 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")
+	model, err = FromStringToLLMProvider("beelzebub-model")
-	assert.Errorf(t, err, "model beelzebub-model not found")
+	assert.Errorf(t, err, "provider beelzebub-model not found")
+}
+
+func TestBuildExecuteModelSSHWithoutPlaintextSection(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: "```plaintext\n```\n",
+				},
+			})
+			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, "", str)
+}
+
+func TestBuildExecuteModelSSHWithoutQuotesSection(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: "```\n```\n",
+				},
+			})
+			if err != nil {
+				return httpmock.NewStringResponse(500, ""), nil
+			}
+			return resp, nil
+		},
+	)
+
+	llmHoneypot := LLMHoneypot{
+		Histories: make([]Message, 0),
+		Protocol:  tracer.SSH,
+		Model:     "llama3",
+		Provider:  Ollama,
+	}
+
+	openAIGPTVirtualTerminal := InitLLMHoneypot(llmHoneypot)
+	openAIGPTVirtualTerminal.client = client
+
+	//When
+	str, err := openAIGPTVirtualTerminal.ExecuteModel("ls")
+
+	//Then
+	assert.Nil(t, err)
+	assert.Equal(t, "", str)
+}
+
+func TestRemoveQuotes(t *testing.T) {
+	plaintext := "```plaintext\n```"
+	bash := "```bash\n```"
+	onlyQuotes := "```\n```"
+	complexText := "```plaintext\ntop - 10:30:48 up 1 day,  4:30,  2 users,  load average: 0.15, 0.10, 0.08\nTasks: 198 total,   1 running, 197 sleeping,   0 stopped,   0 zombie\n```"
+	complexText2 := "```\ntop - 15:06:59 up 10 days,  3:17,  1 user,  load average: 0.10, 0.09, 0.08\nTasks: 285 total\n```"
+
+	assert.Equal(t, "", removeQuotes(plaintext))
+	assert.Equal(t, "", removeQuotes(bash))
+	assert.Equal(t, "", removeQuotes(onlyQuotes))
+	assert.Equal(t, "top - 10:30:48 up 1 day,  4:30,  2 users,  load average: 0.15, 0.10, 0.08\nTasks: 198 total,   1 running, 197 sleeping,   0 stopped,   0 zombie\n", removeQuotes(complexText))
+	assert.Equal(t, "top - 15:06:59 up 10 days,  3:17,  1 user,  load average: 0.10, 0.09, 0.08\nTasks: 285 total\n", removeQuotes(complexText2))
 }

protocols/strategies/http.go

@@ -2,9 +2,6 @@ package strategies
 
 import (
 	"fmt"
-	"github.com/mariocandela/beelzebub/v3/parser"
-	"github.com/mariocandela/beelzebub/v3/plugins"
-	"github.com/mariocandela/beelzebub/v3/tracer"
 	"io"
 	"net"
 	"net/http"
@@ -12,6 +9,9 @@ import (
 	"strings"
 
 	"github.com/google/uuid"
+	"github.com/mariocandela/beelzebub/v3/parser"
+	"github.com/mariocandela/beelzebub/v3/plugins"
+	"github.com/mariocandela/beelzebub/v3/tracer"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -37,19 +37,21 @@ func (httpStrategy HTTPStrategy) Init(beelzebubServiceConfiguration parser.Beelz
 
 				if command.Plugin == plugins.LLMPluginName {
 
-					llmModel, err := plugins.FromStringToLLMModel(beelzebubServiceConfiguration.Plugin.LLMModel)
+					llmProvider, err := plugins.FromStringToLLMProvider(beelzebubServiceConfiguration.Plugin.LLMProvider)
 
 					if err != nil {
-						log.Errorf("Error fromString: %s", err.Error())
+						log.Errorf("Error: %s", err.Error())
 						responseHTTPBody = "404 Not Found!"
 					}
 
 					llmHoneypot := plugins.LLMHoneypot{
-						Histories: make([]plugins.Message, 0),
+						Histories:    make([]plugins.Message, 0),
-						OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
+						OpenAIKey:    beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
-						Protocol:  tracer.HTTP,
+						Protocol:     tracer.HTTP,
-						Host:      beelzebubServiceConfiguration.Plugin.Host,
+						Host:         beelzebubServiceConfiguration.Plugin.Host,
-						Model:     llmModel,
+						Model:        beelzebubServiceConfiguration.Plugin.LLMModel,
+						Provider:     llmProvider,
+						CustomPrompt: beelzebubServiceConfiguration.Plugin.Prompt,
 					}
 
 					llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot)
@@ -66,13 +68,25 @@ func (httpStrategy HTTPStrategy) Init(beelzebubServiceConfiguration parser.Beelz
 				}
 
 				setResponseHeaders(responseWriter, command.Headers, command.StatusCode)
-				fmt.Fprintf(responseWriter, responseHTTPBody)
+				fmt.Fprint(responseWriter, responseHTTPBody)
 				break
 			}
 		}
 	})
 	go func() {
-		err := http.ListenAndServe(httpStrategy.beelzebubServiceConfiguration.Address, serverMux)
+		var err error
+		// Launch a TLS supporting server if we are supplied a TLS Key and Certificate.
+		// If relative paths are supplied, they are relative to the CWD of the binary.
+		// The can be self-signed, only the client will validate this (or not).
+		if httpStrategy.beelzebubServiceConfiguration.TLSKeyPath != "" && httpStrategy.beelzebubServiceConfiguration.TLSCertPath != "" {
+			err = http.ListenAndServeTLS(
+				httpStrategy.beelzebubServiceConfiguration.Address,
+				httpStrategy.beelzebubServiceConfiguration.TLSCertPath,
+				httpStrategy.beelzebubServiceConfiguration.TLSKeyPath,
+				serverMux)
+		} else {
+			err = http.ListenAndServe(httpStrategy.beelzebubServiceConfiguration.Address, serverMux)
+		}
 		if err != nil {
 			log.Errorf("Error during init HTTP Protocol: %s", err.Error())
 			return
@@ -94,7 +108,7 @@ func traceRequest(request *http.Request, tr tracer.Tracer, HoneypotDescription s
 	}
 	host, port, _ := net.SplitHostPort(request.RemoteAddr)
 
-	tr.TraceEvent(tracer.Event{
+	event := tracer.Event{
 		Msg:             "HTTP New request",
 		RequestURI:      request.RequestURI,
 		Protocol:        tracer.HTTP.String(),
@@ -103,26 +117,20 @@ func traceRequest(request *http.Request, tr tracer.Tracer, HoneypotDescription s
 		HostHTTPRequest: request.Host,
 		UserAgent:       request.UserAgent(),
 		Cookies:         mapCookiesToString(request.Cookies()),
-		Headers:         mapHeaderToString(request.Header),
+		Headers:         request.Header,
 		Status:          tracer.Stateless.String(),
 		RemoteAddr:      request.RemoteAddr,
 		SourceIp:        host,
 		SourcePort:      port,
 		ID:              uuid.New().String(),
 		Description:     HoneypotDescription,
-	})
-}
-
-func mapHeaderToString(headers http.Header) string {
-	headersString := ""
-
-	for key := range headers {
-		for _, values := range headers[key] {
-			headersString += fmt.Sprintf("[Key: %s, values: %s],", key, values)
-		}
 	}
-
+	// Capture the TLS details from the request, if provided.
-	return headersString
+	if request.TLS != nil {
+		event.Msg = "HTTPS New Request"
+		event.TLSServerName = request.TLS.ServerName
+	}
+	tr.TraceEvent(event)
 }
 
 func mapCookiesToString(cookies []*http.Cookie) string {

protocols/strategies/ssh.go

@@ -13,13 +13,15 @@ import (
 	"github.com/gliderlabs/ssh"
 	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
-	"golang.org/x/crypto/ssh/terminal"
+	"golang.org/x/term"
 )
 
 type SSHStrategy struct {
+	Sessions map[string][]plugins.Message
 }
 
 func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error {
+	sshStrategy.Sessions = make(map[string][]plugins.Message)
 	go func() {
 		server := &ssh.Server{
 			Addr:        beelzebubServiceConfiguration.Address,
@@ -30,7 +32,9 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 				uuidSession := uuid.New()
 
 				host, port, _ := net.SplitHostPort(sess.RemoteAddr().String())
+				sessionKey := host + sess.User()
 
+				// Inline SSH command
 				if sess.RawCommand() != "" {
 					for _, command := range beelzebubServiceConfiguration.Commands {
 						matched, err := regexp.MatchString(command.Regex, sess.RawCommand())
@@ -44,19 +48,28 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 
 							if command.Plugin == plugins.LLMPluginName {
 
-								llmModel, err := plugins.FromStringToLLMModel(beelzebubServiceConfiguration.Plugin.LLMModel)
+								llmProvider, err := plugins.FromStringToLLMProvider(beelzebubServiceConfiguration.Plugin.LLMProvider)
 
 								if err != nil {
-									log.Errorf("Error fromString: %s", err.Error())
+									log.Errorf("Error: %s", err.Error())
 									commandOutput = "command not found"
+									llmProvider = plugins.OpenAI
+								}
+
+								histories := make([]plugins.Message, 0)
+
+								if sshStrategy.Sessions[sessionKey] != nil {
+									histories = sshStrategy.Sessions[sessionKey]
 								}
 
 								llmHoneypot := plugins.LLMHoneypot{
-									Histories: make([]plugins.Message, 0),
+									Histories:    histories,
-									OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
+									OpenAIKey:    beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
-									Protocol:  tracer.SSH,
+									Protocol:     tracer.SSH,
-									Host:      beelzebubServiceConfiguration.Plugin.Host,
+									Host:         beelzebubServiceConfiguration.Plugin.Host,
-									Model:     llmModel,
+									Model:        beelzebubServiceConfiguration.Plugin.LLMModel,
+									Provider:     llmProvider,
+									CustomPrompt: beelzebubServiceConfiguration.Plugin.Prompt,
 								}
 
 								llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot)
@@ -83,6 +96,10 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 								Command:       sess.RawCommand(),
 								CommandOutput: commandOutput,
 							})
+							var histories []plugins.Message
+							histories = append(histories, plugins.Message{Role: plugins.USER.String(), Content: sess.RawCommand()})
+							histories = append(histories, plugins.Message{Role: plugins.ASSISTANT.String(), Content: commandOutput})
+							sshStrategy.Sessions[sessionKey] = histories
 							tr.TraceEvent(tracer.Event{
 								Msg:    "End SSH Session",
 								Status: tracer.End.String(),
@@ -106,10 +123,14 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 					Description: beelzebubServiceConfiguration.Description,
 				})
 
-				term := terminal.NewTerminal(sess, buildPrompt(sess.User(), beelzebubServiceConfiguration.ServerName))
+				terminal := term.NewTerminal(sess, buildPrompt(sess.User(), beelzebubServiceConfiguration.ServerName))
 				var histories []plugins.Message
+				if sshStrategy.Sessions[sessionKey] != nil {
+					histories = sshStrategy.Sessions[sessionKey]
+				}
+
 				for {
-					commandInput, err := term.ReadLine()
+					commandInput, err := terminal.ReadLine()
 					if err != nil {
 						break
 					}
@@ -129,19 +150,21 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 
 							if command.Plugin == plugins.LLMPluginName {
 
-								llmModel, err := plugins.FromStringToLLMModel(beelzebubServiceConfiguration.Plugin.LLMModel)
+								llmProvider, err := plugins.FromStringToLLMProvider(beelzebubServiceConfiguration.Plugin.LLMProvider)
 
 								if err != nil {
-									log.Errorf("Error fromString: %s", err.Error())
+									log.Errorf("Error: %s, fallback OpenAI", err.Error())
-									commandOutput = "command not found"
+									llmProvider = plugins.OpenAI
 								}
 
 								llmHoneypot := plugins.LLMHoneypot{
-									Histories: histories,
+									Histories:    histories,
-									OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
+									OpenAIKey:    beelzebubServiceConfiguration.Plugin.OpenAISecretKey,
-									Protocol:  tracer.SSH,
+									Protocol:     tracer.SSH,
-									Host:      beelzebubServiceConfiguration.Plugin.Host,
+									Host:         beelzebubServiceConfiguration.Plugin.Host,
-									Model:     llmModel,
+									Model:        beelzebubServiceConfiguration.Plugin.LLMModel,
+									Provider:     llmProvider,
+									CustomPrompt: beelzebubServiceConfiguration.Plugin.Prompt,
 								}
 
 								llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot)
@@ -155,7 +178,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 							histories = append(histories, plugins.Message{Role: plugins.USER.String(), Content: commandInput})
 							histories = append(histories, plugins.Message{Role: plugins.ASSISTANT.String(), Content: commandOutput})
 
-							term.Write(append([]byte(commandOutput), '\n'))
+							terminal.Write(append([]byte(commandOutput), '\n'))
 
 							tr.TraceEvent(tracer.Event{
 								Msg:           "New SSH Terminal Session",
@@ -173,6 +196,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 						}
 					}
 				}
+				sshStrategy.Sessions[sessionKey] = histories
 				tr.TraceEvent(tracer.Event{
 					Msg:    "End SSH Session",
 					Status: tracer.End.String(),

tracer/tracer.go

@@ -2,10 +2,11 @@
 package tracer
 
 import (
-	log "github.com/sirupsen/logrus"
 	"sync"
 	"time"
 
+	log "github.com/sirupsen/logrus"
+
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 )
@@ -26,7 +27,7 @@ type Event struct {
 	User            string
 	Password        string
 	Client          string
-	Headers         string
+	Headers         map[string][]string
 	Cookies         string
 	UserAgent       string
 	HostHTTPRequest string
@@ -36,6 +37,7 @@ type Event struct {
 	Description     string
 	SourceIp        string
 	SourcePort      string
+	TLSServerName   string
 }
 
 type (