mirror of
https://github.com/mariocandela/beelzebub.git
synced 2025-07-01 18:47:26 -04:00
Feat: add LMM Honeypot HTTP Server (#110)
* add LMM Honeypot HTTP Server * improve unit test code coverage * integrate LLM plugin into http honeypot strategy * improve code coverage * fix typos * improve OpenAI plugin with gpt-4, adpt new API amd map new object
This commit is contained in:
@ -12,9 +12,9 @@ Beelzebub is an advanced honeypot framework designed to provide a highly secure
|
|||||||
|
|
||||||
<img src="https://beelzebub.netlify.app/go-beelzebub.png" alt="Beelzebub Logo" width="200"/>
|
<img src="https://beelzebub.netlify.app/go-beelzebub.png" alt="Beelzebub Logo" width="200"/>
|
||||||
|
|
||||||
## OpenAI GPT Integration
|
## LLM Honeypot
|
||||||
|
|
||||||
Learn how to integrate Beelzebub with OpenAI GPT-3 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)
|
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)
|
||||||
|
|
||||||
## Telegram Bot for Real-Time Attacks
|
## Telegram Bot for Real-Time Attacks
|
||||||
|
|
||||||
|
@ -4,10 +4,10 @@ address: ":2222"
|
|||||||
description: "SSH interactive ChatGPT"
|
description: "SSH interactive ChatGPT"
|
||||||
commands:
|
commands:
|
||||||
- regex: "^(.+)$"
|
- regex: "^(.+)$"
|
||||||
plugin: "OpenAIGPTLinuxTerminal"
|
plugin: "LLMHoneypot"
|
||||||
serverVersion: "OpenSSH"
|
serverVersion: "OpenSSH"
|
||||||
serverName: "ubuntu"
|
serverName: "ubuntu"
|
||||||
passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
|
passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$"
|
||||||
deadlineTimeoutSeconds: 60
|
deadlineTimeoutSeconds: 60
|
||||||
plugin:
|
plugin:
|
||||||
openAPIChatGPTSecretKey: ""
|
openAISecretKey: ""
|
@ -49,7 +49,7 @@ type Prometheus struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
type Plugin struct {
|
type Plugin struct {
|
||||||
OpenAPIChatGPTSecretKey string `yaml:"openAPIChatGPTSecretKey"`
|
OpenAISecretKey string `yaml:"openAISecretKey"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// BeelzebubServiceConfiguration is the struct that contains the configurations of the honeypot service
|
// BeelzebubServiceConfiguration is the struct that contains the configurations of the honeypot service
|
||||||
|
@ -3,31 +3,28 @@ package plugins
|
|||||||
import (
|
import (
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
|
||||||
"github.com/go-resty/resty/v2"
|
"github.com/go-resty/resty/v2"
|
||||||
"strings"
|
"github.com/mariocandela/beelzebub/v3/tracer"
|
||||||
|
|
||||||
log "github.com/sirupsen/logrus"
|
log "github.com/sirupsen/logrus"
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
promptVirtualizeLinuxTerminal = "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. Remember previous commands and consider their effects on subsequent outputs.\n\nA:pwd\n\nQ:/home/user\n\n"
|
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. Remember previous commands and consider their effects on subsequent outputs."
|
||||||
ChatGPTPluginName = "OpenAIGPTLinuxTerminal"
|
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."
|
||||||
openAIGPTEndpoint = "https://api.openai.com/v1/completions"
|
ChatGPTPluginName = "LLMHoneypot"
|
||||||
|
openAIGPTEndpoint = "https://api.openai.com/v1/chat/completions"
|
||||||
)
|
)
|
||||||
|
|
||||||
type History struct {
|
type openAIVirtualHoneypot struct {
|
||||||
Input, Output string
|
Histories []Message
|
||||||
}
|
|
||||||
|
|
||||||
type openAIGPTVirtualTerminal struct {
|
|
||||||
Histories []History
|
|
||||||
openAIKey string
|
openAIKey string
|
||||||
client *resty.Client
|
client *resty.Client
|
||||||
|
protocol tracer.Protocol
|
||||||
}
|
}
|
||||||
|
|
||||||
type Choice struct {
|
type Choice struct {
|
||||||
Text string `json:"text"`
|
Message Message `json:"message"`
|
||||||
Index int `json:"index"`
|
Index int `json:"index"`
|
||||||
Logprobs interface{} `json:"logprobs"`
|
Logprobs interface{} `json:"logprobs"`
|
||||||
FinishReason string `json:"finish_reason"`
|
FinishReason string `json:"finish_reason"`
|
||||||
@ -48,60 +45,105 @@ type gptResponse struct {
|
|||||||
|
|
||||||
type gptRequest struct {
|
type gptRequest struct {
|
||||||
Model string `json:"model"`
|
Model string `json:"model"`
|
||||||
Prompt string `json:"prompt"`
|
Messages []Message `json:"messages"`
|
||||||
Temperature int `json:"temperature"`
|
|
||||||
MaxTokens int `json:"max_tokens"`
|
|
||||||
TopP int `json:"top_p"`
|
|
||||||
FrequencyPenalty int `json:"frequency_penalty"`
|
|
||||||
PresencePenalty int `json:"presence_penalty"`
|
|
||||||
Stop []string `json:"stop"`
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func Init(history []History, openAIKey string) *openAIGPTVirtualTerminal {
|
type Message struct {
|
||||||
return &openAIGPTVirtualTerminal{
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type Role int
|
||||||
|
|
||||||
|
const (
|
||||||
|
SYSTEM Role = iota
|
||||||
|
USER
|
||||||
|
ASSISTANT
|
||||||
|
)
|
||||||
|
|
||||||
|
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,
|
Histories: history,
|
||||||
openAIKey: openAIKey,
|
openAIKey: openAIKey,
|
||||||
client: resty.New(),
|
client: resty.New(),
|
||||||
|
protocol: protocol,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func buildPrompt(histories []History, command string) string {
|
func buildPrompt(histories []Message, protocol tracer.Protocol, command string) ([]Message, error) {
|
||||||
var sb strings.Builder
|
var messages []Message
|
||||||
|
|
||||||
sb.WriteString(promptVirtualizeLinuxTerminal)
|
|
||||||
|
|
||||||
|
switch protocol {
|
||||||
|
case tracer.SSH:
|
||||||
|
messages = append(messages, Message{
|
||||||
|
Role: SYSTEM.String(),
|
||||||
|
Content: systemPromptVirtualizeLinuxTerminal,
|
||||||
|
})
|
||||||
|
messages = append(messages, Message{
|
||||||
|
Role: USER.String(),
|
||||||
|
Content: "pwd",
|
||||||
|
})
|
||||||
|
messages = append(messages, Message{
|
||||||
|
Role: ASSISTANT.String(),
|
||||||
|
Content: "/home/user",
|
||||||
|
})
|
||||||
for _, history := range histories {
|
for _, history := range histories {
|
||||||
sb.WriteString(fmt.Sprintf("A:%s\n\nQ:%s\n\n", history.Input, history.Output))
|
messages = append(messages, history)
|
||||||
}
|
}
|
||||||
// Append command to evaluate
|
case tracer.HTTP:
|
||||||
sb.WriteString(fmt.Sprintf("A:%s\n\nQ:", command))
|
messages = append(messages, Message{
|
||||||
|
Role: SYSTEM.String(),
|
||||||
|
Content: systemPromptVirtualizeHTTPServer,
|
||||||
|
})
|
||||||
|
messages = append(messages, Message{
|
||||||
|
Role: USER.String(),
|
||||||
|
Content: "GET /index.html",
|
||||||
|
})
|
||||||
|
messages = append(messages, Message{
|
||||||
|
Role: ASSISTANT.String(),
|
||||||
|
Content: "<html><body>Hello, World!</body></html>",
|
||||||
|
})
|
||||||
|
default:
|
||||||
|
return nil, errors.New("no prompt for protocol selected")
|
||||||
|
}
|
||||||
|
messages = append(messages, Message{
|
||||||
|
Role: USER.String(),
|
||||||
|
Content: command,
|
||||||
|
})
|
||||||
|
|
||||||
return sb.String()
|
return messages, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (openAIVirtualHoneypot *openAIVirtualHoneypot) GetCompletions(command string) (string, error) {
|
||||||
|
var err error
|
||||||
|
|
||||||
|
prompt, err := buildPrompt(openAIVirtualHoneypot.Histories, openAIVirtualHoneypot.protocol, command)
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
}
|
}
|
||||||
|
|
||||||
func (openAIGPTVirtualTerminal *openAIGPTVirtualTerminal) GetCompletions(command string) (string, error) {
|
|
||||||
requestJson, err := json.Marshal(gptRequest{
|
requestJson, err := json.Marshal(gptRequest{
|
||||||
Model: "gpt-3.5-turbo-instruct",
|
Model: "gpt-4", //"gpt-3.5-turbo",
|
||||||
Prompt: buildPrompt(openAIGPTVirtualTerminal.Histories, command),
|
Messages: prompt,
|
||||||
Temperature: 0,
|
|
||||||
MaxTokens: 100,
|
|
||||||
TopP: 1,
|
|
||||||
FrequencyPenalty: 0,
|
|
||||||
PresencePenalty: 0,
|
|
||||||
Stop: []string{"\n"},
|
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
|
|
||||||
if openAIGPTVirtualTerminal.openAIKey == "" {
|
if openAIVirtualHoneypot.openAIKey == "" {
|
||||||
return "", errors.New("openAIKey is empty")
|
return "", errors.New("openAIKey is empty")
|
||||||
}
|
}
|
||||||
|
|
||||||
response, err := openAIGPTVirtualTerminal.client.R().
|
log.Debug(string(requestJson))
|
||||||
|
response, err := openAIVirtualHoneypot.client.R().
|
||||||
SetHeader("Content-Type", "application/json").
|
SetHeader("Content-Type", "application/json").
|
||||||
SetBody(requestJson).
|
SetBody(requestJson).
|
||||||
SetAuthToken(openAIGPTVirtualTerminal.openAIKey).
|
SetAuthToken(openAIVirtualHoneypot.openAIKey).
|
||||||
SetResult(&gptResponse{}).
|
SetResult(&gptResponse{}).
|
||||||
Post(openAIGPTEndpoint)
|
Post(openAIGPTEndpoint)
|
||||||
|
|
||||||
@ -113,5 +155,5 @@ func (openAIGPTVirtualTerminal *openAIGPTVirtualTerminal) GetCompletions(command
|
|||||||
return "", errors.New("no choices")
|
return "", errors.New("no choices")
|
||||||
}
|
}
|
||||||
|
|
||||||
return response.Result().(*gptResponse).Choices[0].Text, nil
|
return response.Result().(*gptResponse).Choices[0].Message.Content, nil
|
||||||
}
|
}
|
||||||
|
@ -3,58 +3,63 @@ package plugins
|
|||||||
import (
|
import (
|
||||||
"github.com/go-resty/resty/v2"
|
"github.com/go-resty/resty/v2"
|
||||||
"github.com/jarcoal/httpmock"
|
"github.com/jarcoal/httpmock"
|
||||||
|
"github.com/mariocandela/beelzebub/v3/tracer"
|
||||||
"github.com/stretchr/testify/assert"
|
"github.com/stretchr/testify/assert"
|
||||||
"net/http"
|
"net/http"
|
||||||
"testing"
|
"testing"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
const SystemPromptLen = 4
|
||||||
|
|
||||||
func TestBuildPromptEmptyHistory(t *testing.T) {
|
func TestBuildPromptEmptyHistory(t *testing.T) {
|
||||||
//Given
|
//Given
|
||||||
var histories []History
|
var histories []Message
|
||||||
command := "pwd"
|
command := "pwd"
|
||||||
|
|
||||||
//When
|
//When
|
||||||
prompt := buildPrompt(histories, command)
|
prompt, err := buildPrompt(histories, tracer.SSH, command)
|
||||||
|
|
||||||
//Then
|
//Then
|
||||||
assert.Equal(t,
|
assert.Nil(t, err)
|
||||||
"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. Remember previous commands and consider their effects on subsequent outputs.\n\nA:pwd\n\nQ:/home/user\n\nA:pwd\n\nQ:",
|
assert.Equal(t, SystemPromptLen, len(prompt))
|
||||||
prompt)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestBuildPromptWithHistory(t *testing.T) {
|
func TestBuildPromptWithHistory(t *testing.T) {
|
||||||
//Given
|
//Given
|
||||||
var histories = []History{
|
var histories = []Message{
|
||||||
{
|
{
|
||||||
Input: "cat hello.txt",
|
Role: "cat hello.txt",
|
||||||
Output: "world",
|
Content: "world",
|
||||||
},
|
|
||||||
{
|
|
||||||
Input: "echo 1234",
|
|
||||||
Output: "1234",
|
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
command := "pwd"
|
command := "pwd"
|
||||||
|
|
||||||
//When
|
//When
|
||||||
prompt := buildPrompt(histories, command)
|
prompt, err := buildPrompt(histories, tracer.SSH, command)
|
||||||
|
|
||||||
//Then
|
//Then
|
||||||
assert.Equal(t,
|
assert.Nil(t, err)
|
||||||
"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. Remember previous commands and consider their effects on subsequent outputs.\n\nA:pwd\n\nQ:/home/user\n\nA:cat hello.txt\n\nQ:world\n\nA:echo 1234\n\nQ:1234\n\nA:pwd\n\nQ:",
|
assert.Equal(t, SystemPromptLen+1, len(prompt))
|
||||||
prompt)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestBuildGetCompletionsFailValidation(t *testing.T) {
|
func TestBuildGetCompletionsFailValidation(t *testing.T) {
|
||||||
openAIGPTVirtualTerminal := Init(make([]History, 0), "")
|
openAIGPTVirtualTerminal := Init(make([]Message, 0), "", tracer.SSH)
|
||||||
|
|
||||||
_, err := openAIGPTVirtualTerminal.GetCompletions("test")
|
_, err := openAIGPTVirtualTerminal.GetCompletions("test")
|
||||||
|
|
||||||
assert.Equal(t, "openAIKey is empty", err.Error())
|
assert.Equal(t, "openAIKey is empty", err.Error())
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestBuildGetCompletionsWithResults(t *testing.T) {
|
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()
|
client := resty.New()
|
||||||
httpmock.ActivateNonDefault(client.GetClient())
|
httpmock.ActivateNonDefault(client.GetClient())
|
||||||
defer httpmock.DeactivateAndReset()
|
defer httpmock.DeactivateAndReset()
|
||||||
@ -65,7 +70,10 @@ func TestBuildGetCompletionsWithResults(t *testing.T) {
|
|||||||
resp, err := httpmock.NewJsonResponse(200, &gptResponse{
|
resp, err := httpmock.NewJsonResponse(200, &gptResponse{
|
||||||
Choices: []Choice{
|
Choices: []Choice{
|
||||||
{
|
{
|
||||||
Text: "prova.txt",
|
Message: Message{
|
||||||
|
Role: SYSTEM.String(),
|
||||||
|
Content: "prova.txt",
|
||||||
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
})
|
})
|
||||||
@ -76,7 +84,7 @@ func TestBuildGetCompletionsWithResults(t *testing.T) {
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
openAIGPTVirtualTerminal := Init(make([]History, 0), "sdjdnklfjndslkjanfk")
|
openAIGPTVirtualTerminal := Init(make([]Message, 0), "sdjdnklfjndslkjanfk", tracer.SSH)
|
||||||
openAIGPTVirtualTerminal.client = client
|
openAIGPTVirtualTerminal.client = client
|
||||||
|
|
||||||
//When
|
//When
|
||||||
@ -87,7 +95,7 @@ func TestBuildGetCompletionsWithResults(t *testing.T) {
|
|||||||
assert.Equal(t, "prova.txt", str)
|
assert.Equal(t, "prova.txt", str)
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestBuildGetCompletionsWithoutResults(t *testing.T) {
|
func TestBuildGetCompletionsSSHWithoutResults(t *testing.T) {
|
||||||
client := resty.New()
|
client := resty.New()
|
||||||
httpmock.ActivateNonDefault(client.GetClient())
|
httpmock.ActivateNonDefault(client.GetClient())
|
||||||
defer httpmock.DeactivateAndReset()
|
defer httpmock.DeactivateAndReset()
|
||||||
@ -105,7 +113,7 @@ func TestBuildGetCompletionsWithoutResults(t *testing.T) {
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
openAIGPTVirtualTerminal := Init(make([]History, 0), "sdjdnklfjndslkjanfk")
|
openAIGPTVirtualTerminal := Init(make([]Message, 0), "sdjdnklfjndslkjanfk", tracer.SSH)
|
||||||
openAIGPTVirtualTerminal.client = client
|
openAIGPTVirtualTerminal.client = client
|
||||||
|
|
||||||
//When
|
//When
|
||||||
@ -114,3 +122,67 @@ func TestBuildGetCompletionsWithoutResults(t *testing.T) {
|
|||||||
//Then
|
//Then
|
||||||
assert.Equal(t, "no choices", err.Error())
|
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())
|
||||||
|
}
|
||||||
|
@ -3,6 +3,7 @@ package strategies
|
|||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
"github.com/mariocandela/beelzebub/v3/parser"
|
"github.com/mariocandela/beelzebub/v3/parser"
|
||||||
|
"github.com/mariocandela/beelzebub/v3/plugins"
|
||||||
"github.com/mariocandela/beelzebub/v3/tracer"
|
"github.com/mariocandela/beelzebub/v3/tracer"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
@ -31,8 +32,24 @@ func (httpStrategy HTTPStrategy) Init(beelzebubServiceConfiguration parser.Beelz
|
|||||||
}
|
}
|
||||||
|
|
||||||
if matched {
|
if matched {
|
||||||
|
responseHTTPBody := command.Handler
|
||||||
|
|
||||||
|
if command.Plugin == plugins.ChatGPTPluginName {
|
||||||
|
openAIGPTVirtualTerminal := plugins.Init(make([]plugins.Message, 0), beelzebubServiceConfiguration.Plugin.OpenAISecretKey, tracer.HTTP)
|
||||||
|
|
||||||
|
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())
|
||||||
|
responseHTTPBody = "404 Not Found!"
|
||||||
|
} else {
|
||||||
|
responseHTTPBody = completions
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
setResponseHeaders(responseWriter, command.Headers, command.StatusCode)
|
setResponseHeaders(responseWriter, command.Headers, command.StatusCode)
|
||||||
fmt.Fprintf(responseWriter, command.Handler)
|
fmt.Fprintf(responseWriter, responseHTTPBody)
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -42,7 +42,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
|
|||||||
})
|
})
|
||||||
|
|
||||||
term := terminal.NewTerminal(sess, buildPrompt(sess.User(), beelzebubServiceConfiguration.ServerName))
|
term := terminal.NewTerminal(sess, buildPrompt(sess.User(), beelzebubServiceConfiguration.ServerName))
|
||||||
var histories []plugins.History
|
var histories []plugins.Message
|
||||||
for {
|
for {
|
||||||
commandInput, err := term.ReadLine()
|
commandInput, err := term.ReadLine()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@ -63,7 +63,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
|
|||||||
commandOutput := command.Handler
|
commandOutput := command.Handler
|
||||||
|
|
||||||
if command.Plugin == plugins.ChatGPTPluginName {
|
if command.Plugin == plugins.ChatGPTPluginName {
|
||||||
openAIGPTVirtualTerminal := plugins.Init(histories, beelzebubServiceConfiguration.Plugin.OpenAPIChatGPTSecretKey)
|
openAIGPTVirtualTerminal := plugins.Init(histories, beelzebubServiceConfiguration.Plugin.OpenAISecretKey, tracer.SSH)
|
||||||
|
|
||||||
if commandOutput, err = openAIGPTVirtualTerminal.GetCompletions(commandInput); err != nil {
|
if commandOutput, err = openAIGPTVirtualTerminal.GetCompletions(commandInput); err != nil {
|
||||||
log.Errorf("Error GetCompletions: %s, %s", commandInput, err.Error())
|
log.Errorf("Error GetCompletions: %s, %s", commandInput, err.Error())
|
||||||
@ -71,7 +71,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
histories = append(histories, plugins.History{Input: commandInput, Output: commandOutput})
|
histories = append(histories, plugins.Message{Role: plugins.USER.String(), Content: commandOutput})
|
||||||
|
|
||||||
term.Write(append([]byte(commandOutput), '\n'))
|
term.Write(append([]byte(commandOutput), '\n'))
|
||||||
|
|
||||||
|
Reference in New Issue
Block a user