From d062435818af94a143f679142e3c6a129e3f0d00 Mon Sep 17 00:00:00 2001 From: Mario Candela Date: Fri, 16 Dec 2022 23:02:16 +0100 Subject: [PATCH] feat: Add ChatBot GPT-3 integration (#16) * Added dependency go-resty * Configured Parser to read plugin configurations * Add example ssh with ChatGPT plugin * Add client ChatBot * Improve logging * Add integration with plugin OpenAIChatGPT * Improve readme with ChatBot Example * Add contributed ChatGPT question * Refactoring * Refactoring and improve unit test --- README.md | 19 +++ configurations/services/ssh-2222.yaml | 13 ++ go.mod | 2 + go.sum | 16 ++- parser/configurationsParser.go | 6 + plugin/OpenAiGPT.go | 113 ++++++++++++++++++ plugin/OpenAiGPT_test.go | 82 +++++++++++++ .../hypertextTransferProtocolStrategy.go | 2 +- protocols/secureShellStrategy.go | 18 ++- 9 files changed, 260 insertions(+), 11 deletions(-) create mode 100644 configurations/services/ssh-2222.yaml create mode 100644 plugin/OpenAiGPT.go create mode 100644 plugin/OpenAiGPT_test.go diff --git a/README.md b/README.md index 8c349ca..249b81d 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,24 @@ commands: ### Example SSH Honeypot +###### Honeypot with ChatBot GPT-3 ssh-2222.yaml + +```yaml +apiVersion: "v1" +protocol: "ssh" +address: ":2222" +description: "SSH interactive ChatGPT" +commands: + - regex: "^(.+)$" + plugin: "OpenAIChatGPT" +serverVersion: "OpenSSH" +serverName: "ubuntu" +passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$" +deadlineTimeoutSeconds: 60 +plugin: + openAPIChatGPTSecretKey: "Here your ChatBot SecretKey " + ``` + ###### ssh-22.yaml ```yaml @@ -120,6 +138,7 @@ deadlineTimeoutSeconds: 60 ## Features +- OpenAPI ChatBot GPT-3 integration - SSH Honeypot - HTTP Honeypot - TCP Honeypot diff --git a/configurations/services/ssh-2222.yaml b/configurations/services/ssh-2222.yaml new file mode 100644 index 0000000..0467708 --- /dev/null +++ b/configurations/services/ssh-2222.yaml @@ -0,0 +1,13 @@ +apiVersion: "v1" +protocol: "ssh" +address: ":2222" +description: "SSH interactive ChatGPT" +commands: + - regex: "^(.+)$" + plugin: "OpenAIChatGPT" +serverVersion: "OpenSSH" +serverName: "ubuntu" +passwordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$" +deadlineTimeoutSeconds: 60 +plugin: + openAPIChatGPTSecretKey: "" \ No newline at end of file diff --git a/go.mod b/go.mod index 5c7b505..797504c 100644 --- a/go.mod +++ b/go.mod @@ -4,7 +4,9 @@ go 1.16 require ( github.com/gliderlabs/ssh v0.3.5 + github.com/go-resty/resty/v2 v2.7.0 github.com/google/uuid v1.3.0 + github.com/jarcoal/httpmock v1.2.0 github.com/rabbitmq/amqp091-go v1.5.0 github.com/sirupsen/logrus v1.9.0 github.com/stretchr/testify v1.8.1 diff --git a/go.sum b/go.sum index 051b215..114d495 100644 --- a/go.sum +++ b/go.sum @@ -5,14 +5,19 @@ 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.5 h1:OcaySEmAQJgyYcArR+gGGTHCyE7nvhEMTlYY+Dp8CpY= github.com/gliderlabs/ssh v0.3.5/go.mod h1:8XB4KraRrX39qHhT6yxPsHedjA08I/uBVwj4xC+/+z4= +github.com/go-resty/resty/v2 v2.7.0 h1:me+K9p3uhSmXtrBZ4k9jcEAfJmuC8IivWHwaLZwPrFY= +github.com/go-resty/resty/v2 v2.7.0/go.mod h1:9PWDzw47qPphMRFfhsyk0NnSgvluHcljSMVIq3w7q0I= github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/jarcoal/httpmock v1.2.0 h1:gSvTxxFR/MEMfsGrvRbdfpRUMBStovlSRLw0Ep1bwwc= +github.com/jarcoal/httpmock v1.2.0/go.mod h1:oCoTsnAz4+UoOUIf5lJOWV2QQIW5UoeUI6aM2YnWAZk= github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= -github.com/kr/pty v1.1.1 h1:VkoXIwSboBpnk99O/KFauAEILuNHv5DVFKZMBN/gUgw= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/maxatome/go-testdeep v1.11.0 h1:Tgh5efyCYyJFGUYiT0qxBSIDeXw0F5zSoatlou685kk= +github.com/maxatome/go-testdeep v1.11.0/go.mod h1:011SgQ6efzZYAen6fDn4BqQ+lUR72ysdyKe7Dyogw70= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/rabbitmq/amqp091-go v1.5.0 h1:VouyHPBu1CrKyJVfteGknGOGCzmOz0zcv/tONLkb7rg= @@ -21,14 +26,12 @@ github.com/sirupsen/logrus v1.9.0 h1:trlNQbNUG3OdDrDil03MCb1H2o9nJ1x4/5LYw7byDE0 github.com/sirupsen/logrus v1.9.0/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= -github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= -github.com/yuin/goldmark v1.3.5 h1:dPmz1Snjq0kmkz159iL7S6WzdahUTHnHB5M56WFVifs= github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k= go.uber.org/goleak v1.1.12 h1:gZAh5/EyT/HQwlpkCy6wTpqfH9H8Lz8zbm3dZh+OyzA= go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ= @@ -36,19 +39,17 @@ golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACk golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20220826181053-bd7e27e6170d h1:3qF+Z8Hkrw9sOhrFHti9TlB1Hkac1x+DNRkv0XQiFjo= golang.org/x/crypto v0.0.0-20220826181053-bd7e27e6170d/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4= -golang.org/x/lint v0.0.0-20190930215403-16217165b5de h1:5hukYrvBGR8/eNkX5mdUezrA6JiaEZDtJb9Ei+1LlBs= golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= -golang.org/x/mod v0.4.2 h1:Gz96sIWK3OalVv/I/qNygP42zyoKp3xptRVCWRFEBvo= golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM= +golang.org/x/net v0.0.0-20211029224645-99673261e6eb/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= golang.org/x/net v0.0.0-20220826154423-83b083e8dc8b h1:ZmngSVLe/wycRns9MKikG9OWIEjGcGAkacif7oYQaUY= golang.org/x/net v0.0.0-20220826154423-83b083e8dc8b/go.mod h1:YDH+HFinaLZZlnHAfSS6ZXJJ9M9t4Dl22yv3iI2vPwk= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ= golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -68,16 +69,13 @@ golang.org/x/term v0.0.0-20220722155259-a9ba230a4035/go.mod h1:jbD1KX2456YbFQfuX 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= -golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= -golang.org/x/tools v0.1.5 h1:ouewzE6p+/VEB31YYnTbEJdi8pFqKp4P4n85vwo3DHA= golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= diff --git a/parser/configurationsParser.go b/parser/configurationsParser.go index 11d375f..d146a63 100644 --- a/parser/configurationsParser.go +++ b/parser/configurationsParser.go @@ -28,6 +28,10 @@ type Tracing struct { RabbitMQURI string `yaml:"rabbitMQURI"` } +type Plugin struct { + OpenAPIChatGPTSecretKey string `yaml:"openAPIChatGPTSecretKey"` +} + type BeelzebubServiceConfiguration struct { ApiVersion string `yaml:"apiVersion"` Protocol string `yaml:"protocol"` @@ -39,6 +43,7 @@ type BeelzebubServiceConfiguration struct { PasswordRegex string `yaml:"passwordRegex"` Description string `yaml:"description"` Banner string `yaml:"banner"` + Plugin Plugin `yaml:"plugin"` } type Command struct { @@ -46,6 +51,7 @@ type Command struct { Handler string `yaml:"handler"` Headers []string `yaml:"headers"` StatusCode int `yaml:"statusCode"` + Plugin string `yaml:"plugin"` } type configurationsParser struct { diff --git a/plugin/OpenAiGPT.go b/plugin/OpenAiGPT.go new file mode 100644 index 0000000..012b030 --- /dev/null +++ b/plugin/OpenAiGPT.go @@ -0,0 +1,113 @@ +package plugin + +import ( + "encoding/json" + "errors" + "fmt" + "github.com/go-resty/resty/v2" + "strings" +) + +const ChatGPTPluginName = "OpenAIChatGPT" +const openAIGPTEndpoint = "https://api.openai.com/v1/completions" + +type History struct { + Input, Output string +} + +type OpenAIGPTVirtualTerminal struct { + Histories []History + OpenAPIChatGPTSecretKey string + client *resty.Client +} + +func (openAIGPTVirtualTerminal *OpenAIGPTVirtualTerminal) InjectDependency() { + if openAIGPTVirtualTerminal.client == nil { + openAIGPTVirtualTerminal.client = resty.New() + } +} + +type Choice struct { + Text string `json:"text"` + Index int `json:"index"` + Logprobs interface{} `json:"logprobs"` + FinishReason string `json:"finish_reason"` +} + +type gptResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int `json:"created"` + Model string `json:"model"` + Choices []Choice `json:"choices"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + +type gptRequest struct { + Model string `json:"model"` + Prompt string `json:"prompt"` + 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"` +} + +//Reference: https://www.engraved.blog/building-a-virtual-machine-inside/ +const chatGPTFirstQuestion = "I want you to act as a Linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. Do no write explanations. Do not type commands unless I instruct you to do so.\n\nA:pwd\n\nQ:/home/user\n\n" + +func buildPrompt(histories []History, command string) string { + var sb strings.Builder + + sb.WriteString(chatGPTFirstQuestion) + + for _, history := range histories { + sb.WriteString(fmt.Sprintf("A:%s\n\nQ:%s\n\n", history.Input, history.Output)) + } + // Append command to evaluate + sb.WriteString(fmt.Sprintf("A:%s\n\nQ:", command)) + + return sb.String() +} + +func (openAIGPTVirtualTerminal *OpenAIGPTVirtualTerminal) GetCompletions(command string) (string, error) { + requestJson, err := json.Marshal(gptRequest{ + Model: "text-davinci-003", + Prompt: buildPrompt(openAIGPTVirtualTerminal.Histories, command), + Temperature: 0, + MaxTokens: 100, + TopP: 1, + FrequencyPenalty: 0, + PresencePenalty: 0, + Stop: []string{"\n"}, + }) + if err != nil { + return "", err + } + + if openAIGPTVirtualTerminal.OpenAPIChatGPTSecretKey == "" { + return "", errors.New("OpenAPIChatGPTSecretKey is empty") + } + + response, err := openAIGPTVirtualTerminal.client.R(). + SetHeader("Content-Type", "application/json"). + SetBody(requestJson). + SetAuthToken(openAIGPTVirtualTerminal.OpenAPIChatGPTSecretKey). + SetResult(&gptResponse{}). + Post(openAIGPTEndpoint) + + if err != nil { + return "", err + } + + if len(response.Result().(*gptResponse).Choices) == 0 { + return "", errors.New("no choices") + } + + return response.Result().(*gptResponse).Choices[0].Text, nil +} diff --git a/plugin/OpenAiGPT_test.go b/plugin/OpenAiGPT_test.go new file mode 100644 index 0000000..311b8a3 --- /dev/null +++ b/plugin/OpenAiGPT_test.go @@ -0,0 +1,82 @@ +package plugin + +import ( + "github.com/go-resty/resty/v2" + "github.com/jarcoal/httpmock" + "github.com/stretchr/testify/assert" + "net/http" + "testing" +) + +func TestBuildPromptEmptyHistory(t *testing.T) { + //Given + var histories []History + command := "pwd" + + //When + prompt := buildPrompt(histories, command) + + //Then + assert.Equal(t, + "I want you to act as a Linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. Do no write explanations. Do not type commands unless I instruct you to do so.\n\nA:pwd\n\nQ:/home/user\n\nA:pwd\n\nQ:", + prompt) +} + +func TestBuildPromptWithHistory(t *testing.T) { + //Given + var histories = []History{ + { + Input: "cat hello.txt", + Output: "world", + }, + { + Input: "echo 1234", + Output: "1234", + }, + } + + command := "pwd" + + //When + prompt := buildPrompt(histories, command) + + //Then + assert.Equal(t, + "I want you to act as a Linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. Do no write explanations. Do not type commands unless I instruct you to do so.\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:", + prompt) +} + +func TestBuildGetCompletions(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{ + { + Text: "prova.txt", + }, + }, + }) + if err != nil { + return httpmock.NewStringResponse(500, ""), nil + } + return resp, nil + }, + ) + + openAIGPTVirtualTerminal := OpenAIGPTVirtualTerminal{ + OpenAPIChatGPTSecretKey: "sdjdnklfjndslkjanfk", + client: client, + } + + //When + str, err := openAIGPTVirtualTerminal.GetCompletions("ls") + + //Then + assert.Nil(t, err) + assert.Equal(t, "prova.txt", str) +} diff --git a/protocols/hypertextTransferProtocolStrategy.go b/protocols/hypertextTransferProtocolStrategy.go index 0d99308..1d79860 100644 --- a/protocols/hypertextTransferProtocolStrategy.go +++ b/protocols/hypertextTransferProtocolStrategy.go @@ -47,7 +47,7 @@ func (httpStrategy HypertextTransferProtocolStrategy) Init(beelzebubServiceConfi log.WithFields(log.Fields{ "port": beelzebubServiceConfiguration.Address, "commands": len(beelzebubServiceConfiguration.Commands), - }).Infof("Init service %s", beelzebubServiceConfiguration.Protocol) + }).Infof("Init service: %s", beelzebubServiceConfiguration.Description) return nil } diff --git a/protocols/secureShellStrategy.go b/protocols/secureShellStrategy.go index d616f93..d5f2ab4 100644 --- a/protocols/secureShellStrategy.go +++ b/protocols/secureShellStrategy.go @@ -2,6 +2,7 @@ package protocols import ( "beelzebub/parser" + "beelzebub/plugin" "beelzebub/tracer" "fmt" "github.com/gliderlabs/ssh" @@ -39,6 +40,7 @@ func (SSHStrategy *SecureShellStrategy) Init(beelzebubServiceConfiguration parse }) term := terminal.NewTerminal(sess, buildPrompt(sess.User(), beelzebubServiceConfiguration.ServerName)) + var histories []plugin.History for { commandInput, err := term.ReadLine() if err != nil { @@ -64,7 +66,21 @@ func (SSHStrategy *SecureShellStrategy) Init(beelzebubServiceConfiguration parse } if matched { - term.Write(append([]byte(command.Handler), '\n')) + commandOutput := command.Handler + + if command.Plugin == plugin.ChatGPTPluginName { + openAIGPTVirtualTerminal := plugin.OpenAIGPTVirtualTerminal{Histories: histories, OpenAPIChatGPTSecretKey: beelzebubServiceConfiguration.Plugin.OpenAPIChatGPTSecretKey} + openAIGPTVirtualTerminal.InjectDependency() + + if commandOutput, err = openAIGPTVirtualTerminal.GetCompletions(commandInput); err != nil { + log.Errorf("Error GetCompletions: %s, %s", commandInput, err.Error()) + commandOutput = "command not found" + } + } + + histories = append(histories, plugin.History{Input: commandInput, Output: commandOutput}) + + term.Write(append([]byte(commandOutput), '\n')) break } }