Merge pull request #2 from mariocandela/ImproveTracerHTTPStrategy

Improve tracer http strategy

.gitignore

@@ -1,2 +1,3 @@
 .DS_Store
 .idea
+logs

protocols/hypertextTransferProtocolStrategy.go

@@ -4,6 +4,7 @@ import (
 	"beelzebub/parser"
 	"beelzebub/tracer"
 	"fmt"
+	"github.com/google/uuid"
 	log "github.com/sirupsen/logrus"
 	"io"
 	"net/http"
@@ -12,38 +13,15 @@ import (
 )
 
 type HypertextTransferProtocolStrategy struct {
-	serverMux                     *http.ServeMux
 	beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration
 }
 
 func (httpStrategy HypertextTransferProtocolStrategy) Init(beelzebubServiceConfiguration parser.BeelzebubServiceConfiguration, tr tracer.Tracer) error {
 	httpStrategy.beelzebubServiceConfiguration = beelzebubServiceConfiguration
-	httpStrategy.serverMux = http.NewServeMux()
+	serverMux := http.NewServeMux()
 
-	httpStrategy.buildHandler()
+	serverMux.HandleFunc("/", func(responseWriter http.ResponseWriter, request *http.Request) {
-
+		traceRequest(request, tr)
-	go func() {
-		httpStrategy.listenAndServe()
-	}()
-
-	log.WithFields(log.Fields{
-		"port":     beelzebubServiceConfiguration.Address,
-		"commands": len(beelzebubServiceConfiguration.Commands),
-	}).Infof("Init service %s", beelzebubServiceConfiguration.Protocol)
-	return nil
-}
-
-func (httpStrategy HypertextTransferProtocolStrategy) listenAndServe() {
-	err := http.ListenAndServe(httpStrategy.beelzebubServiceConfiguration.Address, httpStrategy.serverMux)
-	if err != nil {
-		log.Errorf("Error during init HTTP Protocol: %s", err.Error())
-		return
-	}
-}
-
-func (httpStrategy HypertextTransferProtocolStrategy) buildHandler() {
-	httpStrategy.serverMux.HandleFunc("/", func(responseWriter http.ResponseWriter, request *http.Request) {
-		traceRequest(request)
 		for _, command := range httpStrategy.beelzebubServiceConfiguration.Commands {
 			matched, err := regexp.MatchString(command.Regex, request.RequestURI)
 			if err != nil {
@@ -58,30 +36,48 @@ func (httpStrategy HypertextTransferProtocolStrategy) buildHandler() {
 			}
 		}
 	})
+	go func() {
+		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.Protocol)
+	return nil
 }
 
-func traceRequest(request *http.Request) {
+func traceRequest(request *http.Request, tr tracer.Tracer) {
 	bodyBytes, err := io.ReadAll(request.Body)
 	body := ""
 	if err == nil {
 		body = string(bodyBytes)
 	}
-	log.WithFields(log.Fields{
+	tr.TraceEvent(tracer.Event{
-		"requestURI": request.RequestURI,
+		Msg:             "HTTP New request",
-		"proto":      request.Proto,
+		RequestURI:      request.RequestURI,
-		"method":     request.Method,
+		Protocol:        tracer.HTTP,
-		"body":       body,
+		HTTPMethod:      request.Method,
-		"host":       request.Host,
+		Body:            body,
-		"userAgent":  request.UserAgent(),
+		HostHTTPRequest: request.Host,
-		"cookies":    request.Cookies(),
+		UserAgent:       request.UserAgent(),
-		"ip":         request.RemoteAddr,
+		Cookies:         request.Cookies(),
-		"headers":    request.Header,
+		Headers:         request.Header,
-		"remoteAddr": request.RemoteAddr,
+		Status:          tracer.Stateless,
-	}).Info("New HTTP request")
+		RemoteAddr:      request.RemoteAddr,
+		ID:              uuid.New().String(),
+	})
 }
 
 func setResponseHeaders(responseWriter http.ResponseWriter, headers []string, statusCode int) {
+	// http.StatusText(statusCode): empty string if the code is unknown.
+	if len(http.StatusText(statusCode)) > 0 {
 		responseWriter.WriteHeader(statusCode)
+	}
 	for _, headerStr := range headers {
 		keyValue := strings.Split(headerStr, ":")
 		if len(keyValue) > 1 {

protocols/secureShellStrategy.go

@@ -28,7 +28,7 @@ func (SSHStrategy *SecureShellStrategy) Init(beelzebubServiceConfiguration parse
 
 				tr.TraceEvent(tracer.Event{
 					Msg:        "New SSH Session",
-					Protocol:   beelzebubServiceConfiguration.Protocol,
+					Protocol:   tracer.SSH,
 					RemoteAddr: sess.RemoteAddr().String(),
 					Status:     tracer.Start,
 					ID:         uuidSession.String(),

tracer/tracer.go

@@ -2,6 +2,7 @@ package tracer
 
 import (
 	log "github.com/sirupsen/logrus"
+	"net/http"
 )
 
 type Tracer struct {
@@ -20,7 +21,7 @@ func (tracer *Tracer) TraceEvent(event Event) {
 
 type Event struct {
 	RemoteAddr      string
-	Protocol   string
+	Protocol        Protocol
 	Command         string
 	Status          Status
 	Msg             string
@@ -29,8 +30,22 @@ type Event struct {
 	User            string
 	Password        string
 	Client          string
+	Headers         http.Header
+	Cookies         []*http.Cookie
+	UserAgent       string
+	HostHTTPRequest string
+	Body            string
+	HTTPMethod      string
+	RequestURI      string
 }
 
+type Protocol int
+
+const (
+	HTTP Protocol = iota
+	SSH
+)
+
 type Status int
 
 const (