chrissewell/beelzebub
1
0
Fork 0
mirror of https://github.com/mariocandela/beelzebub.git synced 2025-07-01 18:47:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

Feature: Enhance Performance, Logging & Stability: Precompile Regex, Command Matching, Golang 1.24, History Cleanup & memLimitMiB Flag. (#182)

Browse Source 
* 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:
Bryan Nolen
2025-03-24 05:16:34 +11:00
committed by GitHub
parent 16b012784c
commit d677cd20b9
14 changed files with 249 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
protocols/strategies/HTTP/http.go
View File

@ -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 {

71
protocols/strategies/SSH/ssh.go
View File

@ -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 {