refactor: Improve Go docs (#74)

* add go docs, package: parser
* add go docs, package: protocols
* add go docs, package: tracer

.github/workflows/ci.yml

@@ -51,6 +51,11 @@ jobs:
             exit 1
         fi
 
+    - name: Upload coverage reports to Codecov
+      uses: codecov/codecov-action@v3
+      env:
+        CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+
     - name: Start integration test dependencies
       run: |
         make test.dependencies.start

README.md

@@ -1,6 +1,10 @@
 # Beelzebub
 
 [![CI](https://github.com/mariocandela/beelzebub/actions/workflows/ci.yml/badge.svg)](https://github.com/mariocandela/beelzebub/actions/workflows/ci.yml) [![Docker](https://github.com/mariocandela/beelzebub/actions/workflows/docker-image.yml/badge.svg)](https://github.com/mariocandela/beelzebub/actions/workflows/docker-image.yml) [![codeql](https://github.com/mariocandela/beelzebub/actions/workflows/codeql.yml/badge.svg)](https://github.com/mariocandela/beelzebub/actions/workflows/codeql.yml)
+[![Go Report Card](https://goreportcard.com/badge/github.com/mariocandela/beelzebub/v3)](https://goreportcard.com/report/github.com/mariocandela/beelzebub/v3)
+[![codecov](https://codecov.io/gh/mariocandela/beelzebub/graph/badge.svg?token=8XTK7D4WHE)](https://codecov.io/gh/mariocandela/beelzebub)
+[![Go Reference](https://pkg.go.dev/badge/github.com/mariocandela/beelzebub/v3.svg)](https://pkg.go.dev/github.com/mariocandela/beelzebub/v3)
+
 
 ## Overview
 
@@ -90,7 +94,9 @@ $ make test.unit
 To run integration tests:
 
 ```bash
+$ make test.dependencies.start
 $ make test.integration
+$ make test.dependencies.down
 ```
 
 ## Key Features

builder/builder.go

@@ -3,10 +3,10 @@ package builder
 import (
 	"errors"
 	"fmt"
-	"github.com/mariocandela/beelzebub/parser"
+	"github.com/mariocandela/beelzebub/v3/parser"
-	"github.com/mariocandela/beelzebub/protocols"
+	"github.com/mariocandela/beelzebub/v3/protocols"
-	"github.com/mariocandela/beelzebub/protocols/strategies"
+	"github.com/mariocandela/beelzebub/v3/protocols/strategies"
-	"github.com/mariocandela/beelzebub/tracer"
+	"github.com/mariocandela/beelzebub/v3/tracer"
 	"io"
 	"net/http"
 	"os"

builder/director.go

@@ -3,8 +3,8 @@ package builder
 import (
 	"context"
 	"encoding/json"
-	"github.com/mariocandela/beelzebub/parser"
+	"github.com/mariocandela/beelzebub/v3/parser"
-	"github.com/mariocandela/beelzebub/tracer"
+	"github.com/mariocandela/beelzebub/v3/tracer"
 
 	amqp "github.com/rabbitmq/amqp091-go"
 	log "github.com/sirupsen/logrus"

go.mod

@@ -1,4 +1,4 @@
-module github.com/mariocandela/beelzebub
+module github.com/mariocandela/beelzebub/v3
 
 go 1.20
 
@@ -30,7 +30,7 @@ require (
 	github.com/prometheus/client_model v0.4.1-0.20230718164431-9a2bf3000d16 // indirect
 	github.com/prometheus/common v0.44.0 // indirect
 	github.com/prometheus/procfs v0.11.1 // indirect
-	golang.org/x/net v0.15.0 // indirect
+	golang.org/x/net v0.17.0 // indirect
 	golang.org/x/sys v0.13.0 // indirect
 	golang.org/x/term v0.13.0 // indirect
 	google.golang.org/protobuf v1.31.0 // indirect

go.sum

@@ -79,8 +79,9 @@ golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug
 golang.org/x/net v0.0.0-20220826154423-83b083e8dc8b/go.mod h1:YDH+HFinaLZZlnHAfSS6ZXJJ9M9t4Dl22yv3iI2vPwk=
 golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
 golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
-golang.org/x/net v0.15.0 h1:ugBLEUaxABaB5AJqW9enI0ACdci2RUd4eP51NTBvuJ8=
 golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk=
+golang.org/x/net v0.17.0 h1:pVaXccu2ozPjCXewfr1S7xza/zcXTity9cCdXQYSjIM=
+golang.org/x/net v0.17.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE=
 golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=

integration_test/integration_test.go

@@ -2,9 +2,9 @@ package integration
 
 import (
 	"encoding/json"
-	"github.com/mariocandela/beelzebub/builder"
+	"github.com/mariocandela/beelzebub/v3/builder"
-	"github.com/mariocandela/beelzebub/parser"
+	"github.com/mariocandela/beelzebub/v3/parser"
-	"github.com/mariocandela/beelzebub/tracer"
+	"github.com/mariocandela/beelzebub/v3/tracer"
 	"net"
 	"net/http"
 	"os"

main.go

@@ -2,8 +2,8 @@ package main
 
 import (
 	"flag"
-	"github.com/mariocandela/beelzebub/builder"
+	"github.com/mariocandela/beelzebub/v3/builder"
-	"github.com/mariocandela/beelzebub/parser"
+	"github.com/mariocandela/beelzebub/v3/parser"
 
 	log "github.com/sirupsen/logrus"
 )

parser/configurations_parser.go

@@ -1,8 +1,8 @@
+// Package parser is responsible for parsing the configurations of the core and honeypot service
 package parser
 
 import (
 	"fmt"
-	"io/ioutil"
 	"os"
 	"path/filepath"
 	"strings"
@@ -11,6 +11,7 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
+// BeelzebubCoreConfigurations is the struct that contains the configurations of the core
 type BeelzebubCoreConfigurations struct {
 	Core struct {
 		Logging    Logging    `yaml:"logging"`
@@ -19,6 +20,7 @@ type BeelzebubCoreConfigurations struct {
 	}
 }
 
+// Logging is the struct that contains the configurations of the logging
 type Logging struct {
 	Debug               bool   `yaml:"debug"`
 	DebugReportCaller   bool   `yaml:"debugReportCaller"`
@@ -26,6 +28,7 @@ type Logging struct {
 	LogsPath            string `yaml:"logsPath,omitempty"`
 }
 
+// Tracings is the struct that contains the configurations of the tracings
 type Tracings struct {
 	RabbitMQ `yaml:"rabbit-mq"`
 }
@@ -43,6 +46,7 @@ type Plugin struct {
 	OpenAPIChatGPTSecretKey string `yaml:"openAPIChatGPTSecretKey"`
 }
 
+// BeelzebubServiceConfiguration is the struct that contains the configurations of the honeypot service
 type BeelzebubServiceConfiguration struct {
 	ApiVersion             string    `yaml:"apiVersion"`
 	Protocol               string    `yaml:"protocol"`
@@ -57,6 +61,7 @@ type BeelzebubServiceConfiguration struct {
 	Plugin                 Plugin    `yaml:"plugin"`
 }
 
+// Command is the struct that contains the configurations of the commands
 type Command struct {
 	Regex      string   `yaml:"regex"`
 	Handler    string   `yaml:"handler"`
@@ -86,6 +91,7 @@ func Init(configurationsCorePath, configurationsServicesDirectory string) *confi
 	}
 }
 
+// ReadConfigurationsCore is the method that reads the configurations of the core from files
 func (bp configurationsParser) ReadConfigurationsCore() (*BeelzebubCoreConfigurations, error) {
 	buf, err := bp.readFileBytesByFilePathDependency(bp.configurationsCorePath)
 	if err != nil {
@@ -101,6 +107,7 @@ func (bp configurationsParser) ReadConfigurationsCore() (*BeelzebubCoreConfigura
 	return beelzebubConfiguration, nil
 }
 
+// ReadConfigurationsServices is the method that reads the configurations of the honeypot services from files
 func (bp configurationsParser) ReadConfigurationsServices() ([]BeelzebubServiceConfiguration, error) {
 	services, err := bp.gelAllFilesNameByDirNameDependency(bp.configurationsServicesDirectory)
 	if err != nil {
@@ -127,7 +134,7 @@ func (bp configurationsParser) ReadConfigurationsServices() ([]BeelzebubServiceC
 }
 
 func gelAllFilesNameByDirName(dirName string) ([]string, error) {
-	files, err := ioutil.ReadDir(dirName)
+	files, err := os.ReadDir(dirName)
 	if err != nil {
 		return nil, err
 	}

parser/configurations_parser_test.go

@@ -2,6 +2,7 @@ package parser
 
 import (
 	"errors"
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -118,3 +119,54 @@ func TestReadConfigurationsServicesValid(t *testing.T) {
 	assert.Equal(t, len(firstBeelzebubServiceConfiguration.Commands[0].Headers), 1)
 	assert.Equal(t, firstBeelzebubServiceConfiguration.Commands[0].Headers[0], "Content-Type: text/html")
 }
+
+func TestGelAllFilesNameByDirName(t *testing.T) {
+
+	var dir = t.TempDir()
+
+	files, err := gelAllFilesNameByDirName(dir)
+
+	assert.Nil(t, err)
+	assert.Equal(t, 0, len(files))
+}
+
+func TestGelAllFilesNameByDirNameFiles(t *testing.T) {
+
+	var dir = t.TempDir()
+
+	testFiles := []string{"file1.yaml", "file2.yaml", "file3.txt", "subdir", "file4.yaml"}
+	for _, filename := range testFiles {
+		filePath := dir + "/" + filename
+		file, err := os.Create(filePath)
+		assert.NoError(t, err)
+		file.Close()
+	}
+
+	files, err := gelAllFilesNameByDirName(dir)
+
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(files))
+}
+
+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())
+}
+
+func TestReadFileBytesByFilePath(t *testing.T) {
+
+	var dir = t.TempDir()
+	filePath := dir + "/test.yaml"
+
+	f, err := os.Create(filePath)
+	assert.NoError(t, err)
+	f.Close()
+
+	bytes, err := readFileBytesByFilePath(filePath)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "", string(bytes))
+}

plugins/openai-gpt.go

@@ -5,33 +5,27 @@ import (
 	"errors"
 	"fmt"
 	"strings"
-	
+
 	log "github.com/sirupsen/logrus"
-	
+
 	"github.com/go-resty/resty/v2"
 )
 
 const (
 	// Reference: https://www.engraved.blog/building-a-virtual-machine-inside/
 	promptVirtualizeLinuxTerminal = "I want you to act as a Linux terminal. I will type commands and you will reply with what the terminal should show. I want you to only reply with the terminal output inside one unique code block, and nothing else. Do no write explanations. Do not type commands unless I instruct you to do so.\n\nA:pwd\n\nQ:/home/user\n\n"
-	ChatGPTPluginName = "OpenAIGPTLinuxTerminal"
+	ChatGPTPluginName             = "OpenAIGPTLinuxTerminal"
-	openAIGPTEndpoint = "https://api.openai.com/v1/completions"
+	openAIGPTEndpoint             = "https://api.openai.com/v1/completions"
-) 
+)
 
 type History struct {
 	Input, Output string
 }
-	
-type OpenAIGPTVirtualTerminal struct {
-	Histories               []History
-	OpenAPIChatGPTSecretKey string
-	client                  *resty.Client
-}
 
-func (openAIGPTVirtualTerminal *OpenAIGPTVirtualTerminal) InjectDependency() {
+type openAIGPTVirtualTerminal struct {
-	if openAIGPTVirtualTerminal.client == nil {
+	Histories []History
-		openAIGPTVirtualTerminal.client = resty.New()
+	openAIKey string
-	}
+	client    *resty.Client
 }
 
 type Choice struct {
@@ -65,6 +59,14 @@ type gptRequest struct {
 	Stop             []string `json:"stop"`
 }
 
+func Init(history []History, openAIKey string) *openAIGPTVirtualTerminal {
+	return &openAIGPTVirtualTerminal{
+		Histories: history,
+		openAIKey: openAIKey,
+		client:    resty.New(),
+	}
+}
+
 func buildPrompt(histories []History, command string) string {
 	var sb strings.Builder
 
@@ -79,7 +81,7 @@ func buildPrompt(histories []History, command string) string {
 	return sb.String()
 }
 
-func (openAIGPTVirtualTerminal *OpenAIGPTVirtualTerminal) GetCompletions(command string) (string, error) {
+func (openAIGPTVirtualTerminal *openAIGPTVirtualTerminal) GetCompletions(command string) (string, error) {
 	requestJson, err := json.Marshal(gptRequest{
 		Model:            "text-davinci-003",
 		Prompt:           buildPrompt(openAIGPTVirtualTerminal.Histories, command),
@@ -94,14 +96,14 @@ func (openAIGPTVirtualTerminal *OpenAIGPTVirtualTerminal) GetCompletions(command
 		return "", err
 	}
 
-	if openAIGPTVirtualTerminal.OpenAPIChatGPTSecretKey == "" {
+	if openAIGPTVirtualTerminal.openAIKey == "" {
-		return "", errors.New("OpenAPIChatGPTSecretKey is empty")
+		return "", errors.New("openAIKey is empty")
 	}
 
 	response, err := openAIGPTVirtualTerminal.client.R().
 		SetHeader("Content-Type", "application/json").
 		SetBody(requestJson).
-		SetAuthToken(openAIGPTVirtualTerminal.OpenAPIChatGPTSecretKey).
+		SetAuthToken(openAIGPTVirtualTerminal.openAIKey).
 		SetResult(&gptResponse{}).
 		Post(openAIGPTEndpoint)
 

plugins/openai-gpt_test.go

@@ -46,7 +46,15 @@ func TestBuildPromptWithHistory(t *testing.T) {
 		prompt)
 }
 
-func TestBuildGetCompletions(t *testing.T) {
+func TestBuildGetCompletionsFailValidation(t *testing.T) {
+	openAIGPTVirtualTerminal := Init(make([]History, 0), "")
+
+	_, err := openAIGPTVirtualTerminal.GetCompletions("test")
+
+	assert.Equal(t, "openAIKey is empty", err.Error())
+}
+
+func TestBuildGetCompletionsWithResults(t *testing.T) {
 	client := resty.New()
 	httpmock.ActivateNonDefault(client.GetClient())
 	defer httpmock.DeactivateAndReset()
@@ -68,10 +76,8 @@ func TestBuildGetCompletions(t *testing.T) {
 		},
 	)
 
-	openAIGPTVirtualTerminal := OpenAIGPTVirtualTerminal{
+	openAIGPTVirtualTerminal := Init(make([]History, 0), "sdjdnklfjndslkjanfk")
-		OpenAPIChatGPTSecretKey: "sdjdnklfjndslkjanfk",
+	openAIGPTVirtualTerminal.client = client
-		client:                  client,
-	}
 
 	//When
 	str, err := openAIGPTVirtualTerminal.GetCompletions("ls")
@@ -80,3 +86,31 @@ func TestBuildGetCompletions(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, "prova.txt", str)
 }
+
+func TestBuildGetCompletionsWithoutResults(t *testing.T) {
+	client := resty.New()
+	httpmock.ActivateNonDefault(client.GetClient())
+	defer httpmock.DeactivateAndReset()
+
+	// Given
+	httpmock.RegisterResponder("POST", openAIGPTEndpoint,
+		func(req *http.Request) (*http.Response, error) {
+			resp, err := httpmock.NewJsonResponse(200, &gptResponse{
+				Choices: []Choice{},
+			})
+			if err != nil {
+				return httpmock.NewStringResponse(500, ""), nil
+			}
+			return resp, nil
+		},
+	)
+
+	openAIGPTVirtualTerminal := Init(make([]History, 0), "sdjdnklfjndslkjanfk")
+	openAIGPTVirtualTerminal.client = client
+
+	//When
+	_, err := openAIGPTVirtualTerminal.GetCompletions("ls")
+
+	//Then
+	assert.Equal(t, "no choices", err.Error())
+}

protocols/protocol_manager.go

@@ -1,10 +1,12 @@
+// Package protocols is responsible for managing the different protocols
 package protocols
 
 import (
-	"github.com/mariocandela/beelzebub/parser"
+	"github.com/mariocandela/beelzebub/v3/parser"
-	"github.com/mariocandela/beelzebub/tracer"
+	"github.com/mariocandela/beelzebub/v3/tracer"
 )
 
+// ServiceStrategy is the common interface that each protocol honeypot implements
 type ServiceStrategy interface {
 	Init(beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration, tracer tracer.Tracer) error
 }
@@ -14,10 +16,11 @@ type ProtocolManager struct {
 	tracer   tracer.Tracer
 }
 
-func InitProtocolManager(tracerStrategy tracer.Strategy, strategy ServiceStrategy) *ProtocolManager {
+// InitProtocolManager is the method that initializes the protocol manager, receving the concrete tracer and the concrete service
+func InitProtocolManager(tracerStrategy tracer.Strategy, serviceStrategy ServiceStrategy) *ProtocolManager {
 	return &ProtocolManager{
-		tracer:   tracer.Init(tracerStrategy),
+		tracer:   tracer.GetInstance(tracerStrategy),
-		strategy: strategy,
+		strategy: serviceStrategy,
 	}
 }
 
@@ -25,6 +28,7 @@ func (pm *ProtocolManager) SetProtocolStrategy(strategy ServiceStrategy) {
 	pm.strategy = strategy
 }
 
+// InitService is the method that initializes the honeypot
 func (pm *ProtocolManager) InitService(beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration) error {
 	return pm.strategy.Init(beelzebubServiceConfiguration, pm.tracer)
 }

protocols/protocol_manager_test.go

@@ -2,8 +2,8 @@ package protocols
 
 import (
 	"errors"
-	"github.com/mariocandela/beelzebub/parser"
+	"github.com/mariocandela/beelzebub/v3/parser"
-	"github.com/mariocandela/beelzebub/tracer"
+	"github.com/mariocandela/beelzebub/v3/tracer"
 	"github.com/stretchr/testify/assert"
 	"testing"
 )

protocols/strategies/http.go

@@ -2,8 +2,8 @@ package strategies
 
 import (
 	"fmt"
-	"github.com/mariocandela/beelzebub/parser"
+	"github.com/mariocandela/beelzebub/v3/parser"
-	"github.com/mariocandela/beelzebub/tracer"
+	"github.com/mariocandela/beelzebub/v3/tracer"
 	"io"
 	"net/http"
 	"regexp"

protocols/strategies/ssh.go

@@ -2,10 +2,11 @@ package strategies
 
 import (
 	"fmt"
-	"github.com/mariocandela/beelzebub/parser"
+	"github.com/mariocandela/beelzebub/v3/parser"
-	"github.com/mariocandela/beelzebub/plugins"
+	"github.com/mariocandela/beelzebub/v3/plugins"
-	"github.com/mariocandela/beelzebub/tracer"
+	"github.com/mariocandela/beelzebub/v3/tracer"
 	"regexp"
+
 	"strings"
 	"time"
 
@@ -62,8 +63,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 							commandOutput := command.Handler
 
 							if command.Plugin == plugins.ChatGPTPluginName {
-								openAIGPTVirtualTerminal := plugins.OpenAIGPTVirtualTerminal{Histories: histories, OpenAPIChatGPTSecretKey: beelzebubServiceConfiguration.Plugin.OpenAPIChatGPTSecretKey}
+								openAIGPTVirtualTerminal := plugins.Init(histories, beelzebubServiceConfiguration.Plugin.OpenAPIChatGPTSecretKey)
-								openAIGPTVirtualTerminal.InjectDependency()
 
 								if commandOutput, err = openAIGPTVirtualTerminal.GetCompletions(commandInput); err != nil {
 									log.Errorf("Error GetCompletions: %s, %s", commandInput, err.Error())
@@ -124,7 +124,7 @@ func (sshStrategy *SSHStrategy) Init(beelzebubServiceConfiguration parser.Beelze
 	log.WithFields(log.Fields{
 		"port":     beelzebubServiceConfiguration.Address,
 		"commands": len(beelzebubServiceConfiguration.Commands),
-	}).Infof("Init service %s", beelzebubServiceConfiguration.Protocol)
+	}).Infof("GetInstance service %s", beelzebubServiceConfiguration.Protocol)
 	return nil
 }
 

protocols/strategies/tcp.go

@@ -2,8 +2,8 @@ package strategies
 
 import (
 	"fmt"
-	"github.com/mariocandela/beelzebub/parser"
+	"github.com/mariocandela/beelzebub/v3/parser"
-	"github.com/mariocandela/beelzebub/tracer"
+	"github.com/mariocandela/beelzebub/v3/tracer"
 	"net"
 	"time"
 

tracer/tracer.go

@@ -1,13 +1,16 @@
+// Package tracer is responsible for tracing the events that occur in the honeypots
 package tracer
 
 import (
+	log "github.com/sirupsen/logrus"
+	"sync"
 	"time"
 
 	"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
 const Workers = 5
 
 type Event struct {
@@ -44,8 +47,8 @@ const (
 	TCP
 )
 
-func (status Protocol) String() string {
+func (protocol Protocol) String() string {
-	return [...]string{"HTTP", "SSH", "TCP"}[status]
+	return [...]string{"HTTP", "SSH", "TCP"}[protocol]
 }
 
 const (
@@ -66,49 +69,60 @@ type Tracer interface {
 }
 
 type tracer struct {
-	strategy   Strategy
+	strategy        Strategy
-	eventsChan chan Event
+	eventsChan      chan Event
+	eventsTotal     prometheus.Counter
+	eventsSSHTotal  prometheus.Counter
+	eventsTCPTotal  prometheus.Counter
+	eventsHTTPTotal prometheus.Counter
 }
 
-var (
+var lock = &sync.Mutex{}
-	eventsTotal = promauto.NewCounter(prometheus.CounterOpts{
+var singleton *tracer
-		Namespace: "beelzebub",
-		Name:      "events_total",
-		Help:      "The total number of events",
-	})
-	eventsSSHTotal = promauto.NewCounter(prometheus.CounterOpts{
-		Namespace: "beelzebub",
-		Name:      "ssh_events_total",
-		Help:      "The total number of SSH events",
-	})
-	eventsTCPTotal = promauto.NewCounter(prometheus.CounterOpts{
-		Namespace: "beelzebub",
-		Name:      "tcp_events_total",
-		Help:      "The total number of TCP events",
-	})
-	eventsHTTPTotal = promauto.NewCounter(prometheus.CounterOpts{
-		Namespace: "beelzebub",
-		Name:      "http_events_total",
-		Help:      "The total number of HTTP events",
-	})
-)
 
-func Init(strategy Strategy) *tracer {
+func GetInstance(defaultStrategy Strategy) *tracer {
-	tracer := &tracer{
+	if singleton == nil {
-		strategy:   strategy,
+		lock.Lock()
-		eventsChan: make(chan Event, Workers),
+		defer lock.Unlock()
-	}
+		// This is to prevent expensive lock operations every time the GetInstance method is called
-
+		if singleton == nil {
-	for i := 0; i < Workers; i++ {
+			singleton = &tracer{
-		go func(i int) {
+				strategy:   defaultStrategy,
-			log.Debug("Init trace worker: ", i)
+				eventsChan: make(chan Event, Workers),
-			for event := range tracer.eventsChan {
+				eventsTotal: promauto.NewCounter(prometheus.CounterOpts{
-				tracer.strategy(event)
+					Namespace: "beelzebub",
+					Name:      "events_total",
+					Help:      "The total number of events",
+				}),
+				eventsSSHTotal: promauto.NewCounter(prometheus.CounterOpts{
+					Namespace: "beelzebub",
+					Name:      "ssh_events_total",
+					Help:      "The total number of SSH events",
+				}),
+				eventsTCPTotal: promauto.NewCounter(prometheus.CounterOpts{
+					Namespace: "beelzebub",
+					Name:      "tcp_events_total",
+					Help:      "The total number of TCP events",
+				}),
+				eventsHTTPTotal: promauto.NewCounter(prometheus.CounterOpts{
+					Namespace: "beelzebub",
+					Name:      "http_events_total",
+					Help:      "The total number of HTTP events",
+				}),
 			}
-		}(i)
+
+			for i := 0; i < Workers; i++ {
+				go func(i int) {
+					log.Debug("Trace worker: ", i)
+					for event := range singleton.eventsChan {
+						singleton.strategy(event)
+					}
+				}(i)
+			}
+		}
 	}
 
-	return tracer
+	return singleton
 }
 
 func (tracer *tracer) setStrategy(strategy Strategy) {
@@ -120,14 +134,17 @@ func (tracer *tracer) TraceEvent(event Event) {
 
 	tracer.eventsChan <- event
 
-	eventsTotal.Inc()
+	tracer.updatePrometheusCounters(event.Protocol)
-
+}
-	switch event.Protocol {
+
-	case HTTP.String():
+func (tracer *tracer) updatePrometheusCounters(protocol string) {
-		eventsHTTPTotal.Inc()
+	switch protocol {
-	case SSH.String():
+	case HTTP.String():
-		eventsSSHTotal.Inc()
+		tracer.eventsHTTPTotal.Inc()
-	case TCP.String():
+	case SSH.String():
-		eventsTCPTotal.Inc()
+		tracer.eventsSSHTotal.Inc()
-	}
+	case TCP.String():
+		tracer.eventsTCPTotal.Inc()
+	}
+	tracer.eventsTotal.Inc()
 }

tracer/tracer_test.go

@@ -1,6 +1,7 @@
 package tracer
 
 import (
+	"github.com/prometheus/client_golang/prometheus"
 	"sync"
 	"testing"
 
@@ -10,7 +11,7 @@ import (
 func TestInit(t *testing.T) {
 	mockStrategy := func(event Event) {}
 
-	tracer := Init(mockStrategy)
+	tracer := GetInstance(mockStrategy)
 
 	assert.NotNil(t, tracer.strategy)
 }
@@ -25,7 +26,9 @@ func TestTraceEvent(t *testing.T) {
 		eventCalled = event
 	}
 
-	tracer := Init(mockStrategy)
+	tracer := GetInstance(mockStrategy)
+
+	tracer.strategy = mockStrategy
 
 	wg.Add(1)
 	tracer.TraceEvent(Event{
@@ -51,7 +54,7 @@ func TestSetStrategy(t *testing.T) {
 		eventCalled = event
 	}
 
-	tracer := Init(mockStrategy)
+	tracer := GetInstance(mockStrategy)
 
 	tracer.setStrategy(mockStrategy)
 
@@ -75,3 +78,42 @@ func TestStringStatus(t *testing.T) {
 	assert.Equal(t, Stateless.String(), "Stateless")
 	assert.Equal(t, Interaction.String(), "Interaction")
 }
+
+type mockCounter struct {
+	prometheus.Metric
+	prometheus.Collector
+	inc func()
+	add func(float64)
+}
+
+var counter = 0
+
+func (m mockCounter) Inc() {
+	counter += 1
+}
+
+func (m mockCounter) Add(f float64) {
+	counter = int(f)
+}
+
+func TestUpdatePrometheusCounters(t *testing.T) {
+	mockStrategy := func(event Event) {}
+
+	tracer := &tracer{
+		strategy:        mockStrategy,
+		eventsChan:      make(chan Event, Workers),
+		eventsTotal:     mockCounter{},
+		eventsSSHTotal:  mockCounter{},
+		eventsTCPTotal:  mockCounter{},
+		eventsHTTPTotal: mockCounter{},
+	}
+
+	tracer.updatePrometheusCounters(SSH.String())
+	assert.Equal(t, 2, counter)
+
+	tracer.updatePrometheusCounters(HTTP.String())
+	assert.Equal(t, 4, counter)
+
+	tracer.updatePrometheusCounters(TCP.String())
+	assert.Equal(t, 6, counter)
+}