Merge pull request #128 from AikidoSec/add-allowlist-global

bitterpanda63 · web-flow · commit de374b0c4e3b · 2025-03-25T16:03:51.000+01:00
Add support for global IP allowlist (e.g. geo-fencing)
diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/ReportingApi.java b/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/ReportingApi.java
@@ -29,11 +29,16 @@ public ReportingApi(int timeoutInSec) {
      */
     public abstract Optional<APIResponse> report(String token, APIEvent event);
 
-    public record APIListsResponse(List<ListsResponseEntry> blockedIPAddresses, String blockedUserAgents) {}
+    public record APIListsResponse(
+            List<ListsResponseEntry> blockedIPAddresses,
+            List<ListsResponseEntry> allowedIPAddresses,
+            String blockedUserAgents
+    ) {}
     public record ListsResponseEntry(String source, String description, List<String> ips) {}
     /**
      * Fetch blocked lists using a separate API call, these can include :
      * -> blocked IP Addresses (e.g. geo restrictions)
+     * -> allowed IP Addresses (e.g. geo restrictions)
      * -> blocked User-Agents (e.g. bot blocking)
      * @param token the authentication token
      */
diff --git a/agent_api/src/main/java/dev/aikido/agent_api/thread_cache/ThreadCacheObject.java b/agent_api/src/main/java/dev/aikido/agent_api/thread_cache/ThreadCacheObject.java
@@ -14,6 +14,7 @@
 
 import static dev.aikido.agent_api.helpers.IPListBuilder.createIPList;
 import static dev.aikido.agent_api.helpers.UnixTimeMS.getUnixTimeMS;
+import static dev.aikido.agent_api.vulnerabilities.ssrf.IsPrivateIP.isPrivateIp;
 
 public class ThreadCacheObject {
     private final List<Endpoint> endpoints;
@@ -23,9 +24,10 @@ public class ThreadCacheObject {
     private final Hostnames hostnames;
     private final Routes routes;
 
-    // IP Blocking (e.g. Geo-IP Restrictions) :
-    public record BlockedIpEntry(IPList blocklist, String description) {}
-    private List<BlockedIpEntry> blockedIps = new ArrayList<>();
+    // IP restrictions (e.g. Geo-IP Restrictions) :
+    public record IPListEntry(IPList ipList, String description) {}
+    private List<IPListEntry> blockedIps = new ArrayList<>();
+    private List<IPListEntry> allowedIps = new ArrayList<>();
     // User-Agent Blocking (e.g. bot blocking) :
     private Pattern blockedUserAgentRegex;
 
@@ -72,8 +74,24 @@ public boolean isBypassedIP(String ip) {
      * Check if the IP is blocked (e.g. Geo IP Restrictions)
      */
     public BlockedResult isIpBlocked(String ip) {
-        for (BlockedIpEntry entry: blockedIps) {
-            if (entry.blocklist.matches(ip)) {
+        // Check for allowed ip addresses (i.e. only one country is allowed to visit the site)
+        // Always allow access from private IP addresses (those include local IP addresses)
+        if(allowedIps != null && allowedIps.size() > 0 && !isPrivateIp(ip)) {
+            boolean ipAllowed = false;
+            for (IPListEntry entry: allowedIps) {
+                if (entry.ipList.matches(ip)) {
+                    ipAllowed = true; // We allow IP addresses as long as they match with one of the lists.
+                    break;
+                }
+            }
+            if (!ipAllowed) {
+                return new BlockedResult(true, "allowlist");
+            }
+        }
+
+        // Check for blocked ip addresses
+        for (IPListEntry entry: blockedIps) {
+            if (entry.ipList.matches(ip)) {
                 return new BlockedResult(true, entry.description);
             }
         }
@@ -87,7 +105,14 @@ public void updateBlockedLists(Optional<ReportingApi.APIListsResponse> blockedLi
             if (res.blockedIPAddresses() != null) {
                 for (ReportingApi.ListsResponseEntry entry : res.blockedIPAddresses()) {
                     IPList ipList = createIPList(entry.ips());
-                    blockedIps.add(new BlockedIpEntry(ipList, entry.description()));
+                    blockedIps.add(new IPListEntry(ipList, entry.description()));
+                }
+            }
+            // Update allowed IP addresses (e.g. for geo restrictions) :
+            if (res.allowedIPAddresses() != null) {
+                for (ReportingApi.ListsResponseEntry entry: res.allowedIPAddresses()) {
+                    IPList ipList = createIPList(entry.ips());
+                    this.allowedIps.add(new IPListEntry(ipList, entry.description()));
                 }
             }
             // Update Blocked User-Agents regex
diff --git a/agent_api/src/test/java/background/ServiceConfigurationTest.java b/agent_api/src/test/java/background/ServiceConfigurationTest.java
@@ -104,7 +104,7 @@ void updateConfig_ShouldHandleAllUpdates_WhenApiResponseIsComplete() {
     @Test
     void testSetForBlockedIpRes() {
         assertNull(serviceConfiguration.blockedListsRes);
-        serviceConfiguration.storeBlockedListsRes(Optional.of(new ReportingApi.APIListsResponse(null, null)));
+        serviceConfiguration.storeBlockedListsRes(Optional.of(new ReportingApi.APIListsResponse(null, null, null)));
         assertNotNull(serviceConfiguration.blockedListsRes);
         serviceConfiguration.storeBlockedListsRes(Optional.empty());
 
diff --git a/agent_api/src/test/java/collectors/WebRequestCollectorTest.java b/agent_api/src/test/java/collectors/WebRequestCollectorTest.java
@@ -23,7 +23,7 @@
 
 class WebRequestCollectorTest {
 
-    private ContextObject contextObject;
+    private EmptySampleContextObject contextObject;
     private ThreadCacheObject threadCacheObject;
 
     @BeforeEach
@@ -105,7 +105,7 @@ void testReport_noThreadCacheObject() {
     void testReport_ipBlockedTwice() {
         ReportingApi.APIListsResponse blockedListsRes = new ReportingApi.APIListsResponse(List.of(
                 new ReportingApi.ListsResponseEntry("geoip", "geoip restrictions", List.of("192.168.1.1"))
-        ), "");
+        ), null, "");
         // Mock ThreadCache
         threadCacheObject = new ThreadCacheObject(List.of(new Endpoint(
                 "GET", "/api/resource", 100, 100,
@@ -126,7 +126,7 @@ void testReport_ipBlockedTwice() {
     void testReport_ipBlockedUsingLists() {
         ReportingApi.APIListsResponse blockedListsRes = new ReportingApi.APIListsResponse(List.of(
                 new ReportingApi.ListsResponseEntry("geoip", "geoip restrictions", List.of("bullshit.ip", "192.168.1.1"))
-        ), "");
+        ), null, "");
         // Mock ThreadCache
         threadCacheObject = new ThreadCacheObject(List.of(), Set.of(), Set.of(), new Routes(), Optional.of(blockedListsRes));
         ThreadCache.set(threadCacheObject);
@@ -138,13 +138,49 @@ void testReport_ipBlockedUsingLists() {
         assertEquals("Your IP address is not allowed to access this resource. (Your IP: 192.168.1.1)", response.msg());
         assertEquals(403, response.status());
     }
+    @SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "test-token")
+    @Test
+    void testReport_ipNotAllowedUsingLists() {
+        ReportingApi.APIListsResponse blockedListsRes = new ReportingApi.APIListsResponse(null, List.of(
+                new ReportingApi.ListsResponseEntry("geoip", "geoip restrictions", List.of("192.168.2.1"))
+        ), "");
+        // Mock ThreadCache
+        threadCacheObject = new ThreadCacheObject(List.of(), Set.of(), Set.of(), new Routes(), Optional.of(blockedListsRes));
+        ThreadCache.set(threadCacheObject);
+
+
+        WebRequestCollector.Res response = WebRequestCollector.report(contextObject);
+
+        assertNull(response); // Private IP
+
+        contextObject.setIp("4.4.4.4");
+        response = WebRequestCollector.report(contextObject);
+        assertNotNull(response);
+        assertEquals("Your IP address is not allowed to access this resource. (Your IP: 4.4.4.4)", response.msg());
+        assertEquals(403, response.status());
+    }
+    @SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "test-token")
+    @Test
+    void testReport_ipInAllowlist() {
+        ReportingApi.APIListsResponse blockedListsRes = new ReportingApi.APIListsResponse(null, List.of(
+                new ReportingApi.ListsResponseEntry("geoip", "geoip restrictions", List.of("192.168.1.1", "10.0.0.0/24"))
+        ), "");
+        // Mock ThreadCache
+        threadCacheObject = new ThreadCacheObject(List.of(), Set.of(), Set.of(), new Routes(), Optional.of(blockedListsRes));
+        ThreadCache.set(threadCacheObject);
+
+
+        WebRequestCollector.Res response = WebRequestCollector.report(contextObject);
+
+        assertNull(response);
+    }
 
     @SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "test-token")
     @Test
     void testReport_ipNotBlockedUsingListsNorUserAgent() {
         ReportingApi.APIListsResponse blockedListsRes = new ReportingApi.APIListsResponse(List.of(
                 new ReportingApi.ListsResponseEntry("geoip", "geoip restrictions", List.of("192.168.1.2", "192.168.1.3"))
-        ), "Unrelated|random");
+        ), null, "Unrelated|random");
         // Mock ThreadCache
         threadCacheObject = new ThreadCacheObject(List.of(), Set.of(), Set.of(), new Routes(), Optional.of(blockedListsRes));
         ThreadCache.set(threadCacheObject);
@@ -160,7 +196,7 @@ void testReport_ipNotBlockedUsingListsNorUserAgent() {
     void testReport_userAgentBlocked() {
         ReportingApi.APIListsResponse blockedListsRes = new ReportingApi.APIListsResponse(List.of(
                 new ReportingApi.ListsResponseEntry("geoip", "geoip restrictions", List.of("192.168.1.2", "192.168.1.3"))
-        ), "AI2Bot|hacker");
+        ), null, "AI2Bot|hacker");
         // Mock ThreadCache
         threadCacheObject = new ThreadCacheObject(List.of(), Set.of(), Set.of(), new Routes(), Optional.of(blockedListsRes));
         ThreadCache.set(threadCacheObject);
@@ -178,7 +214,7 @@ void testReport_userAgentBlocked() {
     void testReport_userAgentBlocked_Ip_Bypassed() {
         ReportingApi.APIListsResponse blockedListsRes = new ReportingApi.APIListsResponse(List.of(
                 new ReportingApi.ListsResponseEntry("geoip", "geoip restrictions", List.of("192.168.1.2", "192.168.1.3"))
-        ), "AI2Bot|hacker");
+        ), null, "AI2Bot|hacker");
         // Mock ThreadCache
         threadCacheObject = new ThreadCacheObject(List.of(), Set.of(),
                 /* bypassedIps : */ Set.of("192.168.1.1"), new Routes(), Optional.of(blockedListsRes));
@@ -196,7 +232,27 @@ void testReport_userAgentBlocked_Ip_Bypassed() {
     void testReport_ipBlockedUsingLists_Ip_Bypassed() {
         ReportingApi.APIListsResponse blockedListsRes = new ReportingApi.APIListsResponse(List.of(
                 new ReportingApi.ListsResponseEntry("geoip", "geoip restrictions", List.of("bullshit.ip", "192.168.1.1"))
-        ), "");
+        ), null, "");
+        // Mock ThreadCache
+        threadCacheObject = new ThreadCacheObject(List.of(), Set.of(),
+                /* bypassedIps : */ Set.of("192.168.1.1"), new Routes(), Optional.of(blockedListsRes));
+        ThreadCache.set(threadCacheObject);
+
+
+        WebRequestCollector.Res response = WebRequestCollector.report(contextObject);
+
+        assertNull(response);
+        assertNull(Context.get());
+    }
+
+    @SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "test-token")
+    @Test
+    void testReport_ipNotAllowedUsingLists_Ip_Bypassed() {
+        ReportingApi.APIListsResponse blockedListsRes = new ReportingApi.APIListsResponse(
+            null,
+            List.of(new ReportingApi.ListsResponseEntry("geoip", "geoip restrictions", List.of("1.2.3.4"))),
+            ""
+        );
         // Mock ThreadCache
         threadCacheObject = new ThreadCacheObject(List.of(), Set.of(),
                 /* bypassedIps : */ Set.of("192.168.1.1"), new Routes(), Optional.of(blockedListsRes));
diff --git a/agent_api/src/test/java/thread_cache/ThreadCacheObjectTest.java b/agent_api/src/test/java/thread_cache/ThreadCacheObjectTest.java
@@ -1,10 +1,8 @@
 package thread_cache;
 
-import dev.aikido.agent_api.background.ServiceConfiguration;
 import dev.aikido.agent_api.background.cloud.api.ReportingApi;
 import dev.aikido.agent_api.thread_cache.ThreadCacheObject;
 import org.junit.jupiter.api.Test;
-import utils.EmtpyThreadCacheObject;
 
 import java.util.List;
 import java.util.Optional;
@@ -24,7 +22,7 @@ public void update() {
                         "fd00:3234:5678:9abc::1/64",
                         "5.6.7.8/32"
                 ))
-        ), "Test|One")));
+        ), null, "Test|One")));
 
         assertEquals(new ThreadCacheObject.BlockedResult(true, "description"), tCache.isIpBlocked("1.2.3.4"));
         assertEquals(new ThreadCacheObject.BlockedResult(false, null), tCache.isIpBlocked("2.3.4.5"));
