Merge pull request #263 from AikidoSec/attack-wave-samples

bitterpanda63 · web-flow · commit 38d0ffcbeaf0 · 2025-12-04T01:13:03.000+01:00
Attack wave samples
diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/events/DetectedAttackWave.java b/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/events/DetectedAttackWave.java
@@ -1,14 +1,17 @@
 package dev.aikido.agent_api.background.cloud.api.events;
 
+import com.google.gson.Gson;
 import dev.aikido.agent_api.background.cloud.GetManagerInfo;
 import dev.aikido.agent_api.context.ContextObject;
 import dev.aikido.agent_api.context.User;
+import dev.aikido.agent_api.storage.attack_wave_detector.AttackWaveDetector;
+import dev.aikido.agent_api.storage.attack_wave_detector.AttackWaveDetectorStore;
 
+import java.util.List;
 import java.util.Map;
 
 import static dev.aikido.agent_api.background.cloud.GetManagerInfo.getManagerInfo;
 import static dev.aikido.agent_api.helpers.UnixTimeMS.getUnixTimeMS;
-import static dev.aikido.agent_api.storage.ServiceConfigStore.getConfig;
 
 public final class DetectedAttackWave {
     private DetectedAttackWave() {
@@ -42,9 +45,15 @@ public static DetectedAttackWaveEvent createAPIEvent(ContextObject context) {
             context.getHeader("user-agent"), // userAgent
             context.getSource() // source
         );
+
+        String ip = context.getRemoteAddress();
+        List<AttackWaveDetector.Sample> samples = AttackWaveDetectorStore.getSamplesForIp(ip);
+        Map<String, String> metadata = Map.of(
+            "samples", new Gson().toJson(samples)
+        );
+
         AttackWaveData attackData = new AttackWaveData(
-            Map.of(),
-            context.getUser()
+            metadata, context.getUser()
         );
 
         return new DetectedAttackWaveEvent(
diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/attack_wave_detector/AttackWaveDetector.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/attack_wave_detector/AttackWaveDetector.java
@@ -3,19 +3,26 @@
 import dev.aikido.agent_api.context.ContextObject;
 import dev.aikido.agent_api.ratelimiting.LRUCache;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
 import static dev.aikido.agent_api.vulnerabilities.attack_wave_detection.WebScanDetector.isWebScanner;
 
 public class AttackWaveDetector {
-    private final LRUCache<String, Integer> suspiciousRequestsMap;
+    private final LRUCache<String, Integer> suspiciousRequestsCounts;
+    private final LRUCache<String, List<Sample>> suspiciousRequestsSamples;
     private final LRUCache<String, Long> sentEventsMap;
     private final int attackWaveThreshold;
+    private final int maxSamplesPerIP;
 
     public AttackWaveDetector() {
         this(
             /* attackWaveThreshold, default: 15 requests */ 15,
             /* attackWaveTimeFrame, default: 1 min */       60_000L,
             /* minTimeBetweenEvents, default: 20 min */     20 * 60_000L,
-            /* maxLRUEntries, default: 10,000 entries */    10_000
+            /* maxLRUEntries, default: 10,000 entries */    10_000,
+            /* maxSamplesPerIP, default: 15 samples */ 15
         );
     }
 
@@ -26,9 +33,11 @@ public AttackWaveDetector() {
      * @param maxLRUEntries        Maximum number of entries in the LRU cache
      */
     public AttackWaveDetector(int attackWaveThreshold, long attackWaveTimeFrame,
-                              long minTimeBetweenEvents, int maxLRUEntries) {
+                              long minTimeBetweenEvents, int maxLRUEntries, int maxSamplesPerIP) {
         this.attackWaveThreshold = attackWaveThreshold;
-        this.suspiciousRequestsMap = new LRUCache<>(maxLRUEntries, attackWaveTimeFrame);
+        this.maxSamplesPerIP = maxSamplesPerIP;
+        this.suspiciousRequestsCounts = new LRUCache<>(maxLRUEntries, attackWaveTimeFrame);
+        this.suspiciousRequestsSamples = new LRUCache<>(maxLRUEntries, attackWaveTimeFrame);
         this.sentEventsMap = new LRUCache<>(maxLRUEntries, minTimeBetweenEvents);
     }
 
@@ -56,16 +65,42 @@ public boolean check(ContextObject ctx) {
 
         // Add 1 to the suspiciousRequests counter.
         int suspiciousRequests = 1;
-        Integer existingSuspiciousRequests = this.suspiciousRequestsMap.get(ip);
+        Integer existingSuspiciousRequests = this.suspiciousRequestsCounts.get(ip);
         if (existingSuspiciousRequests != null) {
             suspiciousRequests = existingSuspiciousRequests + 1;
         }
-        this.suspiciousRequestsMap.set(ip, suspiciousRequests);
+        this.suspiciousRequestsCounts.set(ip, suspiciousRequests);
+
+        this.trackSample(ip, ctx.getMethod(), ctx.getUrl());
 
         if (suspiciousRequests < this.attackWaveThreshold) {
             return false;
         }
         this.sentEventsMap.set(ip, System.currentTimeMillis());
         return true;
     }
+
+    public List<Sample> getSamplesForIp(String ip) {
+        List<Sample> samples = this.suspiciousRequestsSamples.get(ip);
+        if (samples == null) {
+            samples = new ArrayList<>();
+        }
+        return samples;
+    }
+
+    public record Sample(String method, String url) {}
+    private void trackSample(String ip, String method, String url) {
+        List<Sample> samples = getSamplesForIp(ip);
+        if (samples.size() >= this.maxSamplesPerIP) {
+            return;
+        }
+
+        for (Sample sample : samples) {
+            if (Objects.equals(sample.method, method) && Objects.equals(sample.url, url)) {
+                return; // an equivalent entry already exists, skipping
+            }
+        }
+        samples.add(new Sample(method, url));
+        this.suspiciousRequestsSamples.set(ip, samples);
+    }
 }
diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/attack_wave_detector/AttackWaveDetectorStore.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/attack_wave_detector/AttackWaveDetectorStore.java
@@ -4,6 +4,7 @@
 import dev.aikido.agent_api.helpers.logging.LogManager;
 import dev.aikido.agent_api.helpers.logging.Logger;
 
+import java.util.List;
 import java.util.concurrent.locks.ReentrantLock;
 
 public final class AttackWaveDetectorStore {
@@ -25,4 +26,13 @@ public static boolean check(ContextObject ctx) {
             mutex.unlock();
         }
     }
+
+    public static List<AttackWaveDetector.Sample> getSamplesForIp(String ip) {
+        mutex.lock();
+        try {
+            return detector.getSamplesForIp(ip);
+        } finally {
+            mutex.unlock();
+        }
+    }
 }
diff --git a/agent_api/src/test/java/attack_wave_detection/AttackWaveDetectorTest.java b/agent_api/src/test/java/attack_wave_detection/AttackWaveDetectorTest.java
@@ -4,14 +4,15 @@
 import org.junit.jupiter.api.Test;
 import utils.EmptySampleContextObject;
 
-import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.*;
 
 class AttackWaveDetectorTest {
 
     private AttackWaveDetector newAttackWaveDetector() {
         // Use much smaller time frames for testing (e.g., 100ms instead of 60s)
-        return new AttackWaveDetector(6, 100L, 200L, 10_000);
+        return new AttackWaveDetector(6, 100L, 200L, 10_000, 3);
     }
 
     private static boolean checkDetector(AttackWaveDetector detector, String ip, boolean isWebScanner) {
@@ -59,6 +60,12 @@ void testWebScannerWithDelays() throws InterruptedException {
         assertFalse(checkDetector(detector, "::1", true));
         assertFalse(checkDetector(detector, "::1", true));
         assertFalse(checkDetector(detector, "::1", true));
+        assertArrayEquals(
+            List.of(
+                new AttackWaveDetector.Sample("BADMETHOD", "https://example.com/wp-config.php")
+            ).toArray(),
+            detector.getSamplesForIp("::1").toArray()
+        );
 
         // Small delay (50ms)
         Thread.sleep(50);
@@ -118,4 +125,24 @@ void testSlowWebScannerThirdInterval() throws InterruptedException {
         assertFalse(checkDetector(detector, "::1", true));
         assertTrue(checkDetector(detector, "::1", true));
     }
+
+    @Test
+    void testItRespectsSamplesLimit() {
+        AttackWaveDetector detector = newAttackWaveDetector();
+        detector.check(new EmptySampleContextObject("", "/../etc/passwd", "GET"));
+        detector.check(new EmptySampleContextObject("", "/../etc/passwd", "GET"));
+        detector.check(new EmptySampleContextObject("", "/test2", "GET"));
+        detector.check(new EmptySampleContextObject("", "/../etc/passwd", "POST"));
+        detector.check(new EmptySampleContextObject("", "/test3", "PUT"));
+        detector.check(new EmptySampleContextObject("", "/.env", "GET"));
+        detector.check(new EmptySampleContextObject("", "/test4", "BADMETHOD"));
+        assertArrayEquals(
+            List.of(
+                new AttackWaveDetector.Sample("GET", "https://example.com/../etc/passwd"),
+                new AttackWaveDetector.Sample("POST", "https://example.com/../etc/passwd"),
+                new AttackWaveDetector.Sample("GET", "https://example.com/.env")
+            ).toArray(),
+            detector.getSamplesForIp("192.168.1.1").toArray()
+        );
+    }
 }
diff --git a/agent_api/src/test/java/collectors/WebResponseCollectorTest.java b/agent_api/src/test/java/collectors/WebResponseCollectorTest.java
@@ -135,7 +135,9 @@ void testReport_WithAttackWaveContext() throws InterruptedException {
         assertEquals("123", event.attack().user().id());
         assertEquals("Jane Doe", event.attack().user().name());
 
-        assertEquals(0, event.attack().metadata().size());
+        assertEquals(1, event.attack().metadata().size());
+        assertEquals("[{\"method\":\"BADMETHOD\",\"url\":\"https://example.com/api/resource\"}]", event.attack().metadata().get("samples"));
+
         // check stats changed
         assertEquals(1, StatisticsStore.getStatsRecord().requests().attackWaves().total());
     }
diff --git a/agent_api/src/test/java/utils/EmptySampleContextObject.java b/agent_api/src/test/java/utils/EmptySampleContextObject.java
@@ -36,6 +36,7 @@ public EmptySampleContextObject(String argument, String route, String method) {
         this();
         this.query.put("arg", List.of(argument));
         this.route = route;
+        this.url = "https://example.com" + route;
         this.method = method;
     }
     public EmptySampleContextObject(String route, String method, Map<String, List<String>> queryParams) {

Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@`
`4`	`4`	`import dev.aikido.agent_api.helpers.logging.LogManager;`
`5`	`5`	`import dev.aikido.agent_api.helpers.logging.Logger;`
`6`	`6`
	`7`	`+import java.util.List;`
`7`	`8`	`import java.util.concurrent.locks.ReentrantLock;`
`8`	`9`
`9`	`10`	`public final class AttackWaveDetectorStore {`
`@@ -25,4 +26,13 @@ public static boolean check(ContextObject ctx) {`
`25`	`26`	`mutex.unlock();`
`26`	`27`	`}`
`27`	`28`	`}`
	`29`	`+`
	`30`	`+ public static List<AttackWaveDetector.Sample> getSamplesForIp(String ip) {`
	`31`	`+ mutex.lock();`
	`32`	`+ try {`
	`33`	`+ return detector.getSamplesForIp(ip);`
	`34`	`+ } finally {`
	`35`	`+ mutex.unlock();`
	`36`	`+ }`
	`37`	`+ }`
`28`	`38`	`}`
Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,7 @@ public EmptySampleContextObject(String argument, String route, String method) {`
`36`	`36`	`this();`
`37`	`37`	`this.query.put("arg", List.of(argument));`
`38`	`38`	`this.route = route;`
	`39`	`+ this.url = "https://example.com" + route;`
`39`	`40`	`this.method = method;`
`40`	`41`	`}`
`41`	`42`	`public EmptySampleContextObject(String route, String method, Map<String, List<String>> queryParams) {`