make rawnet jobs use packetgen

arriven · Apr 4, 2022 · 2c15d5f · 2c15d5f
1 parent d034065
commit 2c15d5f
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 138 deletions.
diff --git a/examples/config/advanced/packetgen-tcp.yaml b/examples/config/advanced/packetgen-tcp.yaml
@@ -7,7 +7,7 @@ jobs:
           protocol: "tcp"
           address: "localhost:1234"
           tls_config:
-            insecure_skip_verify: true # this is not snake case because I decode it straight into tls config
+            insecure_skip_verify: true
       packet:
         payload:
           type: raw

diff --git a/src/job/rawnet.go b/src/job/rawnet.go
@@ -25,157 +25,40 @@ package job
 import (
 	"context"
 	"fmt"
-	"log"
-	"net"
-	"strings"
-	"text/template"
 	"time"
 
-	"github.com/google/uuid"
 	"go.uber.org/zap"
 
-	"github.com/Arriven/db1000n/src/core/packetgen"
 	"github.com/Arriven/db1000n/src/job/config"
-	"github.com/Arriven/db1000n/src/utils"
-	"github.com/Arriven/db1000n/src/utils/metrics"
-	"github.com/Arriven/db1000n/src/utils/templates"
 )
 
-type rawnetConfig struct {
-	BasicJobConfig
-	addr      string
-	bodyTpl   *template.Template
-	proxyURLs string
-	timeout   time.Duration
-}
-
 func tcpJob(ctx context.Context, logger *zap.Logger, globalConfig *GlobalConfig, args config.Args) (data any, err error) {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	jobConfig, err := parseRawNetJobArgs(ctx, logger, globalConfig, args)
+	packetgenArgs, err := parseRawNetJobArgs(ctx, logger, globalConfig, args, "tcp")
 	if err != nil {
 		return nil, err
 	}
 
-	backoffController := utils.BackoffController{BackoffConfig: utils.NonNilOrDefault(jobConfig.Backoff, globalConfig.Backoff)}
-
-	if globalConfig.ProxyURLs != "" {
-		jobConfig.proxyURLs = templates.ParseAndExecute(logger, globalConfig.ProxyURLs, ctx)
-	}
-
-	trafficMonitor := metrics.Default.NewWriter(metrics.Traffic, uuid.New().String())
-	go trafficMonitor.Update(ctx, time.Second)
-
-	processedTrafficMonitor := metrics.Default.NewWriter(metrics.ProcessedTraffic, uuid.NewString())
-	go processedTrafficMonitor.Update(ctx, time.Second)
-
-	if !isInEncryptedContext(ctx) {
-		log.Printf("Attacking %v", jobConfig.addr)
-	}
-
-	for jobConfig.Next(ctx) {
-		err = sendTCP(ctx, logger, jobConfig, trafficMonitor, processedTrafficMonitor)
-		if err != nil {
-			utils.Sleep(ctx, backoffController.Increment().GetTimeout())
-		} else {
-			backoffController.Reset()
-		}
-	}
-
-	return nil, nil
-}
-
-func sendTCP(ctx context.Context, logger *zap.Logger, jobConfig *rawnetConfig, trafficMonitor, processedTrafficMonitor *metrics.Writer) error {
-	// track sending of SYN packet
-	trafficMonitor.Add(packetgen.TCPHeaderSize + packetgen.IPHeaderSize)
-
-	conn, err := utils.GetProxyFunc(jobConfig.proxyURLs, jobConfig.timeout)("tcp", jobConfig.addr)
-	if err != nil {
-		logger.Debug("error connecting via tcp", zap.String("addr", jobConfig.addr), zap.Error(err))
-		metrics.IncRawnetTCP(jobConfig.addr, metrics.StatusFail)
-
-		return err
-	}
-
-	defer conn.Close()
-
-	// if we got here the connection was successful and thus we need to track SYN
-	processedTrafficMonitor.Add(packetgen.TCPHeaderSize + packetgen.IPHeaderSize)
-
-	// if we got here we had to also send ACK packet (and get a response)
-	trafficMonitor.Add(packetgen.TCPHeaderSize + packetgen.IPHeaderSize)
-	processedTrafficMonitor.Add(packetgen.TCPHeaderSize + packetgen.IPHeaderSize)
-
-	// Write to conn until error
-	for jobConfig.Next(ctx) {
-		n, err := conn.Write([]byte(templates.Execute(logger, jobConfig.bodyTpl, ctx)))
-		trafficMonitor.Add(uint64(n) + packetgen.TCPHeaderSize + packetgen.IPHeaderSize)
-
-		if err != nil {
-			metrics.IncRawnetTCP(jobConfig.addr, metrics.StatusFail)
-
-			return err
-		}
-
-		processedTrafficMonitor.Add(uint64(n))
-		metrics.IncRawnetTCP(jobConfig.addr, metrics.StatusSuccess)
-	}
-
-	return nil
+	return packetgenJob(ctx, logger, globalConfig, packetgenArgs)
 }
 
 func udpJob(ctx context.Context, logger *zap.Logger, globalConfig *GlobalConfig, args config.Args) (data any, err error) {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	jobConfig, err := parseRawNetJobArgs(ctx, logger, globalConfig, args)
+	packetgenArgs, err := parseRawNetJobArgs(ctx, logger, globalConfig, args, "udp")
 	if err != nil {
 		return nil, err
 	}
 
-	udpAddr, err := net.ResolveUDPAddr("udp", jobConfig.addr)
-	if err != nil {
-		return nil, err
-	}
-
-	trafficMonitor := metrics.Default.NewWriter(metrics.Traffic, uuid.New().String())
-	go trafficMonitor.Update(ctx, time.Second)
-
-	if !isInEncryptedContext(ctx) {
-		log.Printf("Attacking %v", jobConfig.addr)
-	}
-
-	conn, err := net.DialUDP("udp", nil, udpAddr)
-	if err != nil {
-		logger.Debug("error connecting via tcp", zap.Reflect("addr", udpAddr), zap.Error(err))
-		metrics.IncRawnetUDP(udpAddr.String(), metrics.StatusFail)
-
-		return nil, err
-	}
-
-	defer conn.Close()
-
-	for jobConfig.Next(ctx) {
-		sendUDP(ctx, logger, udpAddr, conn, jobConfig.bodyTpl, trafficMonitor)
-	}
-
-	return nil, nil
-}
-
-func sendUDP(ctx context.Context, logger *zap.Logger, a *net.UDPAddr, conn *net.UDPConn, bodyTpl *template.Template, trafficMonitor *metrics.Writer) {
-	n, err := conn.Write([]byte(templates.Execute(logger, bodyTpl, ctx)))
-	if err != nil {
-		metrics.IncRawnetUDP(a.String(), metrics.StatusFail)
-
-		return
-	}
-
-	trafficMonitor.Add(uint64(n) + packetgen.UDPHeaderSize + packetgen.IPHeaderSize)
-	metrics.IncRawnetUDP(a.String(), metrics.StatusSuccess)
+	return packetgenJob(ctx, logger, globalConfig, packetgenArgs)
 }
 
