Merge pull request #512 from AikidoSec/AIK-4345-Add-a-way-to-only-all…

…ow-IP-traffic-from-certain-countries-Zen-for-Node

Implement ip allowlist (e.g. geo based)

end2end/server/src/handlers/lists.js

@@ -1,6 +1,7 @@
 const {
   getBlockedIPAddresses,
   getBlockedUserAgents,
+  getAllowedIPAddresses,
 } = require("../zen/config");
 
 module.exports = function lists(req, res) {
@@ -10,6 +11,7 @@ module.exports = function lists(req, res) {
 
   const blockedIps = getBlockedIPAddresses(req.app);
   const blockedUserAgents = getBlockedUserAgents(req.app);
+  const allowedIps = getAllowedIPAddresses(req.app);
 
   res.json({
     success: true,
@@ -25,5 +27,15 @@ module.exports = function lists(req, res) {
           ]
         : [],
     blockedUserAgents: blockedUserAgents,
+    allowedIPAddresses:
+      allowedIps.length > 0
+        ? [
+            {
+              source: "geoip",
+              description: "geo restrictions",
+              ips: allowedIps,
+            },
+          ]
+        : [],
   });
 };

end2end/server/src/handlers/updateLists.js

@@ -1,6 +1,7 @@
 const {
   updateBlockedIPAddresses,
   updateBlockedUserAgents,
+  updateAllowedIPAddresses,
 } = require("../zen/config");
 
 module.exports = function updateIPLists(req, res) {
@@ -38,5 +39,12 @@ module.exports = function updateIPLists(req, res) {
     updateBlockedUserAgents(req.app, req.body.blockedUserAgents);
   }
 
+  if (
+    req.body.allowedIPAddresses &&
+    Array.isArray(req.body.allowedIPAddresses)
+  ) {
+    updateAllowedIPAddresses(req.app, req.body.allowedIPAddresses);
+  }
+
   res.json({ success: true });
 };

end2end/server/src/zen/config.js

@@ -39,6 +39,7 @@ function updateAppConfig(app, newConfig) {
 
 const blockedIPAddresses = [];
 const blockedUserAgents = [];
+const allowedIPAddresses = [];
 
 function updateBlockedIPAddresses(app, ips) {
   let entry = blockedIPAddresses.find((ip) => ip.serviceId === app.serviceId);
@@ -64,6 +65,30 @@ function getBlockedIPAddresses(app) {
   return { serviceId: app.serviceId, ipAddresses: [] };
 }
 
+function updateAllowedIPAddresses(app, ips) {
+  let entry = allowedIPAddresses.find((ip) => ip.serviceId === app.serviceId);
+
+  if (entry) {
+    entry.ipAddresses = ips;
+  } else {
+    entry = { serviceId: app.serviceId, ipAddresses: ips };
+    allowedIPAddresses.push(entry);
+  }
+
+  // Bump lastUpdatedAt
+  updateAppConfig(app, {});
+}
+
+function getAllowedIPAddresses(app) {
+  const entry = allowedIPAddresses.find((ip) => ip.serviceId === app.serviceId);
+
+  if (entry) {
+    return entry.ipAddresses;
+  }
+
+  return { serviceId: app.serviceId, ipAddresses: [] };
+}
+
 function updateBlockedUserAgents(app, uas) {
   let entry = blockedUserAgents.find((e) => e.serviceId === e.serviceId);
 
@@ -95,4 +120,6 @@ module.exports = {
   getBlockedIPAddresses,
   updateBlockedUserAgents,
   getBlockedUserAgents,
+  getAllowedIPAddresses,
+  updateAllowedIPAddresses,
 };

end2end/tests/hono-xml-allowlists.test.ts

@@ -0,0 +1,144 @@
+const t = require("tap");
+const { spawn } = require("child_process");
+const { resolve } = require("path");
+const timeout = require("../timeout");
+
+const pathToApp = resolve(__dirname, "../../sample-apps/hono-xml", "app.js");
+const testServerUrl = "http://localhost:5874";
+
+let token;
+t.beforeEach(async () => {
+  const response = await fetch(`${testServerUrl}/api/runtime/apps`, {
+    method: "POST",
+  });
+  const body = await response.json();
+  token = body.token;
+
+  const config = await fetch(`${testServerUrl}/api/runtime/config`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: token,
+    },
+    body: JSON.stringify({
+      allowedIPAddresses: ["5.6.7.8"],
+      endpoints: [
+        {
+          route: "/admin",
+          method: "GET",
+          forceProtectionOff: false,
+          allowedIPAddresses: [],
+          rateLimiting: {
+            enabled: false,
+          },
+        },
+      ],
+    }),
+  });
+  t.same(config.status, 200);
+
+  const lists = await fetch(`${testServerUrl}/api/runtime/firewall/lists`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: token,
+    },
+    body: JSON.stringify({
+      allowedIPAddresses: ["4.3.2.1/32", "fe80::1234:5678:abcd:ef12/64"],
+      blockedIPAddresses: [],
+      blockedUserAgents: "hacker|attacker|GPTBot",
+    }),
+  });
+  t.same(lists.status, 200);
+});
+
+t.test("it blocks non-allowed IP addresses", (t) => {
+  const server = spawn(`node`, [pathToApp, "4002"], {
+    env: {
+      ...process.env,
+      AIKIDO_DEBUG: "true",
+      AIKIDO_BLOCK: "true",
+      AIKIDO_TOKEN: token,
+      AIKIDO_URL: testServerUrl,
+    },
+  });
+
+  server.on("close", () => {
+    t.end();
+  });
+
+  server.on("error", (err) => {
+    t.fail(err);
+  });
+
+  let stdout = "";
+  server.stdout.on("data", (data) => {
+    stdout += data.toString();
+  });
+
+  let stderr = "";
+  server.stderr.on("data", (data) => {
+    stderr += data.toString();
+  });
+
+  // Wait for the server to start
+  timeout(2000)
+    .then(async () => {
+      const resp1 = await fetch("http://127.0.0.1:4002/add", {
+        method: "POST",
+        body: "<cat><name>Njuska</name></cat>",
+        headers: {
+          "Content-Type": "application/xml",
+          "X-Forwarded-For": "1.3.2.4",
+        },
+        signal: AbortSignal.timeout(5000),
+      });
+      t.same(resp1.status, 403);
+      t.same(
+        await resp1.text(),
+        "Your IP address is not allowed to access this resource. (Your IP: 1.3.2.4)"
+      );
+
+      const resp2 = await fetch("http://127.0.0.1:4002/add", {
+        method: "POST",
+        body: "<cat><name>Harry</name></cat>",
+        headers: {
+          "Content-Type": "application/xml",
+          "X-Forwarded-For": "fe80::1234:5678:abcd:ef12",
+        },
+        signal: AbortSignal.timeout(5000),
+      });
+      t.same(resp2.status, 200);
+      t.same(await resp2.text(), JSON.stringify({ success: true }));
+
+      const resp3 = await fetch("http://127.0.0.1:4002/add", {
+        method: "POST",
+        body: "<cat><name>Harry</name></cat>",
+        headers: {
+          "Content-Type": "application/xml",
+          "X-Forwarded-For": "4.3.2.1",
+        },
+        signal: AbortSignal.timeout(5000),
+      });
+      t.same(resp3.status, 200);
+      t.same(await resp3.text(), JSON.stringify({ success: true }));
+
+      const resp4 = await fetch("http://127.0.0.1:4002/add", {
+        method: "POST",
+        body: "<cat><name>Harry2</name></cat>",
+        headers: {
+          "Content-Type": "application/xml",
+          "X-Forwarded-For": "5.6.7.8",
+        },
+        signal: AbortSignal.timeout(5000),
+      });
+      t.same(resp4.status, 200);
+      t.same(await resp4.text(), JSON.stringify({ success: true }));
+    })
+    .catch((error) => {
+      t.fail(error);
+    })
+    .finally(() => {
+      server.kill();
+    });
+});