@@ -59,9 +57,9 @@ public void updateEmpty() {
                         "fd00:3234:5678:9abc::1/64",
                         "5.6.7.8/32"
                 ))
-        ), "Test|One")));
+        ), null, "Test|One")));
 
-        tCache.updateBlockedLists(Optional.of(new ReportingApi.APIListsResponse(null, null)));
+        tCache.updateBlockedLists(Optional.of(new ReportingApi.APIListsResponse(null,  null,null)));
 
         assertEquals(new ThreadCacheObject.BlockedResult(true, "description"), tCache.isIpBlocked("1.2.3.4"));
         assertEquals(new ThreadCacheObject.BlockedResult(false, null), tCache.isIpBlocked("2.3.4.5"));
@@ -96,16 +94,16 @@ public void updateRegexes() {
                         "fd00:3234:5678:9abc::1/64",
                         "5.6.7.8/32"
                 ))
-        ), "Test|One")));
+        ), null, "Test|One")));
 
-        tCache.updateBlockedLists(Optional.of(new ReportingApi.APIListsResponse(null, "")));
+        tCache.updateBlockedLists(Optional.of(new ReportingApi.APIListsResponse(null, null, "")));
 
         assertTrue(tCache.isBlockedUserAgent("This is my TEST user agent"));
         assertTrue(tCache.isBlockedUserAgent("Test"));
         assertTrue(tCache.isBlockedUserAgent("TEst and ONE"));
         assertFalse(tCache.isBlockedUserAgent("Est|On"));
         assertFalse(tCache.isBlockedUserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"));
-        tCache.updateBlockedLists(Optional.of(new ReportingApi.APIListsResponse(null, "Mozilla")));
+        tCache.updateBlockedLists(Optional.of(new ReportingApi.APIListsResponse(null, null, "Mozilla")));
 
         assertFalse(tCache.isBlockedUserAgent("This is my TEST user agent"));
         assertFalse(tCache.isBlockedUserAgent("Test"));
@@ -155,4 +153,93 @@ public void testThreadCacheBypassedIPsSubnet() {
         assertFalse(tCache.isBypassedIP("10.0.1.1"));
         assertFalse(tCache.isBypassedIP("1.2.3.4"));
     }
+
+    @Test
+    public void testIsIpBlockedWithAllowedAndBlockedIPs() {
+        // Create a ThreadCacheObject with both allowed and blocked IPs
+        ThreadCacheObject tCache = new ThreadCacheObject(null, null, null, null, Optional.of(new ReportingApi.APIListsResponse(List.of(
+                new ReportingApi.ListsResponseEntry("geoip", "description", List.of(
+                        "1.2.3.4", // Blocked IP
+                        "192.168.1.1" // Blocked IP
+                ))
+        ), List.of(
+                new ReportingApi.ListsResponseEntry("geoip", "description", List.of(
+                        "10.0.0.1", // Allowed IP
+                        "1.2.3.4"
+                ))
+        ), "Test|One")));
+
+        // Test blocked IPs
+        assertEquals(new ThreadCacheObject.BlockedResult(true, "description"), tCache.isIpBlocked("1.2.3.4"));
+        assertEquals(new ThreadCacheObject.BlockedResult(true, "description"), tCache.isIpBlocked("192.168.1.1"));
+
+        // Test allowed IPs
+        /// Private IP :
+        assertEquals(new ThreadCacheObject.BlockedResult(false, null), tCache.isIpBlocked("10.0.0.2"));
+        /// Not in allowlist
+        assertEquals(new ThreadCacheObject.BlockedResult(true, "allowlist"), tCache.isIpBlocked("1.2.3.3"));
+    }
+
+    @Test
+    public void testIsIpBlockedWithOnlyAllowedIPs() {
+        // Create a ThreadCacheObject with only allowed IPs
+        ThreadCacheObject tCache = new ThreadCacheObject(null, null, null, null, Optional.of(new ReportingApi.APIListsResponse(null, List.of(
+                new ReportingApi.ListsResponseEntry("geoip", "description", List.of(
+                        "10.0.0.1" // Allowed IP
+                ))
+        ), "Test|One")));
+
+        // Test allowed IP
+        assertEquals(new ThreadCacheObject.BlockedResult(false, null), tCache.isIpBlocked("10.0.0.1"));
+        // Test a non-allowed private-IP
+        assertEquals(new ThreadCacheObject.BlockedResult(false, null), tCache.isIpBlocked("10.0.0.2"));
+        // Test a non-allowed IP
+        assertEquals(new ThreadCacheObject.BlockedResult(true, "allowlist"), tCache.isIpBlocked("1.2.3.4"));
+    }
+
+    @Test
+    public void testIsIpBlockedWithOnlyBlockedIPs() {
+        // Create a ThreadCacheObject with only blocked IPs
+        ThreadCacheObject tCache = new ThreadCacheObject(null, null, null, null, Optional.of(new ReportingApi.APIListsResponse(List.of(
+                new ReportingApi.ListsResponseEntry("geoip", "description", List.of(
+                        "1.2.3.4", // Blocked IP
+                        "192.168.1.1" // Blocked IP
+                ))
+        ), null, "Test|One")));
+
+        // Test blocked IPs
+        assertEquals(new ThreadCacheObject.BlockedResult(true, "description"), tCache.isIpBlocked("1.2.3.4"));
+        assertEquals(new ThreadCacheObject.BlockedResult(true, "description"), tCache.isIpBlocked("192.168.1.1"));
+        // Test a non-blocked IP
+        assertEquals(new ThreadCacheObject.BlockedResult(false, null), tCache.isIpBlocked("10.0.0.1"));
+    }
+
+    @Test
+    public void testIsIpBlockedWithAllowedIPsAndBlockedIPs() {
+        // Create a ThreadCacheObject with multiple allowed and blocked IPs
+        ThreadCacheObject tCache = new ThreadCacheObject(null, null, null, null, Optional.of(new ReportingApi.APIListsResponse(null, List.of(
+                new ReportingApi.ListsResponseEntry("geoip1", "description", List.of(
+                        "1.2.3.4" // Blocked IP
+                )),
+                new ReportingApi.ListsResponseEntry("geoip2", "description", List.of(
+                        "8.8.8.0/24"
+                )),
+                new ReportingApi.ListsResponseEntry("geoip3", "description", List.of(
+                        "4.4.4.4" // Another allowed IP
+                ))
+        ), "Test|One")));
+
+        // Test allowed IPs
+        assertEquals(new ThreadCacheObject.BlockedResult(false, null), tCache.isIpBlocked("10.0.0.1"));
+        assertEquals(new ThreadCacheObject.BlockedResult(false, null), tCache.isIpBlocked("4.4.4.4"));
+        assertEquals(new ThreadCacheObject.BlockedResult(false, null), tCache.isIpBlocked("1.2.3.4"));
+        assertEquals(new ThreadCacheObject.BlockedResult(false, null), tCache.isIpBlocked("8.8.8.1"));
+
+        // Test a non-allowed IP
+        assertEquals(new ThreadCacheObject.BlockedResult(true, "allowlist"), tCache.isIpBlocked("4.4.4.1"));
+        assertEquals(new ThreadCacheObject.BlockedResult(true, "allowlist"), tCache.isIpBlocked("8.8.7.8"));
+
+    }
+
+
 }
diff --git a/agent_api/src/test/java/utils/EmptySampleContextObject.java b/agent_api/src/test/java/utils/EmptySampleContextObject.java
@@ -33,4 +33,7 @@ public EmptySampleContextObject(String argument, String route, String method) {
         this.route = route;
         this.method = method;
     }
+    public void setIp(String ip) {
+        this.remoteAddress = ip;
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -33,4 +33,7 @@ public EmptySampleContextObject(String argument, String route, String method) {`
`33`	`33`	`this.route = route;`
`34`	`34`	`this.method = method;`
`35`	`35`	`}`
	`36`	`+ public void setIp(String ip) {`
	`37`	`+ this.remoteAddress = ip;`
	`38`	`+ }`
`36`	`39`	`}`