refactor]

protocols/helpers/helpers.go

@@ -1,9 +1,38 @@
 package helpers
 
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"os"
+	"path/filepath"
+)
+
 func FirstOrEmpty[T any](s []T) T {
 	if len(s) > 0 {
 		return s[0]
 	}
 	var t T
 	return t
 }
+
+func StorePayload(data []byte) (string, error) {
+	sum := sha256.Sum256(data)
+	if err := os.MkdirAll("payloads", os.ModePerm); err != nil {
+		return "", err
+	}
+	sha256Hash := hex.EncodeToString(sum[:])
+	path := filepath.Join("payloads", sha256Hash)
+	if _, err := os.Stat(path); err == nil {
+		// file already exists
+		return "", nil
+	}
+	out, err := os.Create(path)
+	if err != nil {
+		return "", err
+	}
+	defer out.Close()
+	if _, err = out.Write(data); err != nil {
+		return "", err
+	}
+	return sha256Hash, nil
+}

protocols/tcp/tcp.go

@@ -3,14 +3,11 @@ package tcp
 import (
 	"context"
 	"crypto/rand"
-	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
 	"log/slog"
 	"math/big"
 	"net"
-	"os"
-	"path/filepath"
 	"strconv"
 
 	"github.com/mushorg/glutton/connection"
@@ -30,28 +27,6 @@ type tcpServer struct {
 	events []parsedTCP
 }
 
-func storePayload(data []byte) (string, error) {
-	sum := sha256.Sum256(data)
-	if err := os.MkdirAll("payloads", os.ModePerm); err != nil {
-		return "", err
-	}
-	sha256Hash := hex.EncodeToString(sum[:])
-	path := filepath.Join("payloads", sha256Hash)
-	if _, err := os.Stat(path); err == nil {
-		return "", nil
-	}
-	out, err := os.Create(path)
-	if err != nil {
-		return "", err
-	}
-	defer out.Close()
-	_, err = out.Write(data)
-	if err != nil {
-		return "", err
-	}
-	return sha256Hash, nil
-}
-
 func (s *tcpServer) sendRandom(conn net.Conn) error {
 	randomInt, err := rand.Int(rand.Reader, big.NewInt(500))
 	if err != nil {
@@ -89,7 +64,7 @@ func HandleTCP(ctx context.Context, conn net.Conn, md connection.Metadata, logge
 
 	defer func() {
 		if msgLength > 0 {
-			payloadHash, err := storePayload(data)
+			payloadHash, err := helpers.StorePayload(data)
 			if err != nil {
 				logger.Error("failed to store payload", slog.String("handler", "tcp"), producer.ErrAttr(err))
 			}

protocols/udp/udp.go

@@ -2,44 +2,19 @@ package udp
 
 import (
 	"context"
-	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
 	"net"
-	"os"
-	"path/filepath"
 
 	"github.com/mushorg/glutton/connection"
 	"github.com/mushorg/glutton/producer"
+	"github.com/mushorg/glutton/protocols/helpers"
 	"github.com/mushorg/glutton/protocols/interfaces"
 )
 
-func storePayload(data []byte) (string, error) {
-	sum := sha256.Sum256(data)
-	if err := os.MkdirAll("payloads", os.ModePerm); err != nil {
-		return "", err
-	}
-	sha256Hash := hex.EncodeToString(sum[:])
-	path := filepath.Join("payloads", sha256Hash)
-	if _, err := os.Stat(path); err == nil {
-		return "", nil
-	}
-	out, err := os.Create(path)
-	if err != nil {
-		return "", err
-	}
-	defer out.Close()
-	_, err = out.Write(data)
-	if err != nil {
-		return "", err
-	}
-	return sha256Hash, nil
-}
-
 func HandleUDP(ctx context.Context, srcAddr, dstAddr *net.UDPAddr, data []byte, md connection.Metadata, log interfaces.Logger, h interfaces.Honeypot) error {
 	log.Info(fmt.Sprintf("UDP payload:\n%s", hex.Dump(data[:len(data)%1024])))
-	_, err := storePayload(data[:len(data)%1024])
-	if err != nil {
+	if _, err := helpers.StorePayload(data[:len(data)%1024]); err != nil {
 		log.Error("failed to store UDP payload", producer.ErrAttr(err))
 	}
 	if err := h.ProduceUDP("udp", srcAddr, dstAddr, md, data[:len(data)%1024], nil); err != nil {