diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8579fb6..04747da 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,9 +2,9 @@ name: CI on: push: - branches: [ main ] + branches: [ "main" ] pull_request: - branches: [ main ] + branches: [ "main" ] jobs: @@ -13,8 +13,7 @@ jobs: fail-fast: false matrix: go-version: - - "1.20.0" - - "stable" + - "1.24.1" runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 @@ -41,7 +40,7 @@ jobs: - name: Quality Gate - Test coverage shall be above threshold env: - TESTCOVERAGE_THRESHOLD: 75 + TESTCOVERAGE_THRESHOLD: 80 run: | echo "Quality Gate: checking test coverage is above threshold ..." echo "Threshold : $TESTCOVERAGE_THRESHOLD %" diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 0dd6e3c..0426cef 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -29,7 +29,7 @@ jobs: - name: Set up Go uses: actions/setup-go@v3 with: - go-version: 1.20.0 + go-version: 1.24.1 - name: Initialize CodeQL uses: github/codeql-action/init@v3 diff --git a/go.mod b/go.mod index e629597..3a63e91 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,8 @@ module github.com/mariocandela/beelzebub/v3 -go 1.20 +go 1.24 + +toolchain go1.24.1 require ( github.com/gliderlabs/ssh v0.3.8 diff --git a/go.sum b/go.sum index 41c23ec..526b3bb 100644 --- a/go.sum +++ b/go.sum @@ -13,6 +13,7 @@ github.com/gliderlabs/ssh v0.3.8/go.mod h1:xYoytBv1sV0aL3CavoDuJIQNURXkkfPA/wxQ1 github.com/go-resty/resty/v2 v2.16.5 h1:hBKqmWrr7uRc3euHVqmh1HTHcKn99Smr7o5spptdhTM= 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/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= 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= github.com/jarcoal/httpmock v1.3.1 h1:iUx3whfZWVf3jT01hQTO/Eo5sAYtB2/rqaUuOtpInww= @@ -22,10 +23,13 @@ github.com/klauspost/compress v1.17.9/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ib github.com/kr/fs v0.1.0 h1:Jskdu9ieNAYnjxsi0LbQp1ulIKZV1LAFgK1tWhpZgl8= github.com/kr/fs v0.1.0/go.mod h1:FFnZGqtBN9Gxj7eW1uZ42v5BccTP0vu6NEaFoC2HwRg= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= github.com/maxatome/go-testdeep v1.12.0 h1:Ql7Go8Tg0C1D/uMMX59LAoYK7LffeJQ6X2T04nTH68g= +github.com/maxatome/go-testdeep v1.12.0/go.mod h1:lPZc/HAcJMP92l7yI6TRz1aZN5URwUBUAfUNvrclaNM= github.com/melbahja/goph v1.4.0 h1:z0PgDbBFe66lRYl3v5dGb9aFgPy0kotuQ37QOwSQFqs= github.com/melbahja/goph v1.4.0/go.mod h1:uG+VfK2Dlhk+O32zFrRlc3kYKTlV6+BtvPWd/kK7U68= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= @@ -47,6 +51,7 @@ github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoG github.com/rabbitmq/amqp091-go v1.10.0 h1:STpn5XsHlHGcecLmMFCtg7mqq0RnD+zFr4uzukfVhBw= github.com/rabbitmq/amqp091-go v1.10.0/go.mod h1:Hy4jKW5kQART1u+JkDTF9YYOQUHXqMuhrgxOEeS7G4o= github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ= +github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog= github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= @@ -55,6 +60,7 @@ github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOf github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= 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= @@ -93,6 +99,7 @@ golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/time v0.6.0 h1:eTDhh4ZXt5Qf0augr54TN6suAUudPcawVZeIAPU7D4U= +golang.org/x/time v0.6.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= @@ -101,6 +108,7 @@ google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6h google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/historystore/history_store.go b/historystore/history_store.go index 0df46eb..869aa44 100644 --- a/historystore/history_store.go +++ b/historystore/history_store.go @@ -2,20 +2,32 @@ package historystore import ( "sync" + "time" "github.com/mariocandela/beelzebub/v3/plugins" ) +var ( + MaxHistoryAge = 60 * time.Minute + CleanerInterval = 1 * time.Minute +) + // HistoryStore is a thread-safe structure for storing Messages used to build LLM Context. type HistoryStore struct { sync.RWMutex - sessions map[string][]plugins.Message + sessions map[string]HistoryEvent +} + +// HistoryEvent is a container for storing messages +type HistoryEvent struct { + LastSeen time.Time + Messages []plugins.Message } // NewHistoryStore returns a prepared HistoryStore func NewHistoryStore() *HistoryStore { return &HistoryStore{ - sessions: make(map[string][]plugins.Message), + sessions: make(map[string]HistoryEvent), } } @@ -31,7 +43,7 @@ func (hs *HistoryStore) HasKey(key string) bool { func (hs *HistoryStore) Query(key string) []plugins.Message { hs.RLock() defer hs.RUnlock() - return hs.sessions[key] + return hs.sessions[key].Messages } // Append will add the slice of Mesages to the entry for the key. @@ -41,7 +53,30 @@ func (hs *HistoryStore) Append(key string, message ...plugins.Message) { defer hs.Unlock() // In the unexpected case that the map has not yet been initalised, create it. if hs.sessions == nil { - hs.sessions = make(map[string][]plugins.Message) + hs.sessions = make(map[string]HistoryEvent) } - hs.sessions[key] = append(hs.sessions[key], message...) + e, ok := hs.sessions[key] + if !ok { + e = HistoryEvent{} + } + e.LastSeen = time.Now() + e.Messages = append(e.Messages, message...) + hs.sessions[key] = e +} + +// HistoryCleaner is a function that will periodically remove records from the HistoryStore +// that are older than MaxHistoryAge. +func (hs *HistoryStore) HistoryCleaner() { + cleanerTicker := time.NewTicker(CleanerInterval) + go func() { + for range cleanerTicker.C { + hs.Lock() + for k, v := range hs.sessions { + if time.Since(v.LastSeen) > MaxHistoryAge { + delete(hs.sessions, k) + } + } + hs.Unlock() + } + }() } diff --git a/historystore/history_store_test.go b/historystore/history_store_test.go index 7bed948..7028783 100644 --- a/historystore/history_store_test.go +++ b/historystore/history_store_test.go @@ -2,6 +2,7 @@ package historystore import ( "testing" + "time" "github.com/mariocandela/beelzebub/v3/plugins" "github.com/stretchr/testify/assert" @@ -15,7 +16,7 @@ func TestNewHistoryStore(t *testing.T) { func TestHasKey(t *testing.T) { hs := NewHistoryStore() - hs.sessions["testKey"] = []plugins.Message{} + hs.sessions["testKey"] = HistoryEvent{Messages: []plugins.Message{}} assert.True(t, hs.HasKey("testKey")) assert.False(t, hs.HasKey("nonExistentKey")) } @@ -23,7 +24,7 @@ func TestHasKey(t *testing.T) { func TestQuery(t *testing.T) { hs := NewHistoryStore() expectedMessages := []plugins.Message{{Role: "user", Content: "Hello"}} - hs.sessions["testKey"] = expectedMessages + hs.sessions["testKey"] = HistoryEvent{Messages: expectedMessages} actualMessages := hs.Query("testKey") assert.Equal(t, expectedMessages, actualMessages) } @@ -33,9 +34,9 @@ func TestAppend(t *testing.T) { message1 := plugins.Message{Role: "user", Content: "Hello"} message2 := plugins.Message{Role: "assistant", Content: "Hi"} hs.Append("testKey", message1) - assert.Equal(t, []plugins.Message{message1}, hs.sessions["testKey"]) + assert.Equal(t, []plugins.Message{message1}, hs.sessions["testKey"].Messages) hs.Append("testKey", message2) - assert.Equal(t, []plugins.Message{message1, message2}, hs.sessions["testKey"]) + assert.Equal(t, []plugins.Message{message1, message2}, hs.sessions["testKey"].Messages) } func TestAppendNilSessions(t *testing.T) { @@ -43,5 +44,23 @@ func TestAppendNilSessions(t *testing.T) { message1 := plugins.Message{Role: "user", Content: "Hello"} hs.Append("testKey", message1) assert.NotNil(t, hs.sessions) - assert.Equal(t, []plugins.Message{message1}, hs.sessions["testKey"]) + assert.Equal(t, []plugins.Message{message1}, hs.sessions["testKey"].Messages) +} + +func TestHistoryCleaner(t *testing.T) { + hs := NewHistoryStore() + hs.Append("testKey", plugins.Message{Role: "user", Content: "Hello"}) + hs.Append("testKey2", plugins.Message{Role: "user", Content: "Hello"}) + + // Make key older than MaxHistoryAge + e := hs.sessions["testKey"] + e.LastSeen = time.Now().Add(-MaxHistoryAge * 2) + hs.sessions["testKey"] = e + + CleanerInterval = 5 * time.Second // Override for the test. + hs.HistoryCleaner() + time.Sleep(CleanerInterval + (1 * time.Second)) + + assert.False(t, hs.HasKey("testKey")) + assert.True(t, hs.HasKey("testKey2")) } diff --git a/main.go b/main.go index e5dc933..a701d9d 100644 --- a/main.go +++ b/main.go @@ -2,6 +2,8 @@ package main import ( "flag" + "runtime/debug" + "github.com/mariocandela/beelzebub/v3/builder" "github.com/mariocandela/beelzebub/v3/parser" @@ -13,12 +15,20 @@ func main() { quit = make(chan struct{}) configurationsCorePath string configurationsServicesDirectory string + memLimitMiB int ) flag.StringVar(&configurationsCorePath, "confCore", "./configurations/beelzebub.yaml", "Provide the path of configurations core") flag.StringVar(&configurationsServicesDirectory, "confServices", "./configurations/services/", "Directory config services") + flag.IntVar(&memLimitMiB, "memLimitMiB", 100, "Process Memory in MiB (default 100, set to -1 to use system default)") flag.Parse() + if memLimitMiB > 0 { + // SetMemoryLimit takes an int64 value for the number of bytes. + // bytes value = MiB value * 1024 * 1024 + debug.SetMemoryLimit(int64(memLimitMiB * 1024 * 1024)) + } + parser := parser.Init(configurationsCorePath, configurationsServicesDirectory) coreConfigurations, err := parser.ReadConfigurationsCore() diff --git a/parser/configurations_parser.go b/parser/configurations_parser.go index 1691ab9..1320dfb 100644 --- a/parser/configurations_parser.go +++ b/parser/configurations_parser.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "path/filepath" + "regexp" "strings" log "github.com/sirupsen/logrus" @@ -76,11 +77,13 @@ type BeelzebubServiceConfiguration struct { // Command is the struct that contains the configurations of the commands type Command struct { - Regex string `yaml:"regex"` - Handler string `yaml:"handler"` - Headers []string `yaml:"headers"` - StatusCode int `yaml:"statusCode"` - Plugin string `yaml:"plugin"` + RegexStr string `yaml:"regex"` + Regex *regexp.Regexp `yaml:"-"` // This field is parsed, not stored in the config itself. + Handler string `yaml:"handler"` + Headers []string `yaml:"headers"` + StatusCode int `yaml:"statusCode"` + Plugin string `yaml:"plugin"` + Name string `yaml:"name"` } type configurationsParser struct { @@ -140,12 +143,29 @@ func (bp configurationsParser) ReadConfigurationsServices() ([]BeelzebubServiceC return nil, fmt.Errorf("in file %s: %v", filePath, err) } log.Debug(beelzebubServiceConfiguration) + if err := beelzebubServiceConfiguration.CompileCommandRegex(); err != nil { + return nil, fmt.Errorf("in file %s: invalid regex: %v", filePath, err) + } servicesConfiguration = append(servicesConfiguration, *beelzebubServiceConfiguration) } return servicesConfiguration, nil } +// CompileCommandRegex is the method that compiles the regular expression for each configured Command. +func (c *BeelzebubServiceConfiguration) CompileCommandRegex() error { + for i, command := range c.Commands { + if command.RegexStr != "" { + rex, err := regexp.Compile(command.RegexStr) + if err != nil { + return err + } + c.Commands[i].Regex = rex + } + } + return nil +} + func gelAllFilesNameByDirName(dirName string) ([]string, error) { files, err := os.ReadDir(dirName) if err != nil { diff --git a/parser/configurations_parser_test.go b/parser/configurations_parser_test.go index 16d9805..17ac586 100644 --- a/parser/configurations_parser_test.go +++ b/parser/configurations_parser_test.go @@ -3,6 +3,7 @@ package parser import ( "errors" "os" + "regexp" "testing" "github.com/stretchr/testify/assert" @@ -56,6 +57,11 @@ commands: handler: "login" headers: - "Content-Type: text/html" + - name: "wp-admin" + regex: "wp-admin" + handler: "login" + headers: + - "Content-Type: text/html" fallbackCommand: handler: "404 Not Found!" statusCode: 404 @@ -131,12 +137,14 @@ func TestReadConfigurationsServicesValid(t *testing.T) { assert.Equal(t, firstBeelzebubServiceConfiguration.Protocol, "http") assert.Equal(t, firstBeelzebubServiceConfiguration.ApiVersion, "v1") assert.Equal(t, firstBeelzebubServiceConfiguration.Address, ":8080") - assert.Equal(t, len(firstBeelzebubServiceConfiguration.Commands), 1) - assert.Equal(t, len(firstBeelzebubServiceConfiguration.Commands), 1) - assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].Regex, "wp-admin") + assert.Equal(t, len(firstBeelzebubServiceConfiguration.Commands), 2) + assert.Equal(t, len(firstBeelzebubServiceConfiguration.Commands), 2) + assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].RegexStr, "wp-admin") + assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].Regex.String(), "wp-admin") assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].Handler, "login") assert.Equal(t, len(firstBeelzebubServiceConfiguration.Commands[0].Headers), 1) assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].Headers[0], "Content-Type: text/html") + assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[1].Name, "wp-admin") assert.Equal(t, firstBeelzebubServiceConfiguration.FallbackCommand.Handler, "404 Not Found!") assert.Equal(t, firstBeelzebubServiceConfiguration.FallbackCommand.StatusCode, 404) assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.OpenAISecretKey, "qwerty") @@ -199,3 +207,79 @@ func TestReadFileBytesByFilePath(t *testing.T) { assert.Equal(t, "", string(bytes)) } + +func TestCompileCommandRegex(t *testing.T) { + tests := []struct { + name string + config BeelzebubServiceConfiguration + expectedError bool + }{ + { + name: "Valid Regex", + config: BeelzebubServiceConfiguration{ + Commands: []Command{ + {RegexStr: "^/api/v1/.*$"}, + {RegexStr: "wp-admin"}, + }, + }, + expectedError: false, + }, + { + name: "Empty Regex", + config: BeelzebubServiceConfiguration{ + Commands: []Command{ + {RegexStr: ""}, + {RegexStr: ""}, + }, + }, + expectedError: false, + }, + { + name: "Invalid Regex", + config: BeelzebubServiceConfiguration{ + Commands: []Command{ + {RegexStr: "["}, + }, + }, + expectedError: true, + }, + { + name: "Mixed valid and Invalid Regex", + config: BeelzebubServiceConfiguration{ + Commands: []Command{ + {RegexStr: "^/api/v1/.*$"}, + {RegexStr: "["}, + {RegexStr: "test"}, + }, + }, + expectedError: true, + }, + { + name: "No commands", + config: BeelzebubServiceConfiguration{}, + expectedError: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.config.CompileCommandRegex() + + if tt.expectedError { + assert.Error(t, err) + } else { + assert.NoError(t, err) + for _, command := range tt.config.Commands { + if command.RegexStr != "" { + assert.NotNil(t, command.Regex) + _, err := regexp.Compile(command.RegexStr) + assert.NoError(t, err) + + } else { + assert.Nil(t, command.Regex) + } + } + } + }) + } +} diff --git a/plugins/beelzebub-cloud.go b/plugins/beelzebub-cloud.go index 0e397e8..07577e4 100644 --- a/plugins/beelzebub-cloud.go +++ b/plugins/beelzebub-cloud.go @@ -4,6 +4,7 @@ import ( "encoding/json" "errors" "fmt" + "github.com/go-resty/resty/v2" "github.com/mariocandela/beelzebub/v3/parser" "github.com/mariocandela/beelzebub/v3/tracer" @@ -91,6 +92,9 @@ func (beelzebubCloud *beelzebubCloud) GetHoneypotsConfigurations() ([]parser.Bee if err = yaml.Unmarshal([]byte(honeypotConfig.Config), &honeypotsConfig); err != nil { return nil, err } + if err := honeypotsConfig.CompileCommandRegex(); err != nil { + return nil, fmt.Errorf("unable to load service config from cloud: invalid regex: %v", err) + } servicesConfiguration = append(servicesConfiguration, honeypotsConfig) } diff --git a/plugins/beelzebub-cloud_test.go b/plugins/beelzebub-cloud_test.go index 5d36548..7012eb7 100644 --- a/plugins/beelzebub-cloud_test.go +++ b/plugins/beelzebub-cloud_test.go @@ -2,13 +2,15 @@ package plugins import ( "fmt" + "net/http" + "regexp" + "testing" + "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" - "testing" ) func TestBuildSendEventFailValidation(t *testing.T) { @@ -111,8 +113,9 @@ func TestGetHoneypotsConfigurationsWithResults(t *testing.T) { Description: "SSH interactive ChatGPT", Commands: []parser.Command{ { - Regex: "^(.+)$", - Plugin: "LLMHoneypot", + RegexStr: "^(.+)$", + Regex: regexp.MustCompile("^(.+)$"), + Plugin: "LLMHoneypot", }, }, ServerVersion: "OpenSSH", diff --git a/protocols/strategies/HTTP/http.go b/protocols/strategies/HTTP/http.go index ff2480e..97ebd9c 100644 --- a/protocols/strategies/HTTP/http.go +++ b/protocols/strategies/HTTP/http.go @@ -5,7 +5,6 @@ import ( "io" "net" "net/http" - "regexp" "strings" "github.com/mariocandela/beelzebub/v3/parser" @@ -28,22 +27,14 @@ func (httpStrategy HTTPStrategy) Init(servConf parser.BeelzebubServiceConfigurat serverMux := http.NewServeMux() serverMux.HandleFunc("/", func(responseWriter http.ResponseWriter, request *http.Request) { - traceRequest(request, tr, servConf.Description) var matched bool var resp httpResponse var err error for _, command := range servConf.Commands { var err error - matched, err = regexp.MatchString(command.Regex, request.RequestURI) - if err != nil { - log.Errorf("error parsing regex: %s, %s", command.Regex, err.Error()) - resp.StatusCode = 500 - resp.Body = "500 Internal Server Error" - continue - } - + matched = command.Regex.MatchString(request.RequestURI) if matched { - resp, err = buildHTTPResponse(servConf, command, request) + resp, err = buildHTTPResponse(servConf, tr, command, request) if err != nil { log.Errorf("error building http response: %s: %v", request.RequestURI, err) resp.StatusCode = 500 @@ -57,7 +48,7 @@ func (httpStrategy HTTPStrategy) Init(servConf parser.BeelzebubServiceConfigurat if !matched { command := servConf.FallbackCommand if command.Handler != "" || command.Plugin != "" { - resp, err = buildHTTPResponse(servConf, command, request) + resp, err = buildHTTPResponse(servConf, tr, command, request) if err != nil { log.Errorf("error building http response: %s: %v", request.RequestURI, err) resp.StatusCode = 500 @@ -92,12 +83,13 @@ func (httpStrategy HTTPStrategy) Init(servConf parser.BeelzebubServiceConfigurat return nil } -func buildHTTPResponse(servConf parser.BeelzebubServiceConfiguration, command parser.Command, request *http.Request) (httpResponse, error) { +func buildHTTPResponse(servConf parser.BeelzebubServiceConfiguration, tr tracer.Tracer, command parser.Command, request *http.Request) (httpResponse, error) { resp := httpResponse{ Body: command.Handler, Headers: command.Headers, StatusCode: command.StatusCode, } + traceRequest(request, tr, command, servConf.Description) if command.Plugin == plugins.LLMPluginName { llmProvider, err := plugins.FromStringToLLMProvider(servConf.Plugin.LLMProvider) @@ -129,7 +121,7 @@ func buildHTTPResponse(servConf parser.BeelzebubServiceConfiguration, command pa return resp, nil } -func traceRequest(request *http.Request, tr tracer.Tracer, HoneypotDescription string) { +func traceRequest(request *http.Request, tr tracer.Tracer, command parser.Command, HoneypotDescription string) { bodyBytes, err := io.ReadAll(request.Body) body := "" if err == nil { @@ -153,6 +145,7 @@ func traceRequest(request *http.Request, tr tracer.Tracer, HoneypotDescription s SourcePort: port, ID: uuid.New().String(), Description: HoneypotDescription, + Handler: command.Name, } // Capture the TLS details from the request, if provided. if request.TLS != nil { diff --git a/protocols/strategies/SSH/ssh.go b/protocols/strategies/SSH/ssh.go index 722d8b0..64ef661 100644 --- a/protocols/strategies/SSH/ssh.go +++ b/protocols/strategies/SSH/ssh.go @@ -26,6 +26,7 @@ func (sshStrategy *SSHStrategy) Init(servConf parser.BeelzebubServiceConfigurati if sshStrategy.Sessions == nil { sshStrategy.Sessions = historystore.NewHistoryStore() } + go sshStrategy.Sessions.HistoryCleaner() go func() { server := &ssh.Server{ Addr: servConf.Address, @@ -40,14 +41,12 @@ func (sshStrategy *SSHStrategy) Init(servConf parser.BeelzebubServiceConfigurati // Inline SSH command if sess.RawCommand() != "" { + var histories []plugins.Message + if sshStrategy.Sessions.HasKey(sessionKey) { + histories = sshStrategy.Sessions.Query(sessionKey) + } for _, command := range servConf.Commands { - matched, err := regexp.MatchString(command.Regex, sess.RawCommand()) - if err != nil { - log.Errorf("error regex: %s, %s", command.Regex, err.Error()) - continue - } - - if matched { + if command.Regex.MatchString(sess.RawCommand()) { commandOutput := command.Handler if command.Plugin == plugins.LLMPluginName { llmProvider, err := plugins.FromStringToLLMProvider(servConf.Plugin.LLMProvider) @@ -56,11 +55,6 @@ func (sshStrategy *SSHStrategy) Init(servConf parser.BeelzebubServiceConfigurati commandOutput = "command not found" llmProvider = plugins.OpenAI } - - var histories []plugins.Message - if sshStrategy.Sessions.HasKey(sessionKey) { - histories = sshStrategy.Sessions.Query(sessionKey) - } llmHoneypot := plugins.LLMHoneypot{ Histories: histories, OpenAIKey: servConf.Plugin.OpenAISecretKey, @@ -76,11 +70,16 @@ func (sshStrategy *SSHStrategy) Init(servConf parser.BeelzebubServiceConfigurati commandOutput = "command not found" } } + var newEntries []plugins.Message + newEntries = append(newEntries, plugins.Message{Role: plugins.USER.String(), Content: sess.RawCommand()}) + newEntries = append(newEntries, plugins.Message{Role: plugins.ASSISTANT.String(), Content: commandOutput}) + // Append the new entries to the store. + sshStrategy.Sessions.Append(sessionKey, newEntries...) sess.Write(append([]byte(commandOutput), '\n')) tr.TraceEvent(tracer.Event{ - Msg: "New SSH Raw Command Session", + Msg: "SSH Raw Command", Protocol: tracer.SSH.String(), RemoteAddr: sess.RemoteAddr().String(), SourceIp: host, @@ -92,19 +91,7 @@ func (sshStrategy *SSHStrategy) Init(servConf parser.BeelzebubServiceConfigurati Description: servConf.Description, Command: sess.RawCommand(), CommandOutput: commandOutput, - }) - - var histories []plugins.Message - if sshStrategy.Sessions.HasKey(sessionKey) { - histories = sshStrategy.Sessions.Query(sessionKey) - } - 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.Append(sessionKey, histories...) - tr.TraceEvent(tracer.Event{ - Msg: "End SSH Raw Command Session", - Status: tracer.End.String(), - ID: uuidSession.String(), + Handler: command.Name, }) return } @@ -139,13 +126,7 @@ func (sshStrategy *SSHStrategy) Init(servConf parser.BeelzebubServiceConfigurati break } for _, command := range servConf.Commands { - matched, err := regexp.MatchString(command.Regex, commandInput) - if err != nil { - log.Errorf("error regex: %s, %s", command.Regex, err.Error()) - continue - } - - if matched { + if command.Regex.MatchString(commandInput) { commandOutput := command.Handler if command.Plugin == plugins.LLMPluginName { llmProvider, err := plugins.FromStringToLLMProvider(servConf.Plugin.LLMProvider) @@ -168,14 +149,17 @@ func (sshStrategy *SSHStrategy) Init(servConf parser.BeelzebubServiceConfigurati commandOutput = "command not found" } } - - histories = append(histories, plugins.Message{Role: plugins.USER.String(), Content: commandInput}) - histories = append(histories, plugins.Message{Role: plugins.ASSISTANT.String(), Content: commandOutput}) + var newEntries []plugins.Message + newEntries = append(newEntries, plugins.Message{Role: plugins.USER.String(), Content: commandInput}) + newEntries = append(newEntries, plugins.Message{Role: plugins.ASSISTANT.String(), Content: commandOutput}) + // Stash the new entries to the store, and update the history for this running session. + sshStrategy.Sessions.Append(sessionKey, newEntries...) + histories = append(histories, newEntries...) terminal.Write(append([]byte(commandOutput), '\n')) tr.TraceEvent(tracer.Event{ - Msg: "New SSH Terminal Session", + Msg: "SSH Terminal Session Interaction", RemoteAddr: sess.RemoteAddr().String(), SourceIp: host, SourcePort: port, @@ -185,19 +169,18 @@ func (sshStrategy *SSHStrategy) Init(servConf parser.BeelzebubServiceConfigurati ID: uuidSession.String(), Protocol: tracer.SSH.String(), Description: servConf.Description, + Handler: command.Name, }) - break + break // Inner range over commands. } } } - // Add all history events for the terminal session to the store. - // This is done at the end of the session to avoid excess lock operations. - sshStrategy.Sessions.Append(sessionKey, histories...) tr.TraceEvent(tracer.Event{ - Msg: "End SSH Session", - Status: tracer.End.String(), - ID: uuidSession.String(), + Msg: "End SSH Session", + Status: tracer.End.String(), + ID: uuidSession.String(), + Protocol: tracer.SSH.String(), }) }, PasswordHandler: func(ctx ssh.Context, password string) bool { diff --git a/tracer/tracer.go b/tracer/tracer.go index 9b72bdf..1731d9c 100644 --- a/tracer/tracer.go +++ b/tracer/tracer.go @@ -5,10 +5,9 @@ import ( "sync" "time" - log "github.com/sirupsen/logrus" - "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promauto" + log "github.com/sirupsen/logrus" ) // Workers is the number of workers that will @@ -38,6 +37,7 @@ type Event struct { SourceIp string SourcePort string TLSServerName string + Handler string } type (