diff --git a/builder/builder.go b/builder/builder.go index 25aa0a5..daa6f72 100644 --- a/builder/builder.go +++ b/builder/builder.go @@ -3,15 +3,18 @@ package builder import ( "errors" "fmt" - "github.com/mariocandela/beelzebub/v3/parser" - "github.com/mariocandela/beelzebub/v3/plugins" - "github.com/mariocandela/beelzebub/v3/protocols" - "github.com/mariocandela/beelzebub/v3/protocols/strategies" - "github.com/mariocandela/beelzebub/v3/tracer" "io" "net/http" "os" + "github.com/mariocandela/beelzebub/v3/parser" + "github.com/mariocandela/beelzebub/v3/plugins" + "github.com/mariocandela/beelzebub/v3/protocols" + "github.com/mariocandela/beelzebub/v3/protocols/strategies/HTTP" + "github.com/mariocandela/beelzebub/v3/protocols/strategies/SSH" + "github.com/mariocandela/beelzebub/v3/protocols/strategies/TCP" + "github.com/mariocandela/beelzebub/v3/tracer" + "github.com/prometheus/client_golang/prometheus/promhttp" amqp "github.com/rabbitmq/amqp091-go" log "github.com/sirupsen/logrus" @@ -106,9 +109,9 @@ Honeypot Framework, happy hacking!`) }() // Init Protocol strategies - secureShellStrategy := &strategies.SSHStrategy{} - hypertextTransferProtocolStrategy := &strategies.HTTPStrategy{} - transmissionControlProtocolStrategy := &strategies.TCPStrategy{} + secureShellStrategy := &SSH.SSHStrategy{} + hypertextTransferProtocolStrategy := &HTTP.HTTPStrategy{} + transmissionControlProtocolStrategy := &TCP.TCPStrategy{} // Init Tracer strategies, and set the trace strategy default HTTP protocolManager := protocols.InitProtocolManager(b.traceStrategy, hypertextTransferProtocolStrategy) @@ -122,7 +125,7 @@ Honeypot Framework, happy hacking!`) return err } else { if len(honeypotsConfiguration) == 0 { - return errors.New("No honeypots configuration found") + return errors.New("no honeypots configuration found") } b.beelzebubServicesConfiguration = honeypotsConfiguration } @@ -132,20 +135,16 @@ Honeypot Framework, happy hacking!`) switch beelzebubServiceConfiguration.Protocol { case "http": protocolManager.SetProtocolStrategy(hypertextTransferProtocolStrategy) - break case "ssh": protocolManager.SetProtocolStrategy(secureShellStrategy) - break case "tcp": protocolManager.SetProtocolStrategy(transmissionControlProtocolStrategy) - break default: - log.Fatalf("Protocol %s not managed", beelzebubServiceConfiguration.Protocol) - continue + log.Fatalf("protocol %s not managed", beelzebubServiceConfiguration.Protocol) } if err := protocolManager.InitService(beelzebubServiceConfiguration); err != nil { - return errors.New(fmt.Sprintf("Error during init protocol: %s, %s", beelzebubServiceConfiguration.Protocol, err.Error())) + return fmt.Errorf("error during init protocol: %s, %s", beelzebubServiceConfiguration.Protocol, err.Error()) } } diff --git a/go.mod b/go.mod index ec25a3f..1127382 100644 --- a/go.mod +++ b/go.mod @@ -13,6 +13,7 @@ require ( github.com/sirupsen/logrus v1.9.3 github.com/stretchr/testify v1.10.0 golang.org/x/crypto v0.33.0 + golang.org/x/term v0.29.0 gopkg.in/yaml.v3 v3.0.1 ) @@ -33,6 +34,5 @@ require ( github.com/prometheus/procfs v0.15.1 // indirect golang.org/x/net v0.33.0 // indirect golang.org/x/sys v0.30.0 // indirect - golang.org/x/term v0.29.0 // indirect google.golang.org/protobuf v1.34.2 // indirect ) diff --git a/historystore/history_store.go b/historystore/history_store.go new file mode 100644 index 0000000..0df46eb --- /dev/null +++ b/historystore/history_store.go @@ -0,0 +1,47 @@ +package historystore + +import ( + "sync" + + "github.com/mariocandela/beelzebub/v3/plugins" +) + +// HistoryStore is a thread-safe structure for storing Messages used to build LLM Context. +type HistoryStore struct { + sync.RWMutex + sessions map[string][]plugins.Message +} + +// NewHistoryStore returns a prepared HistoryStore +func NewHistoryStore() *HistoryStore { + return &HistoryStore{ + sessions: make(map[string][]plugins.Message), + } +} + +// HasKey returns true if the supplied key exists in the map. +func (hs *HistoryStore) HasKey(key string) bool { + hs.RLock() + defer hs.RUnlock() + _, ok := hs.sessions[key] + return ok +} + +// Query returns the value stored at the map +func (hs *HistoryStore) Query(key string) []plugins.Message { + hs.RLock() + defer hs.RUnlock() + return hs.sessions[key] +} + +// Append will add the slice of Mesages to the entry for the key. +// If the map has not yet been initalised, then a new map is created. +func (hs *HistoryStore) Append(key string, message ...plugins.Message) { + hs.Lock() + 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[key] = append(hs.sessions[key], message...) +} diff --git a/historystore/history_store_test.go b/historystore/history_store_test.go new file mode 100644 index 0000000..7bed948 --- /dev/null +++ b/historystore/history_store_test.go @@ -0,0 +1,47 @@ +package historystore + +import ( + "testing" + + "github.com/mariocandela/beelzebub/v3/plugins" + "github.com/stretchr/testify/assert" +) + +func TestNewHistoryStore(t *testing.T) { + hs := NewHistoryStore() + assert.NotNil(t, hs) + assert.NotNil(t, hs.sessions) +} + +func TestHasKey(t *testing.T) { + hs := NewHistoryStore() + hs.sessions["testKey"] = []plugins.Message{} + assert.True(t, hs.HasKey("testKey")) + assert.False(t, hs.HasKey("nonExistentKey")) +} + +func TestQuery(t *testing.T) { + hs := NewHistoryStore() + expectedMessages := []plugins.Message{{Role: "user", Content: "Hello"}} + hs.sessions["testKey"] = expectedMessages + actualMessages := hs.Query("testKey") + assert.Equal(t, expectedMessages, actualMessages) +} + +func TestAppend(t *testing.T) { + hs := NewHistoryStore() + 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"]) + hs.Append("testKey", message2) + assert.Equal(t, []plugins.Message{message1, message2}, hs.sessions["testKey"]) +} + +func TestAppendNilSessions(t *testing.T) { + hs := &HistoryStore{} + 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"]) +} diff --git a/parser/configurations_parser.go b/parser/configurations_parser.go index 3b6db89..1691ab9 100644 --- a/parser/configurations_parser.go +++ b/parser/configurations_parser.go @@ -62,6 +62,7 @@ type BeelzebubServiceConfiguration struct { Protocol string `yaml:"protocol"` Address string `yaml:"address"` Commands []Command `yaml:"commands"` + FallbackCommand Command `yaml:"fallbackCommand"` ServerVersion string `yaml:"serverVersion"` ServerName string `yaml:"serverName"` DeadlineTimeoutSeconds int `yaml:"deadlineTimeoutSeconds"` diff --git a/parser/configurations_parser_test.go b/parser/configurations_parser_test.go index 1ba472e..16d9805 100644 --- a/parser/configurations_parser_test.go +++ b/parser/configurations_parser_test.go @@ -56,6 +56,9 @@ commands: handler: "login" headers: - "Content-Type: text/html" +fallbackCommand: + handler: "404 Not Found!" + statusCode: 404 plugin: openAISecretKey: "qwerty" llmModel: "llama3" @@ -134,6 +137,8 @@ func TestReadConfigurationsServicesValid(t *testing.T) { 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.FallbackCommand.Handler, "404 Not Found!") + assert.Equal(t, firstBeelzebubServiceConfiguration.FallbackCommand.StatusCode, 404) assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.OpenAISecretKey, "qwerty") assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.LLMModel, "llama3") assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.LLMProvider, "ollama") @@ -176,7 +181,8 @@ func TestGelAllFilesNameByDirNameError(t *testing.T) { files, err := gelAllFilesNameByDirName("nosuchfile") assert.Nil(t, files) - assert.Equal(t, "open nosuchfile: no such file or directory", err.Error()) + // Windows and Linux return slightly different error strings, but share a common prefix, so check for that. + assert.Contains(t, err.Error(), "open nosuchfile: ") } func TestReadFileBytesByFilePath(t *testing.T) { diff --git a/protocols/strategies/HTTP/http.go b/protocols/strategies/HTTP/http.go new file mode 100644 index 0000000..ff2480e --- /dev/null +++ b/protocols/strategies/HTTP/http.go @@ -0,0 +1,186 @@ +package HTTP + +import ( + "fmt" + "io" + "net" + "net/http" + "regexp" + "strings" + + "github.com/mariocandela/beelzebub/v3/parser" + "github.com/mariocandela/beelzebub/v3/plugins" + "github.com/mariocandela/beelzebub/v3/tracer" + + "github.com/google/uuid" + log "github.com/sirupsen/logrus" +) + +type HTTPStrategy struct{} + +type httpResponse struct { + StatusCode int + Headers []string + Body string +} + +func (httpStrategy HTTPStrategy) Init(servConf parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error { + 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 + } + + if matched { + resp, err = buildHTTPResponse(servConf, command, request) + if err != nil { + log.Errorf("error building http response: %s: %v", request.RequestURI, err) + resp.StatusCode = 500 + resp.Body = "500 Internal Server Error" + } + break + } + } + // If none of the main commands matched, and we have a fallback command configured, process it here. + // The regexp is ignored for fallback commands, as they are catch-all for any request. + if !matched { + command := servConf.FallbackCommand + if command.Handler != "" || command.Plugin != "" { + resp, err = buildHTTPResponse(servConf, command, request) + if err != nil { + log.Errorf("error building http response: %s: %v", request.RequestURI, err) + resp.StatusCode = 500 + resp.Body = "500 Internal Server Error" + } + } + } + setResponseHeaders(responseWriter, resp.Headers, resp.StatusCode) + fmt.Fprint(responseWriter, resp.Body) + + }) + go func() { + var err error + // Launch a TLS supporting server if we are supplied a TLS Key and Certificate. + // If relative paths are supplied, they are relative to the CWD of the binary. + // The can be self-signed, only the client will validate this (or not). + if servConf.TLSKeyPath != "" && servConf.TLSCertPath != "" { + err = http.ListenAndServeTLS(servConf.Address, servConf.TLSCertPath, servConf.TLSKeyPath, serverMux) + } else { + err = http.ListenAndServe(servConf.Address, serverMux) + } + if err != nil { + log.Errorf("error during init HTTP Protocol: %v", err) + return + } + }() + + log.WithFields(log.Fields{ + "port": servConf.Address, + "commands": len(servConf.Commands), + }).Infof("Init service: %s", servConf.Description) + return nil +} + +func buildHTTPResponse(servConf parser.BeelzebubServiceConfiguration, command parser.Command, request *http.Request) (httpResponse, error) { + resp := httpResponse{ + Body: command.Handler, + Headers: command.Headers, + StatusCode: command.StatusCode, + } + + if command.Plugin == plugins.LLMPluginName { + llmProvider, err := plugins.FromStringToLLMProvider(servConf.Plugin.LLMProvider) + if err != nil { + log.Errorf("error: %v", err) + resp.Body = "404 Not Found!" + return resp, err + } + + llmHoneypot := plugins.LLMHoneypot{ + Histories: make([]plugins.Message, 0), + OpenAIKey: servConf.Plugin.OpenAISecretKey, + Protocol: tracer.HTTP, + Host: servConf.Plugin.Host, + Model: servConf.Plugin.LLMModel, + Provider: llmProvider, + CustomPrompt: servConf.Plugin.Prompt, + } + llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot) + command := fmt.Sprintf("%s %s", request.Method, request.RequestURI) + + completions, err := llmHoneypotInstance.ExecuteModel(command) + if err != nil { + resp.Body = "404 Not Found!" + return resp, fmt.Errorf("ExecuteModel error: %s, %v", command, err) + } + resp.Body = completions + } + return resp, nil +} + +func traceRequest(request *http.Request, tr tracer.Tracer, HoneypotDescription string) { + bodyBytes, err := io.ReadAll(request.Body) + body := "" + if err == nil { + body = string(bodyBytes) + } + host, port, _ := net.SplitHostPort(request.RemoteAddr) + + event := tracer.Event{ + Msg: "HTTP New request", + RequestURI: request.RequestURI, + Protocol: tracer.HTTP.String(), + HTTPMethod: request.Method, + Body: body, + HostHTTPRequest: request.Host, + UserAgent: request.UserAgent(), + Cookies: mapCookiesToString(request.Cookies()), + Headers: request.Header, + Status: tracer.Stateless.String(), + RemoteAddr: request.RemoteAddr, + SourceIp: host, + SourcePort: port, + ID: uuid.New().String(), + Description: HoneypotDescription, + } + // Capture the TLS details from the request, if provided. + if request.TLS != nil { + event.Msg = "HTTPS New Request" + event.TLSServerName = request.TLS.ServerName + } + tr.TraceEvent(event) +} + +func mapCookiesToString(cookies []*http.Cookie) string { + cookiesString := "" + + for _, cookie := range cookies { + cookiesString += cookie.String() + } + + return cookiesString +} + +func setResponseHeaders(responseWriter http.ResponseWriter, headers []string, statusCode int) { + for _, headerStr := range headers { + keyValue := strings.Split(headerStr, ":") + if len(keyValue) > 1 { + responseWriter.Header().Add(keyValue[0], keyValue[1]) + } + } + // http.StatusText(statusCode): empty string if the code is unknown. + if len(http.StatusText(statusCode)) > 0 { + responseWriter.WriteHeader(statusCode) + } +} diff --git a/protocols/strategies/ssh.go b/protocols/strategies/SSH/ssh.go similarity index 62% rename from protocols/strategies/ssh.go rename to protocols/strategies/SSH/ssh.go index 6c1ecf5..722d8b0 100644 --- a/protocols/strategies/ssh.go +++ b/protocols/strategies/SSH/ssh.go @@ -1,15 +1,17 @@ -package strategies +package SSH import ( "fmt" - "github.com/mariocandela/beelzebub/v3/parser" - "github.com/mariocandela/beelzebub/v3/plugins" - "github.com/mariocandela/beelzebub/v3/tracer" "net" "regexp" "strings" "time" + "github.com/mariocandela/beelzebub/v3/historystore" + "github.com/mariocandela/beelzebub/v3/parser" + "github.com/mariocandela/beelzebub/v3/plugins" + "github.com/mariocandela/beelzebub/v3/tracer" + "github.com/gliderlabs/ssh" "github.com/google/uuid" log "github.com/sirupsen/logrus" @@ -17,65 +19,60 @@ import ( ) type SSHStrategy struct { - Sessions map[string][]plugins.Message + Sessions *historystore.HistoryStore } -func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error { - sshStrategy.Sessions = make(map[string][]plugins.Message) +func (sshStrategy *SSHStrategy) Init(servConf parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error { + if sshStrategy.Sessions == nil { + sshStrategy.Sessions = historystore.NewHistoryStore() + } go func() { server := &ssh.Server{ - Addr: beelzebubServiceConfiguration.Address, - MaxTimeout: time.Duration(beelzebubServiceConfiguration.DeadlineTimeoutSeconds) * time.Second, - IdleTimeout: time.Duration(beelzebubServiceConfiguration.DeadlineTimeoutSeconds) * time.Second, - Version: beelzebubServiceConfiguration.ServerVersion, + Addr: servConf.Address, + MaxTimeout: time.Duration(servConf.DeadlineTimeoutSeconds) * time.Second, + IdleTimeout: time.Duration(servConf.DeadlineTimeoutSeconds) * time.Second, + Version: servConf.ServerVersion, Handler: func(sess ssh.Session) { uuidSession := uuid.New() host, port, _ := net.SplitHostPort(sess.RemoteAddr().String()) - sessionKey := host + sess.User() + sessionKey := "SSH" + host + sess.User() // Inline SSH command if sess.RawCommand() != "" { - for _, command := range beelzebubServiceConfiguration.Commands { + 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()) + log.Errorf("error regex: %s, %s", command.Regex, err.Error()) continue } if matched { commandOutput := command.Handler - if command.Plugin == plugins.LLMPluginName { - - llmProvider, err := plugins.FromStringToLLMProvider(beelzebubServiceConfiguration.Plugin.LLMProvider) - + llmProvider, err := plugins.FromStringToLLMProvider(servConf.Plugin.LLMProvider) if err != nil { - log.Errorf("Error: %s", err.Error()) + log.Errorf("error: %s", err.Error()) commandOutput = "command not found" llmProvider = plugins.OpenAI } - histories := make([]plugins.Message, 0) - - if sshStrategy.Sessions[sessionKey] != nil { - histories = sshStrategy.Sessions[sessionKey] + var histories []plugins.Message + if sshStrategy.Sessions.HasKey(sessionKey) { + histories = sshStrategy.Sessions.Query(sessionKey) } - llmHoneypot := plugins.LLMHoneypot{ Histories: histories, - OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey, + OpenAIKey: servConf.Plugin.OpenAISecretKey, Protocol: tracer.SSH, - Host: beelzebubServiceConfiguration.Plugin.Host, - Model: beelzebubServiceConfiguration.Plugin.LLMModel, + Host: servConf.Plugin.Host, + Model: servConf.Plugin.LLMModel, Provider: llmProvider, - CustomPrompt: beelzebubServiceConfiguration.Plugin.Prompt, + CustomPrompt: servConf.Plugin.Prompt, } - llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot) - if commandOutput, err = llmHoneypotInstance.ExecuteModel(sess.RawCommand()); err != nil { - log.Errorf("Error ExecuteModel: %s, %s", sess.RawCommand(), err.Error()) + log.Errorf("error ExecuteModel: %s, %s", sess.RawCommand(), err.Error()) commandOutput = "command not found" } } @@ -83,7 +80,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze sess.Write(append([]byte(commandOutput), '\n')) tr.TraceEvent(tracer.Event{ - Msg: "New SSH Session", + Msg: "New SSH Raw Command Session", Protocol: tracer.SSH.String(), RemoteAddr: sess.RemoteAddr().String(), SourceIp: host, @@ -92,16 +89,20 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze ID: uuidSession.String(), Environ: strings.Join(sess.Environ(), ","), User: sess.User(), - Description: beelzebubServiceConfiguration.Description, + 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[sessionKey] = histories + sshStrategy.Sessions.Append(sessionKey, histories...) tr.TraceEvent(tracer.Event{ - Msg: "End SSH Session", + Msg: "End SSH Raw Command Session", Status: tracer.End.String(), ID: uuidSession.String(), }) @@ -111,7 +112,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze } tr.TraceEvent(tracer.Event{ - Msg: "New SSH Session", + Msg: "New SSH Terminal Session", Protocol: tracer.SSH.String(), RemoteAddr: sess.RemoteAddr().String(), SourceIp: host, @@ -120,13 +121,13 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze ID: uuidSession.String(), Environ: strings.Join(sess.Environ(), ","), User: sess.User(), - Description: beelzebubServiceConfiguration.Description, + Description: servConf.Description, }) - terminal := term.NewTerminal(sess, buildPrompt(sess.User(), beelzebubServiceConfiguration.ServerName)) + terminal := term.NewTerminal(sess, buildPrompt(sess.User(), servConf.ServerName)) var histories []plugins.Message - if sshStrategy.Sessions[sessionKey] != nil { - histories = sshStrategy.Sessions[sessionKey] + if sshStrategy.Sessions.HasKey(sessionKey) { + histories = sshStrategy.Sessions.Query(sessionKey) } for { @@ -134,43 +135,36 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze if err != nil { break } - if commandInput == "exit" { break } - for _, command := range beelzebubServiceConfiguration.Commands { + 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()) + log.Errorf("error regex: %s, %s", command.Regex, err.Error()) continue } if matched { commandOutput := command.Handler - if command.Plugin == plugins.LLMPluginName { - - llmProvider, err := plugins.FromStringToLLMProvider(beelzebubServiceConfiguration.Plugin.LLMProvider) - + llmProvider, err := plugins.FromStringToLLMProvider(servConf.Plugin.LLMProvider) if err != nil { - log.Errorf("Error: %s, fallback OpenAI", err.Error()) + log.Errorf("error: %s, fallback OpenAI", err.Error()) llmProvider = plugins.OpenAI } - llmHoneypot := plugins.LLMHoneypot{ Histories: histories, - OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey, + OpenAIKey: servConf.Plugin.OpenAISecretKey, Protocol: tracer.SSH, - Host: beelzebubServiceConfiguration.Plugin.Host, - Model: beelzebubServiceConfiguration.Plugin.LLMModel, + Host: servConf.Plugin.Host, + Model: servConf.Plugin.LLMModel, Provider: llmProvider, - CustomPrompt: beelzebubServiceConfiguration.Plugin.Prompt, + CustomPrompt: servConf.Plugin.Prompt, } - llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot) - if commandOutput, err = llmHoneypotInstance.ExecuteModel(commandInput); err != nil { - log.Errorf("Error ExecuteModel: %s, %s", commandInput, err.Error()) + log.Errorf("error ExecuteModel: %s, %s", commandInput, err.Error()) commandOutput = "command not found" } } @@ -190,13 +184,16 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze CommandOutput: commandOutput, ID: uuidSession.String(), Protocol: tracer.SSH.String(), - Description: beelzebubServiceConfiguration.Description, + Description: servConf.Description, }) break } } } - sshStrategy.Sessions[sessionKey] = histories + + // 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(), @@ -207,7 +204,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze host, port, _ := net.SplitHostPort(ctx.RemoteAddr().String()) tr.TraceEvent(tracer.Event{ - Msg: "New SSH attempt", + Msg: "New SSH Login Attempt", Protocol: tracer.SSH.String(), Status: tracer.Stateless.String(), User: ctx.User(), @@ -217,11 +214,11 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze SourceIp: host, SourcePort: port, ID: uuid.New().String(), - Description: beelzebubServiceConfiguration.Description, + Description: servConf.Description, }) - matched, err := regexp.MatchString(beelzebubServiceConfiguration.PasswordRegex, password) + matched, err := regexp.MatchString(servConf.PasswordRegex, password) if err != nil { - log.Errorf("Error regex: %s, %s", beelzebubServiceConfiguration.PasswordRegex, err.Error()) + log.Errorf("error regex: %s, %s", servConf.PasswordRegex, err.Error()) return false } return matched @@ -229,14 +226,14 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze } err := server.ListenAndServe() if err != nil { - log.Errorf("Error during init SSH Protocol: %s", err.Error()) + log.Errorf("error during init SSH Protocol: %s", err.Error()) } }() log.WithFields(log.Fields{ - "port": beelzebubServiceConfiguration.Address, - "commands": len(beelzebubServiceConfiguration.Commands), - }).Infof("GetInstance service %s", beelzebubServiceConfiguration.Protocol) + "port": servConf.Address, + "commands": len(servConf.Commands), + }).Infof("GetInstance service %s", servConf.Protocol) return nil } diff --git a/protocols/strategies/tcp.go b/protocols/strategies/TCP/tcp.go similarity index 61% rename from protocols/strategies/tcp.go rename to protocols/strategies/TCP/tcp.go index d59f32d..dd21255 100644 --- a/protocols/strategies/tcp.go +++ b/protocols/strategies/TCP/tcp.go @@ -1,12 +1,13 @@ -package strategies +package TCP import ( "fmt" - "github.com/mariocandela/beelzebub/v3/parser" - "github.com/mariocandela/beelzebub/v3/tracer" "net" "time" + "github.com/mariocandela/beelzebub/v3/parser" + "github.com/mariocandela/beelzebub/v3/tracer" + "github.com/google/uuid" log "github.com/sirupsen/logrus" ) @@ -14,8 +15,8 @@ import ( type TCPStrategy struct { } -func (tcpStrategy *TCPStrategy) Init(beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error { - listen, err := net.Listen("tcp", beelzebubServiceConfiguration.Address) +func (tcpStrategy *TCPStrategy) Init(servConf parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error { + listen, err := net.Listen("tcp", servConf.Address) if err != nil { log.Errorf("Error during init TCP Protocol: %s", err.Error()) return err @@ -25,8 +26,8 @@ func (tcpStrategy *TCPStrategy) Init(beelzebubServiceConfiguration parser.Beelze for { if conn, err := listen.Accept(); err == nil { go func() { - conn.SetDeadline(time.Now().Add(time.Duration(beelzebubServiceConfiguration.DeadlineTimeoutSeconds) * time.Second)) - conn.Write([]byte(fmt.Sprintf("%s\n", beelzebubServiceConfiguration.Banner))) + conn.SetDeadline(time.Now().Add(time.Duration(servConf.DeadlineTimeoutSeconds) * time.Second)) + conn.Write(fmt.Appendf([]byte{}, "%s\n", servConf.Banner)) buffer := make([]byte, 1024) command := "" @@ -46,7 +47,7 @@ func (tcpStrategy *TCPStrategy) Init(beelzebubServiceConfiguration parser.Beelze SourceIp: host, SourcePort: port, ID: uuid.New().String(), - Description: beelzebubServiceConfiguration.Description, + Description: servConf.Description, }) conn.Close() }() @@ -55,8 +56,8 @@ func (tcpStrategy *TCPStrategy) Init(beelzebubServiceConfiguration parser.Beelze }() log.WithFields(log.Fields{ - "port": beelzebubServiceConfiguration.Address, - "banner": beelzebubServiceConfiguration.Banner, - }).Infof("Init service %s", beelzebubServiceConfiguration.Protocol) + "port": servConf.Address, + "banner": servConf.Banner, + }).Infof("Init service %s", servConf.Protocol) return nil } diff --git a/protocols/strategies/http.go b/protocols/strategies/http.go deleted file mode 100644 index 006ed44..0000000 --- a/protocols/strategies/http.go +++ /dev/null @@ -1,157 +0,0 @@ -package strategies - -import ( - "fmt" - "io" - "net" - "net/http" - "regexp" - "strings" - - "github.com/google/uuid" - "github.com/mariocandela/beelzebub/v3/parser" - "github.com/mariocandela/beelzebub/v3/plugins" - "github.com/mariocandela/beelzebub/v3/tracer" - log "github.com/sirupsen/logrus" -) - -type HTTPStrategy struct { - beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration -} - -func (httpStrategy HTTPStrategy) Init(beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error { - httpStrategy.beelzebubServiceConfiguration = beelzebubServiceConfiguration - serverMux := http.NewServeMux() - - serverMux.HandleFunc("/", func(responseWriter http.ResponseWriter, request *http.Request) { - traceRequest(request, tr, beelzebubServiceConfiguration.Description) - for _, command := range httpStrategy.beelzebubServiceConfiguration.Commands { - matched, err := regexp.MatchString(command.Regex, request.RequestURI) - if err != nil { - log.Errorf("Error regex: %s, %s", command.Regex, err.Error()) - continue - } - - if matched { - responseHTTPBody := command.Handler - - if command.Plugin == plugins.LLMPluginName { - - llmProvider, err := plugins.FromStringToLLMProvider(beelzebubServiceConfiguration.Plugin.LLMProvider) - - if err != nil { - log.Errorf("Error: %s", err.Error()) - responseHTTPBody = "404 Not Found!" - } - - llmHoneypot := plugins.LLMHoneypot{ - Histories: make([]plugins.Message, 0), - OpenAIKey: beelzebubServiceConfiguration.Plugin.OpenAISecretKey, - Protocol: tracer.HTTP, - Host: beelzebubServiceConfiguration.Plugin.Host, - Model: beelzebubServiceConfiguration.Plugin.LLMModel, - Provider: llmProvider, - CustomPrompt: beelzebubServiceConfiguration.Plugin.Prompt, - } - - llmHoneypotInstance := plugins.InitLLMHoneypot(llmHoneypot) - - command := fmt.Sprintf("%s %s", request.Method, request.RequestURI) - - if completions, err := llmHoneypotInstance.ExecuteModel(command); err != nil { - log.Errorf("Error ExecuteModel: %s, %s", command, err.Error()) - responseHTTPBody = "404 Not Found!" - } else { - responseHTTPBody = completions - } - - } - - setResponseHeaders(responseWriter, command.Headers, command.StatusCode) - fmt.Fprint(responseWriter, responseHTTPBody) - break - } - } - }) - go func() { - var err error - // Launch a TLS supporting server if we are supplied a TLS Key and Certificate. - // If relative paths are supplied, they are relative to the CWD of the binary. - // The can be self-signed, only the client will validate this (or not). - if httpStrategy.beelzebubServiceConfiguration.TLSKeyPath != "" && httpStrategy.beelzebubServiceConfiguration.TLSCertPath != "" { - err = http.ListenAndServeTLS( - httpStrategy.beelzebubServiceConfiguration.Address, - httpStrategy.beelzebubServiceConfiguration.TLSCertPath, - httpStrategy.beelzebubServiceConfiguration.TLSKeyPath, - serverMux) - } else { - err = http.ListenAndServe(httpStrategy.beelzebubServiceConfiguration.Address, serverMux) - } - if err != nil { - log.Errorf("Error during init HTTP Protocol: %s", err.Error()) - return - } - }() - - log.WithFields(log.Fields{ - "port": beelzebubServiceConfiguration.Address, - "commands": len(beelzebubServiceConfiguration.Commands), - }).Infof("Init service: %s", beelzebubServiceConfiguration.Description) - return nil -} - -func traceRequest(request *http.Request, tr tracer.Tracer, HoneypotDescription string) { - bodyBytes, err := io.ReadAll(request.Body) - body := "" - if err == nil { - body = string(bodyBytes) - } - host, port, _ := net.SplitHostPort(request.RemoteAddr) - - event := tracer.Event{ - Msg: "HTTP New request", - RequestURI: request.RequestURI, - Protocol: tracer.HTTP.String(), - HTTPMethod: request.Method, - Body: body, - HostHTTPRequest: request.Host, - UserAgent: request.UserAgent(), - Cookies: mapCookiesToString(request.Cookies()), - Headers: request.Header, - Status: tracer.Stateless.String(), - RemoteAddr: request.RemoteAddr, - SourceIp: host, - SourcePort: port, - ID: uuid.New().String(), - Description: HoneypotDescription, - } - // Capture the TLS details from the request, if provided. - if request.TLS != nil { - event.Msg = "HTTPS New Request" - event.TLSServerName = request.TLS.ServerName - } - tr.TraceEvent(event) -} - -func mapCookiesToString(cookies []*http.Cookie) string { - cookiesString := "" - - for _, cookie := range cookies { - cookiesString += cookie.String() - } - - return cookiesString -} - -func setResponseHeaders(responseWriter http.ResponseWriter, headers []string, statusCode int) { - for _, headerStr := range headers { - keyValue := strings.Split(headerStr, ":") - if len(keyValue) > 1 { - responseWriter.Header().Add(keyValue[0], keyValue[1]) - } - } - // http.StatusText(statusCode): empty string if the code is unknown. - if len(http.StatusText(statusCode)) > 0 { - responseWriter.WriteHeader(statusCode) - } -}