Merge pull request #192 from ByteTheCookies/dev

Enhance exploiter with multiprocessing and service-name validation

Author: akiidjk

README.md

@@ -140,7 +140,7 @@ This will generate `your_exploit_name.py` in `~/.cookiefarm/exploits/`.
 
 5. Run your exploit:
 ```bash
-ckc exploit run -e your_exploit_name.py -n CookieService -t 120 -T 40
+ckc exploit run -e your_exploit_name.py -n CookieService -t 120 -W 40
 ```
 > [!NOTE]
 > For more usage examples, check out the [client documentation](https://cookiefarm.bytethecookies.org/docs/client/overview).

benchmarks/ckvsdf/tool/start_benchmark.sh

@@ -49,7 +49,7 @@ echo "==> Starting CookieFarm Exploit..."
 # Example command to run CF exploit:
 just client-build-linux-prod
 ../../../bin/ckc login -P password
-../../../bin/ckc exploit run -e benchmark -n CookieService -t 5 -T 10 &
+../../../bin/ckc exploit run -e benchmark -n CookieService -t 5 -W 10 &
 PIDS+=($!)
 
 echo "==> Starting DestructiveFarm Exploit..."
@@ -75,7 +75,7 @@ echo "==> Benchmark is running. Waiting for $BENCHMARK_DUR seconds for completio
 wait $TIMELINE_PID
 
 echo "==> Benchmark duration completed. Stopping exploits and monitoring..."
-pkill -f "../../../bin/ckc exploit run -e benchmark -n CookieService -t 5 -T 10"
+pkill -f "../../../bin/ckc exploit run -e benchmark -n CookieService -t 5 -W 10"
 pkill -f "python3 ./start_sploit.py benchmark.py -u http://localhost:5000 --attack-period 5"
 pkill -f "python3 /tmp/DestructiveFarm/client/benchmark.py"
 

benchmarks/db_improve/notes.md

@@ -35,7 +35,7 @@ LAN
 LAN
 
 - Build command: `just client-build-prod`
-- Command run: `./bin/ckc exploit run -e benchmark -n CookieService -t 5 -T 10`
+- Command run: `./bin/ckc exploit run -e benchmark -n CookieService -t 5 -W 10`
 - Run command: `time python3 riempire_db_daiii.py`
 
 ### Cks config used:

benchmarks/protocol_beef/notes.md

@@ -34,7 +34,7 @@ LAN
 (same machine)
 
 - Build command: `just client-build-prod`
-- Command run: `./bin/ckc exploit run -e benchmark -n CookieService -t 5 -T 10`
+- Command run: `./bin/ckc exploit run -e benchmark -n CookieService -t 5 -W 10`
 
 ### Cks config used:
 

cookiefarm/client/cmd/exploit.go

@@ -25,6 +25,7 @@ var (
 	teamID      uint16 // Service name and team ID for the submit command
 	detach      bool   // Run the exploit in the background (detached mode)
 	submitValue bool   // This flag is used to directly submit to the game server
+	processMode bool   // Using multiprocessing instead of multi-threading to run the exploit
 	tickTime    int    // Interval in seconds between exploit executions
 	threadCount int    // Number of concurrent threads to run the exploit with
 	pid         int    // PID of the exploit to stop
@@ -125,8 +126,9 @@ func buildExploitCmd() *cobra.Command {
 	runCmd.Flags().StringVarP(&serviceName, "name", "n", "", "Service name to attack")
 	runCmd.Flags().BoolVarP(&detach, "detach", "d", false, "Run the exploit in the background (detached mode)")
 	runCmd.Flags().IntVarP(&tickTime, "tick", "t", 120, "Interval in seconds between exploit executions")
-	runCmd.Flags().IntVarP(&threadCount, "thread", "T", 5, "Number of concurrent threads to run the exploit with")
+	runCmd.Flags().IntVarP(&threadCount, "worker", "W", 10, "Number of concurrent worker to run the exploit with")
 	runCmd.Flags().BoolVarP(&submitValue, "submit", "S", false, "Direct submit to the game server")
+	runCmd.Flags().BoolVarP(&processMode, "process_mode", "P", false, "Using multiprocessing instead of multi-threading to run the exploit")
 	runCmd.MarkFlagRequired("exploit")
 	runCmd.MarkFlagRequired("port")
 	runCmd.MarkFlagRequired("host")
@@ -135,7 +137,8 @@ func buildExploitCmd() *cobra.Command {
 	testCmd.Flags().StringVarP(&serviceName, "name", "n", "", "Service name to attack")
 	testCmd.Flags().BoolVarP(&detach, "detach", "d", false, "Run the exploit in the background (detached mode)")
 	testCmd.Flags().IntVarP(&tickTime, "tick", "t", 120, "Interval in seconds between exploit executions")
-	testCmd.Flags().IntVarP(&threadCount, "thread", "T", 5, "Number of concurrent threads to run the exploit with")
+	testCmd.Flags().IntVarP(&threadCount, "worker", "W", 10, "Number of concurrent worker to run the exploit with")
+	testCmd.Flags().BoolVarP(&processMode, "process_mode", "P", false, "Using multiprocessing instead of multi-threading to run the exploit")
 	testCmd.MarkFlagRequired("exploit")
 	testCmd.MarkFlagRequired("port")
 	testCmd.MarkFlagRequired("host")
@@ -162,9 +165,10 @@ func test(cmd *cobra.Command, args []string) {
 	exploitArgs := exploit.ExploitArgs{
 		ServiceName: serviceName,
 		TickTime:    tickTime,
-		ThreadCount: threadCount,
+		WorkerCount: threadCount,
 		Detach:      detach,
 		ExploitPath: exploitPath,
+		ProcessMode: processMode,
 	}
 
 	exploit.Setup(&exploitArgs, detach)
@@ -185,9 +189,10 @@ func run(cmd *cobra.Command, args []string) {
 	exploitArgs := exploit.ExploitArgs{
 		ServiceName: serviceName,
 		TickTime:    tickTime,
-		ThreadCount: threadCount,
+		WorkerCount: threadCount,
 		Detach:      detach,
 		ExploitPath: exploitPath,
+		ProcessMode: processMode,
 	}
 
 	err := exploit.Setup(&exploitArgs, detach)

cookiefarm/client/internal/api/api.go

@@ -77,7 +77,7 @@ func Login(username, password string) error {
 	return nil
 }
 
-// @IMPORTANT: prefer websockets
+// @IMPORTANT: prefer ckp
 func SubmitBatchDirect(flags []database.Flag) error {
 	payload, err := json.Marshal(models.SubmitFlagsRequest{
 		Flags: flags,

cookiefarm/client/internal/ckp/submitter.go

@@ -16,7 +16,7 @@ func Start(flagsChan <-chan database.Flag) {
 
 	conn, err := NewClient(addr)
 	if err != nil {
-		logger.Log.Fatal().Err(err).Msg("Error connecting to WebSocket")
+		logger.Log.Fatal().Err(err).Msg("Error connecting to ckp")
 	}
 
 	defer conn.Close()

cookiefarm/client/internal/exploit/levenshtein.go

@@ -0,0 +1,126 @@
+package exploit
+
+// Go code from https://github.com/agnivade/levenshtein/blob/master/levenshtein.go
+// Package levenshtein is a Go implementation to calculate Levenshtein Distance.
+//
+// Implementation taken from
+// https://gist.github.com/andrei-m/982927#gistcomment-1931258
+
+import "unicode/utf8"
+
+// minLengthThreshold is the length of the string beyond which
+// an allocation will be made. Strings smaller than this will be
+// zero alloc.
+const minLengthThreshold = 32
+
+// ComputeDistance computes the levenshtein distance between the two
+// strings passed as an argument. The return value is the levenshtein distance
+//
+// Works on runes (Unicode code points) but does not normalize
+// the input strings. See https://blog.golang.org/normalization
+// and the golang.org/x/text/unicode/norm package.
+func ComputeDistance(a, b string) int {
+	if len(a) == 0 {
+		return utf8.RuneCountInString(b)
+	}
+
+	if len(b) == 0 {
+		return utf8.RuneCountInString(a)
+	}
+
+	if a == b {
+		return 0
+	}
+
+	// We need to convert to []rune if the strings are non-ASCII.
+	// This could be avoided by using utf8.RuneCountInString
+	// and then doing some juggling with rune indices,
+	// but leads to far more bounds checks. It is a reasonable trade-off.
+	s1 := []rune(a)
+	s2 := []rune(b)
+
+	// swap to save some memory O(min(a,b)) instead of O(a)
+	if len(s1) > len(s2) {
+		s1, s2 = s2, s1
+	}
+
+	// remove trailing identical runes.
+	s1, s2 = trimLongestCommonSuffix(s1, s2)
+
+	// Remove leading identical runes.
+	s1, s2 = trimLongestCommonPrefix(s1, s2)
+
+	lenS1 := len(s1)
+	lenS2 := len(s2)
+
+	// Init the row.
+	var x []uint16
+	if lenS1+1 > minLengthThreshold {
+		x = make([]uint16, lenS1+1)
+	} else {
+		// We make a small optimization here for small strings.
+		// Because a slice of constant length is effectively an array,
+		// it does not allocate. So we can re-slice it to the right length
+		// as long as it is below a desired threshold.
+		x = make([]uint16, minLengthThreshold)
+		x = x[:lenS1+1]
+	}
+
+	// we start from 1 because index 0 is already 0.
+	for i := 1; i < len(x); i++ {
+		x[i] = uint16(i)
+	}
+
+	// hoist bounds checks out of the loops
+	_ = x[lenS1]
+	y := x[1:]
+	y = y[:lenS1]
+	// fill in the rest
+	for i := range lenS2 {
+		prev := uint16(i + 1)
+		for j := range lenS1 {
+			current := x[j] // match
+			if s2[i] != s1[j] {
+				current = min(x[j], prev, y[j]) + 1
+			}
+			x[j] = prev
+			prev = current
+		}
+		x[lenS1] = prev
+	}
+	return int(x[lenS1])
+}
+
+func trimLongestCommonSuffix(a, b []rune) ([]rune, []rune) {
+	m := min(len(a), len(b))
+	a2 := a[len(a)-m:]
+	b2 := b[len(b)-m:]
+	i := len(a2)
+	b2 = b2[:i] // hoist bounds checks out of the loop
+	// Find the longest common suffix length using slices.IndexFunc
+	suffixLen := 0
+	for k := 1; k <= len(a2); k++ {
+		if a2[len(a2)-k] != b2[len(b2)-k] {
+			break
+		}
+		suffixLen = k
+	}
+	i -= len(a2) - suffixLen
+	return a[:len(a)-len(a2)+i], b[:len(b)-len(b2)+i]
+}
+
+func trimLongestCommonPrefix(a, b []rune) ([]rune, []rune) {
+	m := min(len(a), len(b))
+	i := m
+	for k := range m {
+		if a[k] != b[k] {
+			i = k
+			break
+		}
+	}
+	if i == m {
+		i = m
+	}
+
+	return a[i:], b[i:]
+}

cookiefarm/client/internal/exploit/parse.go

@@ -15,6 +15,7 @@ import (
 const (
 	INFO    = "info"
 	EXPLOIT = "exploit_info"
+	PROCESS = "process_info"
 	FAILED  = "failed"
 	ERROR   = "error"
 	FATAL   = "fatal"
@@ -42,9 +43,25 @@ func (p *Parser) Parse(line string) (database.Flag, string, error) {
 		return p.parseStats(line)
 	}
 
+	if strings.Contains(line, PROCESS) {
+		return p.parseProcessInfo(line)
+	}
+
 	return p.parseFlag(line)
 }
 
+func (*Parser) parseProcessInfo(line string) (database.Flag, string, error) {
+	var out ProcessInfoOutput
+
+	err := json.Unmarshal([]byte(line), &out)
+	if err != nil {
+		return database.Flag{}, INVALID, errors.New(line)
+	}
+
+	msg := fmt.Sprintf("Process Info: PID: %d, TeamID: %d, Event: %s", out.PID, out.TeamID, out.Event)
+	return database.Flag{}, PROCESS, errors.New(msg)
+}
+
 func (*Parser) parseStats(line string) (database.Flag, string, error) {
 	var out StatusBatchOutput
 	if err := json.Unmarshal([]byte(line), &out); err != nil {
@@ -64,7 +81,7 @@ func (p *Parser) parseFlag(line string) (database.Flag, string, error) {
 	}
 
 	switch out.Status {
-	case INFO, EXPLOIT, ERROR, DEBUG, FATAL:
+	case INFO, EXPLOIT, ERROR, DEBUG, FATAL, PROCESS:
 		return database.Flag{}, out.Status, errors.New(out.Message)
 	case FAILED:
 		return database.Flag{}, out.Status, errors.New(out.Message)

cookiefarm/client/internal/exploit/run.go

@@ -141,20 +141,50 @@ func formatMessage(raw string, err error) string {
 	return raw
 }
 
+func checkServiceName(cm *config.ConfigManager, serviceName string) error {
+	minDistance := 1000 // Set a high initial value for minimum distance
+	minDistanceName := ""
+	if cm.MapServiceToPort(serviceName) == 0 {
+		services := cm.Get().Shared.Services
+
+		for name := range services {
+			nameDistance := ComputeDistance(serviceName, name)
+			if nameDistance <= minDistance {
+				minDistance = nameDistance
+				minDistanceName = name
+			}
+		}
+
+		return fmt.Errorf("service name %s is not mapped to any port, did you mean %s (port %d)?", serviceName, minDistanceName, cm.MapServiceToPort(minDistanceName))
+	}
+
+	return nil
+}
+
 func buildCmd(args ExploitArgs, test bool) string {
 	cfg := config.GetInstance()
+	// cfg.Read()
+
+	if err := checkServiceName(cfg, args.ServiceName); err != nil {
+		logger.Log.Fatal().Err(err).Msg("Service name check failed")
+		return ""
+	}
 
 	cmd := fmt.Sprintf(
-		"%s -s %s:%d -t %d -T %d -p %d -n %s",
+		"%s -s %s:%d -t %d -W %d -p %d -n %s",
 		args.ExploitPath,
 		cfg.GetHost(),
 		cfg.GetPort(),
 		args.TickTime,
-		args.ThreadCount,
+		args.WorkerCount,
 		cfg.MapServiceToPort(args.ServiceName),
 		args.ServiceName,
 	)
 
+	if args.ProcessMode {
+		cmd += " -P"
+	}
+
 	if test {
 		cmd += " -x"
 	}