mirror of
https://github.com/mariocandela/beelzebub.git
synced 2025-07-01 18:47:26 -04:00
Feature: Enhance Performance, Logging & Stability: Precompile Regex, Command Matching, Golang 1.24, History Cleanup & memLimitMiB Flag. (#182)
* Feat: Add support for logging which "command" was matched for SSH and HTTP strategies. * Feat: Convert to precompiling regexp at config load time. This allows for errors to be presented to the user during startup, and provides better performance for complex regexp. * Feat:Bump Golang version to latest stable 1.24 * Feat: Add a cleanup routine for HistoryStore, default TTL for events is 1 hour since last interaction. * Feat: Add new command line flag "memLimitMiB" with a default value of 100. --------- Signed-off-by: Bryan Nolen <bryan@arc.net.au> Signed-off-by: Mario Candela <mario.candela.personal@gmail.com> Co-authored-by: Mario Candela <mario.candela.personal@gmail.com>
This commit is contained in:
9
.github/workflows/ci.yml
vendored
9
.github/workflows/ci.yml
vendored
@ -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 %"
|
||||
|
2
.github/workflows/codeql.yml
vendored
2
.github/workflows/codeql.yml
vendored
@ -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
|
||||
|
4
go.mod
4
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
|
||||
|
8
go.sum
8
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=
|
||||
|
@ -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()
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
@ -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"))
|
||||
}
|
||||
|
10
main.go
10
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()
|
||||
|
@ -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 {
|
||||
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
@ -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)
|
||||
}
|
||||
|
||||
|
@ -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",
|
||||
|
@ -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 {
|
||||
|
@ -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 {
|
||||
|
@ -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 (
|
||||
|
Reference in New Issue
Block a user