Feature enable protection for direct tcp

Dockerfile

@@ -1,4 +1,4 @@
-FROM golang:1.16-alpine
+FROM golang:1.22-alpine
 RUN apk add build-base
 RUN apk add libpcap-dev
 RUN apk add tcpdump
@@ -10,6 +10,7 @@ COPY go.sum ./
 RUN go mod download
 
 COPY *.go ./
+COPY protobuf ./protobuf
 COPY db ./db
 COPY utils ./utils
 

generate.go

@@ -0,0 +1,3 @@
+package main
+
+//go:generate protoc --go_out=. --go-grpc_out=. protobuf/traffic_payload/message.proto

go.mod

@@ -1,12 +1,15 @@
 module github.com/akto-api-security/mirroring-api-logging
 
-go 1.17
+go 1.22
+
+toolchain go1.23.3
 
 require (
 	github.com/akto-api-security/gomiddleware v0.1.0
 	github.com/google/gopacket v1.1.19
 	github.com/segmentio/kafka-go v0.4.25
 	go.mongodb.org/mongo-driver v1.11.3
+	google.golang.org/protobuf v1.36.1
 )
 
 require (

go.sum

@@ -10,8 +10,9 @@ github.com/frankban/quicktest v1.11.3/go.mod h1:wRf/ReqHper53s+kmmSZizM8NamnL3IM
 github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4=
 github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
-github.com/google/go-cmp v0.5.4 h1:L8R9j+yAqZuZjsqh/z+F1NCffTKKLShY6zXTItVIZ8M=
 github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
+github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/gopacket v1.1.19 h1:ves8RnFZPGiFnTS0uPQStjwru6uO6h+nlr9j6fL7kF8=
 github.com/google/gopacket v1.1.19/go.mod h1:iJ8V8n6KS+z2U1A8pUwu8bW5SyEMkXJB8Yo/Vo+TKTo=
 github.com/klauspost/compress v1.9.8/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
@@ -31,7 +32,6 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/segmentio/kafka-go v0.4.23/go.mod h1:XzMcoMjSzDGHcIwpWUI7GB43iKZ2fTVmryPSGLf/MPg=
 github.com/segmentio/kafka-go v0.4.25 h1:QVx9yz12syKBFkxR+dVDDwTO0ItHgnjjhIdBfqizj+8=
 github.com/segmentio/kafka-go v0.4.25/go.mod h1:XzMcoMjSzDGHcIwpWUI7GB43iKZ2fTVmryPSGLf/MPg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -83,6 +83,8 @@ golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapK
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+google.golang.org/protobuf v1.36.1 h1:yBPeRvTftaleIgM3PZ/WBIZ7XM/eEYAaEyCwvyjq/gk=
+google.golang.org/protobuf v1.36.1/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

kafka.go

@@ -0,0 +1,97 @@
+package main
+
+import (
+	"context"
+	trafficpb "github.com/akto-api-security/mirroring-api-logging/protobuf/traffic_payload"
+	"github.com/segmentio/kafka-go"
+	"google.golang.org/protobuf/proto"
+	"log"
+	"log/slog"
+	"strings"
+	"time"
+)
+
+var CLIENT_IP_HEADERS = []string{
+	"x-forwarded-for",
+	"x-real-ip",
+	"x-cluster-client-ip",
+	"true-client-ip",
+	"x-original-forwarded-for",
+	"x-client-ip",
+	"client-ip",
+}
+
+func Produce(kafkaWriter *kafka.Writer, ctx context.Context, value *trafficpb.HttpResponseParam) error {
+	// intialize the writer with the broker addresses, and the topic
+	protoBytes, err := proto.Marshal(value)
+	if err != nil {
+		slog.Error("Failed to serialize protobuf message", "error", err)
+		return err
+	}
+
+	if value.Ip == "" {
+		slog.Warn("ip is empty, avoiding kafka push")
+		return nil
+	}
+	// Send serialized message to Kafka
+	topic := "akto.api.logs2"
+	msg := kafka.Message{
+		Topic: topic,
+		Key:   []byte(value.Ip), // what to do when ip is empty?
+		Value: protoBytes,
+	}
+
+	err = kafkaWriter.WriteMessages(ctx, msg)
+
+	if err != nil {
+		slog.Error("Kafka write for threat failed", "topic", topic, "error", err)
+	}
+	return err
+}
+
+func GetSourceIp(reqHeaders map[string]*trafficpb.StringList, packetIp string) string {
+
+	for _, header := range CLIENT_IP_HEADERS {
+		if headerValues, exists := reqHeaders[header]; exists {
+			for _, headerValue := range headerValues.Values {
+				parts := strings.Split(headerValue, ",")
+				for _, part := range parts {
+					ip := strings.TrimSpace(part)
+					if ip != "" {
+						slog.Debug("Ip found in", "the header", header)
+						return ip
+					}
+				}
+			}
+		}
+	}
+
+	slog.Debug("No ip found in headers returning", "packetIp", packetIp)
+	return packetIp
+}
+
+func ProduceStr(kafkaWriter *kafka.Writer, ctx context.Context, message string) error {
+	// intialize the writer with the broker addresses, and the topic
+	topic := "akto.api.logs"
+	msg := kafka.Message{
+		Topic: topic,
+		Value: []byte(message),
+	}
+	err := kafkaWriter.WriteMessages(ctx, msg)
+
+	if err != nil {
+		log.Println("ERROR while writing messages: ", err)
+		slog.Error("Kafka write for runtime failed", "topic", topic, "error", err)
+		return err
+	}
+
+	return nil
+}
+
+func GetKafkaWriter(kafkaURL string, batchSize int, batchTimeout time.Duration) *kafka.Writer {
+	return &kafka.Writer{
+		Addr:         kafka.TCP(kafkaURL),
+		BatchSize:    batchSize,
+		BatchTimeout: batchTimeout,
+	}
+}

main.go

@@ -17,8 +17,8 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"github.com/akto-api-security/gomiddleware"
 	"github.com/akto-api-security/mirroring-api-logging/db"
+	trafficpb "github.com/akto-api-security/mirroring-api-logging/protobuf/traffic_payload"
 	"github.com/akto-api-security/mirroring-api-logging/utils"
 	"github.com/segmentio/kafka-go"
 	"io"
@@ -42,6 +42,7 @@ var printCounter = 500
 var assemblerMap = make(map[int]*tcpassembly.Assembler)
 var incomingCountMap = make(map[string]utils.IncomingCounter)
 var outgoingCountMap = make(map[string]utils.OutgoingCounter)
+var threatEnabled = false
 
 var readFiles = make(map[string]bool)
 
@@ -255,26 +256,71 @@ func tryReadFromBD(bd *bidi, isPending bool) {
 
 		}
 
-		reqHeader := make(map[string]string)
+		reqHeader := make(map[string]*trafficpb.StringList)
 		for name, values := range req.Header {
 			// Loop over all values for the name.
 			for _, value := range values {
-				reqHeader[name] = value
+				reqHeader[strings.ToLower(name)] = &trafficpb.StringList{
+					Values: []string{value},
+				}
 			}
 		}
+		ip := GetSourceIp(reqHeader, bd.key.net.Src().String())
 
-		reqHeader["host"] = req.Host
+		reqHeader["host"] = &trafficpb.StringList{
+			Values: []string{req.Host},
+		}
 
-		respHeader := make(map[string]string)
+		// build req headers for runtime
+		reqHeaderStr := make(map[string]string)
+		for name, values := range req.Header {
+			// Loop over all values for the name.
+			for _, value := range values {
+				reqHeaderStr[name] = value
+			}
+		}
+		reqHeaderStr["host"] = req.Host
+
+		// build resp headers for threat client
+		respHeader := make(map[string]*trafficpb.StringList)
 		for name, values := range resp.Header {
 			// Loop over all values for the name.
 			for _, value := range values {
-				respHeader[name] = value
+				respHeader[strings.ToLower(name)] = &trafficpb.StringList{
+					Values: []string{value},
+				}
 			}
 		}
 
-		reqHeaderString, _ := json.Marshal(reqHeader)
-		respHeaderString, _ := json.Marshal(respHeader)
+		// build resp headers for runtime
+		respHeaderStr := make(map[string]string)
+		for name, values := range resp.Header {
+			// Loop over all values for the name.
+			for _, value := range values {
+				respHeaderStr[name] = value
+			}
+		}
+
+		payload := &trafficpb.HttpResponseParam{
+			Method:          req.Method,
+			Path:            req.URL.String(),
+			RequestHeaders:  reqHeader,
+			ResponseHeaders: respHeader,
+			RequestPayload:  requestsContent[i],
+			ResponsePayload: string(body),
+			Ip:              ip,
+			Time:            int32(time.Now().Unix()),
+			StatusCode:      int32(resp.StatusCode),
+			Type:            string(req.Proto),
+			Status:          resp.Status,
+			AktoAccountId:   fmt.Sprint(1000000),
+			AktoVxlanId:     fmt.Sprint(bd.vxlanID),
+			IsPending:       isPending,
+		}
+		ctx := context.Background()
+
+		reqHeaderString, _ := json.Marshal(reqHeaderStr)
+		respHeaderString, _ := json.Marshal(respHeaderStr)
 
 		if shouldPrintDebugURL(req.URL.String()) {
 			fmt.Printf("Found debug url %s while creating final value", req.URL.String())
@@ -299,13 +345,12 @@ func tryReadFromBD(bd *bidi, isPending bool) {
 		}
 
 		out, _ := json.Marshal(value)
-		ctx := context.Background()
 
 		totalCounter += 1
 
 		// calculating the size of outgoing bytes and requests (1) and saving it in outgoingCounterMap
 		outgoingBytes := len(bd.a.bytes) + len(bd.b.bytes)
-		hostString := reqHeader["host"]
+		hostString := reqHeader["host"].String()
 		if utils.CheckIfIpHost(hostString) {
 			hostString = "ip-host"
 		}
@@ -322,7 +367,11 @@ func tryReadFromBD(bd *bidi, isPending bool) {
 			printCounter--
 			log.Println("req-resp.String()", string(out))
 		}
-		go gomiddleware.Produce(kafkaWriter, ctx, string(out))
+		go ProduceStr(kafkaWriter, ctx, string(out))
+		if threatEnabled {
+			// insert kafka record for threat client
+			go Produce(kafkaWriter, ctx, payload)
+		}
 		if totalCounter%1000 == 0 {
 			println("totalCounter: ", totalCounter)
 		}
@@ -401,7 +450,15 @@ func run(handle *pcap.Handle, apiCollectionId int, source string) {
 	}
 	kafka_batch_time_secs_duration := time.Duration(kafka_batch_time_secs)
 
-	kafkaWriter = gomiddleware.GetKafkaWriter(kafka_url, "akto.api.logs", kafka_batch_size, kafka_batch_time_secs_duration*time.Second)
+	threatEnabledVar := os.Getenv("AKTO_THREAT_ENABLED")
+	if len(threatEnabledVar) > 0 {
+		threatEnabled = strings.ToLower(threatEnabledVar) == "true"
+		log.Println("threatEnabled: ", threatEnabled)
+	} else {
+		log.Println("threatEnabled: missing. defaulting to false")
+	}
+
+	kafkaWriter = GetKafkaWriter(kafka_url, kafka_batch_size, kafka_batch_time_secs_duration*time.Second)
 
 	// Set up pcap packet capture
 	// handle, err = pcap.OpenOffline("/Users/ankushjain/Downloads/dump2.pcap")