From c032489522de4844169b678caaf4f40ac50d96c5 Mon Sep 17 00:00:00 2001 From: Mario Candela Date: Tue, 1 Jul 2025 23:07:49 +0200 Subject: [PATCH] Feat: add MCP honeypot support (#199) * Add MCP honeypot * Add http headers in plain text * Improve code coverage * Refactor README.md --- .github/workflows/ci.yml | 2 +- README.md | 168 +++++++++++++++++++------- builder/builder.go | 4 + configurations/services/mcp-8000.yaml | 41 +++++++ go.mod | 3 + go.sum | 6 + parser/configurations_parser.go | 15 +++ parser/configurations_parser_test.go | 17 +++ protocols/strategies/HTTP/http.go | 15 ++- protocols/strategies/MCP/mcp.go | 86 +++++++++++++ tracer/tracer.go | 15 ++- tracer/tracer_test.go | 4 + 12 files changed, 326 insertions(+), 50 deletions(-) create mode 100644 configurations/services/mcp-8000.yaml create mode 100644 protocols/strategies/MCP/mcp.go diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 04747da..7242281 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -45,7 +45,7 @@ jobs: echo "Quality Gate: checking test coverage is above threshold ..." echo "Threshold : $TESTCOVERAGE_THRESHOLD %" # Excluded the concrete strategy from the unit test coverage, because covered by integration tests - cat coverage.tmp.out | grep -v "ssh.go" | grep -v "http.go" | grep -v "tcp.go" | grep -v "builder.go" | grep -v "director.go" > coverage.out + cat coverage.tmp.out | grep -v "mcp.go" | grep -v "ssh.go" | grep -v "http.go" | grep -v "tcp.go" | grep -v "builder.go" | grep -v "director.go" > coverage.out totalCoverage=`go tool cover -func=coverage.out | grep total | grep -Eo '[0-9]+\.[0-9]+'` echo "Current test coverage : $totalCoverage %" if (( $(echo "$totalCoverage $TESTCOVERAGE_THRESHOLD" | awk '{print ($1 > $2)}') )); then diff --git a/README.md b/README.md index 7d4e6eb..2df558d 100644 --- a/README.md +++ b/README.md @@ -12,22 +12,45 @@ Beelzebub is an advanced honeypot framework designed to provide a highly secure Beelzebub Logo -## LLM Honeypot +## Key Features + +Beelzebub offers a wide range of features to enhance your honeypot environment: + +- Low-code configuration: YAML-based, modular service definition +- LLM integration: The LLM convincingly simulates a real system, creating high-interaction honeypot experiences, while actually maintaining low-interaction architecture for enhanced security and easy management. +- Multi-protocol support: SSH, HTTP, TCP, MCP(Detect prompt injection against LLM agents) +- Prometheus metrics & observability +- Docker & Kubernetes ready +- ELK stack ready, docs: [Official ELK integration](https://www.elastic.co/docs/reference/integrations/beelzebub) + +## LLM SSH Honeypot Demo [![asciicast](https://asciinema.org/a/665295.svg)](https://asciinema.org/a/665295) +## Code Quality -## Telegram Bot for Real-Time Attacks +We are strongly committed to maintaining high code quality in the Beelzebub project. Our development workflow includes comprehensive testing, code reviews, static analysis, and continuous integration to ensure the reliability and maintainability of the codebase. -Stay updated on real-time attacks by joining our dedicated Telegram channel: [Telegram Channel](https://t.me/beelzebubhoneypot) +### What We Do -## Examples +* **Automated Testing:** + Both unit and integration tests are run on every pull request to catch regressions and ensure stability. -To better understand the capabilities of Beelzebub, you can explore our example repository: [mariocandela/beelzebub-example](https://github.com/mariocandela/beelzebub-example) +* **Static Analysis:** + We use tools like Go Report Card and CodeQL to automatically check for code quality, style, and security issues. + +* **Code Coverage:** + Our test coverage is monitored with [Codecov](https://codecov.io/gh/mariocandela/beelzebub), and we aim for extensive coverage of all core components. + +* **Continuous Integration:** + Every commit triggers automated CI pipelines on GitHub Actions, which run all tests and quality checks. + +* **Code Reviews:** + All new contributions undergo peer review to maintain consistency and high standards across the project. ## Quick Start -We provide two quick start options for build and run Beelzebub: using Docker Compose or the Go compiler. +You can run Beelzebub via Docker, Go compiler(cross device), or Helm (Kubernetes). ### Using Docker Compose @@ -43,6 +66,7 @@ We provide two quick start options for build and run Beelzebub: using Docker Com $ docker-compose up -d ``` + ### Using Go Compiler 1. Download the necessary Go modules: @@ -78,41 +102,6 @@ We provide two quick start options for build and run Beelzebub: using Docker Com ```bash $ helm upgrade beelzebub ./beelzebub-chart ``` -## Testing - -We provide two types of tests: unit tests and integration tests. - -### Unit Tests - -To run unit tests: - -```bash -$ make test.unit -``` - -### Integration Tests - -To run integration tests: - -```bash -$ make test.dependencies.start -$ make test.integration -$ make test.dependencies.down -``` - -## Key Features - -Beelzebub offers a wide range of features to enhance your honeypot environment: - -- Support for Ollama -- Support for OpenAI -- SSH Honeypot -- HTTP Honeypot -- TCP Honeypot -- Prometheus openmetrics integration -- Docker integration -- RabbitMQ integration -- kubernetes ## Example Configuration @@ -126,7 +115,73 @@ $ ./beelzebub --confCore ./configurations/beelzebub.yaml --confServices ./config Here are some example configurations for different honeypot scenarios: -#### Example HTTP Honeypot on Port 80 +### MCP Honeypot + +#### Why choose an MCP Honeypot? + +An MCP honeypot is a **decoy tool** that the agent should never invoke under normal circumstances. Integrating this strategy into your agent pipeline offers three key benefits: + +* **Real-time detection of guardrail bypass attempts.** + + Instantly identify when a prompt injection attack successfully convinces the agent to invoke a restricted tool. +* **Automatic collection of real attack prompts for guardrail fine-tuning.** + + Every activation logs genuine malicious prompts, enabling continuous improvement of your filtering mechanisms. +* **Continuous monitoring of attack trends through key metrics (HAR, TPR, MTP).** + + Track exploit frequency and system resilience using objective, actionable measurements. + +##### Example MCP Honeypot Configuration + +###### mcp-8000.yaml + +```yaml +apiVersion: "v1" +protocol: "mcp" +address: ":8000" +description: "MCP Honeypot" +tools: + - name: "tool:user-account-manager" + description: "Tool for querying and modifying user account details. Requires administrator privileges." + params: + - name: "user_id" + description: "The ID of the user account to manage." + - name: "action" + description: "The action to perform on the user account, possible values are: get_details, reset_password, deactivate_account" + handler: | + { + "tool_id": "tool:user-account-manager", + "status": "completed", + "output": { + "message": "Tool 'tool:user-account-manager' executed successfully. Results are pending internal processing and will be logged.", + "result": { + "operation_status": "success", + "details": "email: kirsten@gmail.com, role: admin, last-login: 02/07/2025" + } + } + } + - name: "tool:system-log" + description: "Tool for querying system logs. Requires administrator privileges." + params: + - name: "filter" + description: "The input used to filter the logs." + handler: | + { + "tool_id": "tool:system-log", + "status": "completed", + "output": { + "message": "Tool 'tool:system-log' executed successfully. Results are pending internal processing and will be logged.", + "result": { + "operation_status": "success", + "details": "Info: email: kirsten@gmail.com, last-login: 02/07/2025" + } + } + } +``` + +#### Invoke remotely: beelzebub:port/mcp (Streamable HTTPServer). + +### HTTP Honeypot ###### http-80.yaml @@ -191,7 +246,7 @@ commands: statusCode: 404 ``` -#### Example HTTP Honeypot on Port 8080 +### HTTP Honeypot ###### http-8080.yaml @@ -209,7 +264,7 @@ commands: statusCode: 401 ``` -#### Example SSH Honeypot +### SSH Honeypot ###### LLM Honeypots @@ -273,7 +328,7 @@ plugin: prompt: "You will act as an Ubuntu Linux terminal. The user will type commands, and you are to reply with what the terminal should show. Your responses must be contained within a single code block." ``` -###### SSH Honeypot on Port 22 +###### SSH Honeypot ###### ssh-22.yaml @@ -307,6 +362,29 @@ passwordRegex: "^(root|qwerty|Smoker666)$" deadlineTimeoutSeconds: 60 ``` +## Testing + +Maintaining excellent code quality is essential for security-focused projects like Beelzebub. We welcome all contributors who share our commitment to robust, readable, and reliable code! + +### Unit Tests + +For contributor, we have a comprehensive suite of unit/integration tests that cover the core functionality of Beelzebub. To run the unit tests, use the following command: + +```bash +$ make test.unit +``` + +### Integration Tests + +To run integration tests: + +```bash +$ make test.dependencies.start +$ make test.integration +$ make test.dependencies.down +``` + + ## Roadmap Our future plans for Beelzebub include developing it into a robust PaaS platform. diff --git a/builder/builder.go b/builder/builder.go index daa6f72..954561c 100644 --- a/builder/builder.go +++ b/builder/builder.go @@ -3,6 +3,7 @@ package builder import ( "errors" "fmt" + "github.com/mariocandela/beelzebub/v3/protocols/strategies/MCP" "io" "net/http" "os" @@ -112,6 +113,7 @@ Honeypot Framework, happy hacking!`) secureShellStrategy := &SSH.SSHStrategy{} hypertextTransferProtocolStrategy := &HTTP.HTTPStrategy{} transmissionControlProtocolStrategy := &TCP.TCPStrategy{} + modelContextProtocolStrategy := &MCP.MCPStrategy{} // Init Tracer strategies, and set the trace strategy default HTTP protocolManager := protocols.InitProtocolManager(b.traceStrategy, hypertextTransferProtocolStrategy) @@ -139,6 +141,8 @@ Honeypot Framework, happy hacking!`) protocolManager.SetProtocolStrategy(secureShellStrategy) case "tcp": protocolManager.SetProtocolStrategy(transmissionControlProtocolStrategy) + case "mcp": + protocolManager.SetProtocolStrategy(modelContextProtocolStrategy) default: log.Fatalf("protocol %s not managed", beelzebubServiceConfiguration.Protocol) } diff --git a/configurations/services/mcp-8000.yaml b/configurations/services/mcp-8000.yaml new file mode 100644 index 0000000..9a8025d --- /dev/null +++ b/configurations/services/mcp-8000.yaml @@ -0,0 +1,41 @@ +apiVersion: "v1" +protocol: "mcp" +address: ":8000" +description: "MCP Honeypot" +tools: + - name: "tool:user-account-manager" + description: "Tool for querying and modifying user account details. Requires administrator privileges." + params: + - name: "user_id" + description: "The ID of the user account to manage." + - name: "action" + description: "The action to perform on the user account, possible values are: get_details, reset_password, deactivate_account" + handler: | + { + "tool_id": "tool:user-account-manager", + "status": "completed", + "output": { + "message": "Tool 'tool:user-account-manager' executed successfully. Results are pending internal processing and will be logged.", + "result": { + "operation_status": "success", + "details": "email: kirsten_12345@gmail.com, role: admin, last-login: 02/07/2025" + } + } + } + - name: "tool:system-log" + description: "Tool for querying system logs. Requires administrator privileges." + params: + - name: "filter" + description: "The input used to filter the logs." + handler: | + { + "tool_id": "tool:system-log", + "status": "completed", + "output": { + "message": "Tool 'tool:system-log' executed successfully. Results are pending internal processing and will be logged.", + "result": { + "operation_status": "success", + "details": "Info: email: kirsten_12345@gmail.com, last-login: 02/07/2025" + } + } + } \ No newline at end of file diff --git a/go.mod b/go.mod index c30307a..f14cce0 100644 --- a/go.mod +++ b/go.mod @@ -26,6 +26,7 @@ require ( github.com/davecgh/go-spew v1.1.1 // indirect github.com/kr/fs v0.1.0 // indirect github.com/kr/text v0.2.0 // indirect + github.com/mark3labs/mcp-go v0.32.0 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pkg/sftp v1.13.5 // indirect @@ -33,6 +34,8 @@ require ( github.com/prometheus/client_model v0.6.1 // indirect github.com/prometheus/common v0.62.0 // indirect github.com/prometheus/procfs v0.15.1 // indirect + github.com/spf13/cast v1.7.1 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect golang.org/x/net v0.38.0 // indirect golang.org/x/sys v0.33.0 // indirect google.golang.org/protobuf v1.36.5 // indirect diff --git a/go.sum b/go.sum index 44fd840..40382da 100644 --- a/go.sum +++ b/go.sum @@ -28,6 +28,8 @@ 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/mark3labs/mcp-go v0.32.0 h1:fgwmbfL2gbd67obg57OfV2Dnrhs1HtSdlY/i5fn7MU8= +github.com/mark3labs/mcp-go v0.32.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4= github.com/maxatome/go-testdeep v1.14.0 h1:rRlLv1+kI8eOI3OaBXZwb3O7xY3exRzdW5QyX48g9wI= github.com/maxatome/go-testdeep v1.14.0/go.mod h1:lPZc/HAcJMP92l7yI6TRz1aZN5URwUBUAfUNvrclaNM= github.com/melbahja/goph v1.4.0 h1:z0PgDbBFe66lRYl3v5dGb9aFgPy0kotuQ37QOwSQFqs= @@ -54,10 +56,14 @@ github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjR 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/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y= +github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= 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= diff --git a/parser/configurations_parser.go b/parser/configurations_parser.go index 1320dfb..5c3263a 100644 --- a/parser/configurations_parser.go +++ b/parser/configurations_parser.go @@ -63,6 +63,7 @@ type BeelzebubServiceConfiguration struct { Protocol string `yaml:"protocol"` Address string `yaml:"address"` Commands []Command `yaml:"commands"` + Tools []Tool `yaml:"tools"` FallbackCommand Command `yaml:"fallbackCommand"` ServerVersion string `yaml:"serverVersion"` ServerName string `yaml:"serverName"` @@ -86,6 +87,20 @@ type Command struct { Name string `yaml:"name"` } +// Tool is the struct that contains the configurations of the MCP Honeypot +type Tool struct { + Name string `yaml:"name"` + Description string `yaml:"description"` + Params []Param `yaml:"params"` + Handler string `yaml:"handler"` +} + +// Param is the struct that contains the configurations of the parameters of the tools +type Param struct { + Name string `yaml:"name"` + Description string `yaml:"description"` +} + type configurationsParser struct { configurationsCorePath string configurationsServicesDirectory string diff --git a/parser/configurations_parser_test.go b/parser/configurations_parser_test.go index 17ac586..5945a8e 100644 --- a/parser/configurations_parser_test.go +++ b/parser/configurations_parser_test.go @@ -52,6 +52,15 @@ protocol: "http" address: ":8080" tlsCertPath: "/tmp/cert.crt" tlsKeyPath: "/tmp/cert.key" +tools: + - name: "tool:user-account-manager" + description: "Tool for querying and modifying user account details. Requires administrator privileges." + params: + - name: "user_id" + description: "The ID of the user account to manage." + - name: "action" + description: "The action to perform on the user account, possible values are: get_details, reset_password, deactivate_account" + handler: "reset_password ok" commands: - regex: "wp-admin" handler: "login" @@ -154,6 +163,14 @@ func TestReadConfigurationsServicesValid(t *testing.T) { assert.Equal(t, firstBeelzebubServiceConfiguration.Plugin.Prompt, "hello world") assert.Equal(t, firstBeelzebubServiceConfiguration.TLSCertPath, "/tmp/cert.crt") assert.Equal(t, firstBeelzebubServiceConfiguration.TLSKeyPath, "/tmp/cert.key") + assert.Equal(t, firstBeelzebubServiceConfiguration.TLSKeyPath, "/tmp/cert.key") + assert.Equal(t, len(firstBeelzebubServiceConfiguration.Tools), 1) + assert.Equal(t, firstBeelzebubServiceConfiguration.Tools[0].Name, "tool:user-account-manager") + assert.Equal(t, firstBeelzebubServiceConfiguration.Tools[0].Description, "Tool for querying and modifying user account details. Requires administrator privileges.") + assert.Equal(t, len(firstBeelzebubServiceConfiguration.Tools[0].Params), 2) + assert.Equal(t, firstBeelzebubServiceConfiguration.Tools[0].Params[0].Name, "user_id") + assert.Equal(t, firstBeelzebubServiceConfiguration.Tools[0].Params[0].Description, "The ID of the user account to manage.") + assert.Equal(t, firstBeelzebubServiceConfiguration.Tools[0].Handler, "reset_password ok") } func TestGelAllFilesNameByDirName(t *testing.T) { diff --git a/protocols/strategies/HTTP/http.go b/protocols/strategies/HTTP/http.go index 97ebd9c..bbb13dc 100644 --- a/protocols/strategies/HTTP/http.go +++ b/protocols/strategies/HTTP/http.go @@ -138,7 +138,8 @@ func traceRequest(request *http.Request, tr tracer.Tracer, command parser.Comman HostHTTPRequest: request.Host, UserAgent: request.UserAgent(), Cookies: mapCookiesToString(request.Cookies()), - Headers: request.Header, + Headers: mapHeaderToString(request.Header), + HeadersMap: request.Header, Status: tracer.Stateless.String(), RemoteAddr: request.RemoteAddr, SourceIp: host, @@ -155,6 +156,18 @@ func traceRequest(request *http.Request, tr tracer.Tracer, command parser.Comman tr.TraceEvent(event) } +func mapHeaderToString(headers http.Header) string { + headersString := "" + + for key := range headers { + for _, values := range headers[key] { + headersString += fmt.Sprintf("[Key: %s, values: %s],", key, values) + } + } + + return headersString +} + func mapCookiesToString(cookies []*http.Cookie) string { cookiesString := "" diff --git a/protocols/strategies/MCP/mcp.go b/protocols/strategies/MCP/mcp.go new file mode 100644 index 0000000..6daad34 --- /dev/null +++ b/protocols/strategies/MCP/mcp.go @@ -0,0 +1,86 @@ +package MCP + +import ( + "context" + "fmt" + "github.com/google/uuid" + "github.com/mariocandela/beelzebub/v3/parser" + "github.com/mariocandela/beelzebub/v3/tracer" + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + log "github.com/sirupsen/logrus" + "net" + "net/http" +) + +type remoteAddrCtxKey struct{} + +type MCPStrategy struct { +} + +func (mcpStrategy *MCPStrategy) Init(servConf parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error { + mcpServer := server.NewMCPServer( + servConf.Description, + "1.0.0", + server.WithToolCapabilities(false), + ) + + for _, toolConfig := range servConf.Tools { + if toolConfig.Params == nil || len(toolConfig.Params) == 0 { + log.Errorf("Tool %s has no parameters defined", toolConfig.Name) + continue + } + + opts := []mcp.ToolOption{ + mcp.WithDescription(toolConfig.Description), + } + + for _, param := range toolConfig.Params { + opts = append(opts, + mcp.WithString( + param.Name, + mcp.Required(), + mcp.Description(param.Description), + ), + ) + } + + tool := mcp.NewTool(toolConfig.Name, opts...) + + mcpServer.AddTool(tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + host, port, _ := net.SplitHostPort(ctx.Value(remoteAddrCtxKey{}).(string)) + + tr.TraceEvent(tracer.Event{ + Msg: "New MCP tool invocation", + Protocol: tracer.MCP.String(), + Status: tracer.Stateless.String(), + RemoteAddr: ctx.Value(remoteAddrCtxKey{}).(string), + SourceIp: host, + SourcePort: port, + ID: uuid.New().String(), + Description: servConf.Description, + Command: fmt.Sprintf("%s|%s", request.Params.Name, request.Params.Arguments), + CommandOutput: toolConfig.Handler, + }) + return mcp.NewToolResultText(toolConfig.Handler), nil + }) + } + + go func() { + httpServer := server.NewStreamableHTTPServer( + mcpServer, + server.WithHTTPContextFunc(func(ctx context.Context, r *http.Request) context.Context { + return context.WithValue(ctx, remoteAddrCtxKey{}, r.RemoteAddr) + }), + ) + if err := httpServer.Start(servConf.Address); err != nil { + log.Errorf("Failed to start MCP server on %s: %v", servConf.Address, err) + return + } + }() + log.WithFields(log.Fields{ + "port": servConf.Address, + "description": servConf.Description, + }).Infof("Init service %s", servConf.Protocol) + return nil +} diff --git a/tracer/tracer.go b/tracer/tracer.go index 1731d9c..e08f80c 100644 --- a/tracer/tracer.go +++ b/tracer/tracer.go @@ -10,7 +10,6 @@ import ( log "github.com/sirupsen/logrus" ) -// Workers is the number of workers that will const Workers = 5 type Event struct { @@ -26,7 +25,8 @@ type Event struct { User string Password string Client string - Headers map[string][]string + Headers string + HeadersMap map[string][]string Cookies string UserAgent string HostHTTPRequest string @@ -49,10 +49,11 @@ const ( HTTP Protocol = iota SSH TCP + MCP ) func (protocol Protocol) String() string { - return [...]string{"HTTP", "SSH", "TCP"}[protocol] + return [...]string{"HTTP", "SSH", "TCP", "MCP"}[protocol] } const ( @@ -79,6 +80,7 @@ type tracer struct { eventsSSHTotal prometheus.Counter eventsTCPTotal prometheus.Counter eventsHTTPTotal prometheus.Counter + eventsMCPTotal prometheus.Counter } var lock = &sync.Mutex{} @@ -113,6 +115,11 @@ func GetInstance(defaultStrategy Strategy) *tracer { Name: "http_events_total", Help: "The total number of HTTP events", }), + eventsMCPTotal: promauto.NewCounter(prometheus.CounterOpts{ + Namespace: "beelzebub", + Name: "mcp_events_total", + Help: "The total number of MCP events", + }), } for i := 0; i < Workers; i++ { @@ -149,6 +156,8 @@ func (tracer *tracer) updatePrometheusCounters(protocol string) { tracer.eventsSSHTotal.Inc() case TCP.String(): tracer.eventsTCPTotal.Inc() + case MCP.String(): + tracer.eventsMCPTotal.Inc() } tracer.eventsTotal.Inc() } diff --git a/tracer/tracer_test.go b/tracer/tracer_test.go index 2abf99a..dfee8fd 100644 --- a/tracer/tracer_test.go +++ b/tracer/tracer_test.go @@ -106,6 +106,7 @@ func TestUpdatePrometheusCounters(t *testing.T) { eventsSSHTotal: mockCounter{}, eventsTCPTotal: mockCounter{}, eventsHTTPTotal: mockCounter{}, + eventsMCPTotal: mockCounter{}, } tracer.updatePrometheusCounters(SSH.String()) @@ -116,4 +117,7 @@ func TestUpdatePrometheusCounters(t *testing.T) { tracer.updatePrometheusCounters(TCP.String()) assert.Equal(t, 6, counter) + + tracer.updatePrometheusCounters(MCP.String()) + assert.Equal(t, 8, counter) }