mirror of
https://github.com/mariocandela/beelzebub.git
synced 2025-07-01 18:47:26 -04:00
Feat: beelzebub cloud integrations (#117)
* improve beelzebub cloud integration * refactoring cloud integration, fix unit test * add unit test get honeypots * improve code coverage
This commit is contained in:
Mario Candela
@ -3,9 +3,12 @@ package plugins
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"github.com/go-resty/resty/v2"
|
||||
"github.com/mariocandela/beelzebub/v3/parser"
|
||||
"github.com/mariocandela/beelzebub/v3/tracer"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
type beelzebubCloud struct {
|
||||
@ -14,6 +17,13 @@ type beelzebubCloud struct {
|
||||
client *resty.Client
|
||||
}
|
||||
|
||||
type HoneypotConfigResponseDTO struct {
|
||||
ID string `json:"id"`
|
||||
Config string `json:"config"`
|
||||
TokenID string `json:"tokenId"`
|
||||
LastUpdatedOn string `json:"lastUpdatedOn"`
|
||||
}
|
||||
|
||||
func InitBeelzebubCloud(uri, authToken string) *beelzebubCloud {
|
||||
return &beelzebubCloud{
|
||||
URI: uri,
|
||||
@ -36,7 +46,8 @@ func (beelzebubCloud *beelzebubCloud) SendEvent(event tracer.Event) (bool, error
|
||||
SetHeader("Content-Type", "application/json").
|
||||
SetBody(requestJson).
|
||||
SetHeader("Authorization", beelzebubCloud.AuthToken).
|
||||
Post(beelzebubCloud.URI)
|
||||
SetResult(&tracer.Event{}).
|
||||
Post(fmt.Sprintf("%s/events", beelzebubCloud.URI))
|
||||
|
||||
log.Debug(response)
|
||||
|
||||
@ -46,3 +57,44 @@ func (beelzebubCloud *beelzebubCloud) SendEvent(event tracer.Event) (bool, error
|
||||
|
||||
return response.StatusCode() == 200, nil
|
||||
}
|
||||
|
||||
func (beelzebubCloud *beelzebubCloud) GetHoneypotsConfigurations() ([]parser.BeelzebubServiceConfiguration, error) {
|
||||
if beelzebubCloud.AuthToken == "" {
|
||||
return nil, errors.New("authToken is empty")
|
||||
}
|
||||
|
||||
response, err := beelzebubCloud.client.R().
|
||||
SetHeader("Content-Type", "application/json").
|
||||
SetHeader("Authorization", beelzebubCloud.AuthToken).
|
||||
SetResult([]HoneypotConfigResponseDTO{}).
|
||||
Get(fmt.Sprintf("%s/honeypots", beelzebubCloud.URI))
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if response.StatusCode() != 200 {
|
||||
return nil, errors.New(fmt.Sprintf("Response code: %v, error: %s", response.StatusCode(), string(response.Body())))
|
||||
}
|
||||
|
||||
var honeypotsConfig []HoneypotConfigResponseDTO
|
||||
|
||||
if err = json.Unmarshal(response.Body(), &honeypotsConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var servicesConfiguration = make([]parser.BeelzebubServiceConfiguration, 0)
|
||||
|
||||
for _, honeypotConfig := range honeypotsConfig {
|
||||
var honeypotsConfig parser.BeelzebubServiceConfiguration
|
||||
|
||||
if err = yaml.Unmarshal([]byte(honeypotConfig.Config), &honeypotsConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
servicesConfiguration = append(servicesConfiguration, honeypotsConfig)
|
||||
}
|
||||
|
||||
log.Debug(servicesConfiguration)
|
||||
|
||||
return servicesConfiguration, nil
|
||||
}
|
||||
|
||||
@ -1,8 +1,10 @@
|
||||
package plugins
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"github.com/go-resty/resty/v2"
|
||||
"github.com/jarcoal/httpmock"
|
||||
"github.com/mariocandela/beelzebub/v3/parser"
|
||||
"github.com/mariocandela/beelzebub/v3/tracer"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"net/http"
|
||||
@ -22,10 +24,10 @@ func TestBuildSendEventWithResults(t *testing.T) {
|
||||
httpmock.ActivateNonDefault(client.GetClient())
|
||||
defer httpmock.DeactivateAndReset()
|
||||
|
||||
uri := "localhost:8081/events"
|
||||
uri := "localhost:8081"
|
||||
|
||||
// Given
|
||||
httpmock.RegisterResponder("POST", uri,
|
||||
httpmock.RegisterResponder("POST", fmt.Sprintf("%s/events", uri),
|
||||
func(req *http.Request) (*http.Response, error) {
|
||||
resp, err := httpmock.NewJsonResponse(200, &tracer.Event{})
|
||||
if err != nil {
|
||||
@ -69,3 +71,160 @@ func TestBuildSendEventErro(t *testing.T) {
|
||||
//Then
|
||||
assert.Equal(t, false, result)
|
||||
}
|
||||
|
||||
func TestGetHoneypotsConfigurationsWithResults(t *testing.T) {
|
||||
client := resty.New()
|
||||
httpmock.ActivateNonDefault(client.GetClient())
|
||||
defer httpmock.DeactivateAndReset()
|
||||
|
||||
uri := "localhost:8081"
|
||||
|
||||
// Given
|
||||
httpmock.RegisterResponder("GET", fmt.Sprintf("%s/honeypots", uri),
|
||||
func(req *http.Request) (*http.Response, error) {
|
||||
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",
|
||||
TokenID: "1234567",
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return httpmock.NewStringResponse(500, ""), nil
|
||||
}
|
||||
return resp, nil
|
||||
},
|
||||
)
|
||||
|
||||
beelzebubCloud := InitBeelzebubCloud(uri, "sdjdnklfjndslkjanfk")
|
||||
beelzebubCloud.client = client
|
||||
|
||||
//When
|
||||
result, err := beelzebubCloud.GetHoneypotsConfigurations()
|
||||
|
||||
//Then
|
||||
assert.Equal(t, &[]parser.BeelzebubServiceConfiguration{
|
||||
{
|
||||
ApiVersion: "v1",
|
||||
Protocol: "ssh",
|
||||
Address: ":2222",
|
||||
Description: "SSH interactive ChatGPT",
|
||||
Commands: []parser.Command{
|
||||
{
|
||||
Regex: "^(.+)$",
|
||||
Plugin: "LLMHoneypot",
|
||||
},
|
||||
},
|
||||
ServerVersion: "OpenSSH",
|
||||
ServerName: "ubuntu",
|
||||
PasswordRegex: "^(root|qwerty|Smoker666|123456|jenkins|minecraft|sinus|alex|postgres|Ly123456)$",
|
||||
DeadlineTimeoutSeconds: 60,
|
||||
Plugin: parser.Plugin{
|
||||
LLMModel: "gpt4-o",
|
||||
OpenAISecretKey: "1234",
|
||||
},
|
||||
},
|
||||
}, &result)
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
func TestGetHoneypotsConfigurationsWithErrorValidation(t *testing.T) {
|
||||
//Given
|
||||
beelzebubCloud := InitBeelzebubCloud("", "")
|
||||
|
||||
//When
|
||||
result, err := beelzebubCloud.GetHoneypotsConfigurations()
|
||||
|
||||
//Then
|
||||
assert.Nil(t, result)
|
||||
assert.Equal(t, "authToken is empty", err.Error())
|
||||
}
|
||||
|
||||
func TestGetHoneypotsConfigurationsWithErrorAPI(t *testing.T) {
|
||||
client := resty.New()
|
||||
httpmock.ActivateNonDefault(client.GetClient())
|
||||
defer httpmock.DeactivateAndReset()
|
||||
|
||||
uri := "localhost:8081"
|
||||
|
||||
// Given
|
||||
httpmock.RegisterResponder("GET", fmt.Sprintf("%s/honeypots", uri),
|
||||
func(req *http.Request) (*http.Response, error) {
|
||||
return httpmock.NewStringResponse(500, ""), nil
|
||||
},
|
||||
)
|
||||
|
||||
beelzebubCloud := InitBeelzebubCloud(uri, "sdjdnklfjndslkjanfk")
|
||||
beelzebubCloud.client = client
|
||||
|
||||
//When
|
||||
result, err := beelzebubCloud.GetHoneypotsConfigurations()
|
||||
|
||||
//Then
|
||||
assert.Nil(t, result)
|
||||
assert.Equal(t, "Response code: 500, error: ", err.Error())
|
||||
}
|
||||
|
||||
func TestGetHoneypotsConfigurationsWithErrorUnmarshal(t *testing.T) {
|
||||
client := resty.New()
|
||||
httpmock.ActivateNonDefault(client.GetClient())
|
||||
defer httpmock.DeactivateAndReset()
|
||||
|
||||
uri := "localhost:8081"
|
||||
|
||||
// Given
|
||||
httpmock.RegisterResponder("GET", fmt.Sprintf("%s/honeypots", uri),
|
||||
func(req *http.Request) (*http.Response, error) {
|
||||
resp, err := httpmock.NewJsonResponse(200, "error")
|
||||
if err != nil {
|
||||
return httpmock.NewStringResponse(500, ""), nil
|
||||
}
|
||||
return resp, nil
|
||||
},
|
||||
)
|
||||
|
||||
beelzebubCloud := InitBeelzebubCloud(uri, "sdjdnklfjndslkjanfk")
|
||||
beelzebubCloud.client = client
|
||||
|
||||
//When
|
||||
result, err := beelzebubCloud.GetHoneypotsConfigurations()
|
||||
|
||||
//Then
|
||||
assert.Nil(t, result)
|
||||
assert.Equal(t, "json: cannot unmarshal string into Go value of type []plugins.HoneypotConfigResponseDTO", err.Error())
|
||||
}
|
||||
|
||||
func TestGetHoneypotsConfigurationsWithErrorDeserializeYaml(t *testing.T) {
|
||||
client := resty.New()
|
||||
httpmock.ActivateNonDefault(client.GetClient())
|
||||
defer httpmock.DeactivateAndReset()
|
||||
|
||||
uri := "localhost:8081"
|
||||
|
||||
// Given
|
||||
httpmock.RegisterResponder("GET", fmt.Sprintf("%s/honeypots", uri),
|
||||
func(req *http.Request) (*http.Response, error) {
|
||||
resp, err := httpmock.NewJsonResponse(200, &[]HoneypotConfigResponseDTO{
|
||||
{
|
||||
ID: "123456",
|
||||
Config: "error",
|
||||
TokenID: "1234567",
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return httpmock.NewStringResponse(500, ""), nil
|
||||
}
|
||||
return resp, nil
|
||||
},
|
||||
)
|
||||
|
||||
beelzebubCloud := InitBeelzebubCloud(uri, "sdjdnklfjndslkjanfk")
|
||||
beelzebubCloud.client = client
|
||||
|
||||
//When
|
||||
result, err := beelzebubCloud.GetHoneypotsConfigurations()
|
||||
|
||||
//Then
|
||||
assert.Nil(t, result)
|
||||
assert.Equal(t, "yaml: unmarshal errors:\n line 1: cannot unmarshal !!str `error` into parser.BeelzebubServiceConfiguration", err.Error())
|
||||
}
|
||||
|
||||
@ -3,6 +3,7 @@ package plugins
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"github.com/go-resty/resty/v2"
|
||||
"github.com/mariocandela/beelzebub/v3/tracer"
|
||||
|
||||
@ -76,6 +77,17 @@ const (
|
||||
GPT4O
|
||||
)
|
||||
|
||||
func FromStringToLLMModel(llmModel string) (LLMModel, error) {
|
||||
switch llmModel {
|
||||
case "llama3":
|
||||
return LLAMA3, nil
|
||||
case "gpt4-o":
|
||||
return GPT4O, nil
|
||||
default:
|
||||
return -1, fmt.Errorf("model %s not found", llmModel)
|
||||
}
|
||||
}
|
||||
|
||||
func InitLLMHoneypot(config LLMHoneypot) *LLMHoneypot {
|
||||
// Inject the dependencies
|
||||
config.client = resty.New()
|
||||
|
||||
@ -285,3 +285,16 @@ func TestBuildExecuteModelHTTPWithoutResults(t *testing.T) {
|
||||
//Then
|
||||
assert.Equal(t, "no choices", err.Error())
|
||||
}
|
||||
|
||||
func TestFromString(t *testing.T) {
|
||||
model, err := FromStringToLLMModel("llama3")
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, LLAMA3, model)
|
||||
|
||||
model, err = FromStringToLLMModel("gpt4-o")
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, GPT4O, model)
|
||||
|
||||
model, err = FromStringToLLMModel("beelzebub-model")
|
||||
assert.Errorf(t, err, "model beelzebub-model not found")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user