-func parseRawNetJobArgs(ctx context.Context, logger *zap.Logger, globalConfig *GlobalConfig, args config.Args) (tpl *rawnetConfig, err error) {
+func parseRawNetJobArgs(ctx context.Context, logger *zap.Logger, globalConfig *GlobalConfig, args config.Args, protocol string) (
+	result map[string]any, err error,
+) {
 	var jobConfig struct {
 		BasicJobConfig
 
@@ -189,19 +72,23 @@ func parseRawNetJobArgs(ctx context.Context, logger *zap.Logger, globalConfig *G
 		return nil, fmt.Errorf("error decoding rawnet job config: %w", err)
 	}
 
-	bodyTpl, err := templates.Parse(jobConfig.Body)
-	if err != nil {
-		return nil, fmt.Errorf("error parsing body template %q: %w", jobConfig.Body, err)
+	args["connection"] = map[string]any{
+		"type": "net",
+		"args": map[string]any{
+			"protocol":   protocol,
+			"address":    jobConfig.Address,
+			"timeout":    jobConfig.Timeout,
+			"proxy_urls": jobConfig.ProxyURLs,
+		},
+	}
+	args["packet"] = map[string]any{
+		"payload": map[string]any{
+			"type": "raw",
+			"data": map[string]any{
+				"payload": jobConfig.Body,
+			},
+		},
 	}
 
-	targetAddress := strings.TrimSpace(templates.ParseAndExecute(logger, jobConfig.Address, ctx))
-	proxyURLs := templates.ParseAndExecute(logger, jobConfig.ProxyURLs, ctx)
-
-	return &rawnetConfig{
-		BasicJobConfig: jobConfig.BasicJobConfig,
-		addr:           targetAddress,
-		bodyTpl:        bodyTpl,
-		proxyURLs:      proxyURLs,
-		timeout:        utils.NonNilOrDefault(jobConfig.Timeout, time.Minute),
-	}, nil
+	return args, nil
 }