library/agent/Agent.test.ts

@@ -19,6 +19,8 @@ import { Context } from "./Context";
 import { createTestAgent } from "../helpers/createTestAgent";
 import { setTimeout } from "node:timers/promises";
 
+let shouldOnlyAllowSomeIPAddresses = false;
+
 wrap(fetch, "fetch", function mock() {
   return async function mock() {
     return {
@@ -32,6 +34,15 @@ wrap(fetch, "fetch", function mock() {
           },
         ],
         blockedUserAgents: "AI2Bot|Bytespider",
+        allowedIPAddresses: shouldOnlyAllowSomeIPAddresses
+          ? [
+              {
+                source: "name",
+                description: "Description",
+                ips: ["4.3.2.1"],
+              },
+            ]
+          : [],
       }),
     };
   };
@@ -1099,6 +1110,10 @@ t.test("it does not fetch blocked IPs if serverless", async () => {
     blocked: false,
   });
 
+  t.same(agent.getConfig().isAllowedIPAddress("1.3.2.4"), {
+    allowed: true,
+  });
+
   t.same(
     agent
       .getConfig()
@@ -1110,3 +1125,31 @@ t.test("it does not fetch blocked IPs if serverless", async () => {
     }
   );
 });
+
+t.test("it only allows some IP addresses", async () => {
+  shouldOnlyAllowSomeIPAddresses = true;
+  const agent = createTestAgent({
+    token: new Token("123"),
+    suppressConsoleLog: false,
+  });
+
+  agent.start([]);
+
+  await setTimeout(0);
+
+  t.same(agent.getConfig().isIPAddressBlocked("1.3.2.4"), {
+    blocked: true,
+    reason: "Description",
+  });
+  t.same(agent.getConfig().isIPAddressBlocked("fe80::1234:5678:abcd:ef12"), {
+    blocked: true,
+    reason: "Description",
+  });
+
+  t.same(agent.getConfig().isAllowedIPAddress("1.2.3.4"), {
+    allowed: false,
+  });
+  t.same(agent.getConfig().isAllowedIPAddress("4.3.2.1"), {
+    allowed: true,
+  });
+});

library/agent/Agent.ts

@@ -41,7 +41,15 @@ export class Agent {
   private timeoutInMS = 10000;
   private hostnames = new Hostnames(200);
   private users = new Users(1000);
-  private serviceConfig = new ServiceConfig([], Date.now(), [], [], true, []);
+  private serviceConfig = new ServiceConfig(
+    [],
+    Date.now(),
+    [],
+    [],
+    true,
+    [],
+    []
+  );
   private routes: Routes = new Routes(200);
   private rateLimiter: RateLimiter = new RateLimiter(5000, 120 * 60 * 1000);
   private statistics = new InspectionStatistics({
@@ -363,11 +371,11 @@ export class Agent {
     }
 
     try {
-      const { blockedIPAddresses, blockedUserAgents } = await fetchBlockedLists(
-        this.token
-      );
+      const { blockedIPAddresses, blockedUserAgents, allowedIPAddresses } =
+        await fetchBlockedLists(this.token);
       this.serviceConfig.updateBlockedIPAddresses(blockedIPAddresses);
       this.serviceConfig.updateBlockedUserAgents(blockedUserAgents);
+      this.serviceConfig.updateAllowedIPAddresses(allowedIPAddresses);
     } catch (error: any) {
       console.error(`Aikido: Failed to update blocked lists: ${error.message}`);
     }