Revert "Update all other code to remove IPList class"

This reverts commit fec4645.

Author: bitterpanda63

aikido_zen/background_process/commands/sync_data_test.py

@@ -4,8 +4,8 @@
 from unittest.mock import MagicMock
 from .sync_data import process_sync_data
 from aikido_zen.background_process.routes import Routes
+from aikido_zen.helpers.iplist import IPList
 from ..packages import PackagesStore
-from ...helpers.ip_matcher import IPMatcher
 from ...storage.hostnames import Hostnames
 from ...storage.statistics import Statistics
 
@@ -18,7 +18,8 @@ def setup_connection_manager():
     connection_manager.hostnames = Hostnames()
     connection_manager.conf.endpoints = ["endpoint1", "endpoint2"]
 
-    connection_manager.conf.bypassed_ips = IPMatcher(["192.168.1.1"])
+    connection_manager.conf.bypassed_ips = IPList()
+    connection_manager.conf.bypassed_ips.add("192.168.1.1")
     connection_manager.conf.blocked_uids = ["user1", "user2"]
     connection_manager.conf.last_updated_at = 200
     connection_manager.statistics = Statistics()

aikido_zen/background_process/service_config.py

@@ -5,9 +5,8 @@
 from typing import Pattern
 
 import regex as re
-
-from aikido_zen.helpers.ip_matcher import IPMatcher
 from aikido_zen.helpers.match_endpoints import match_endpoints
+from aikido_zen.helpers.iplist import IPList
 
 
 # noinspection PyAttributeOutsideInit
@@ -51,7 +50,7 @@ def set_endpoints(self, endpoints):
             endpoint for endpoint in endpoints if not endpoint.get("graphql")
         ]
 
-        # Create an IPMatcher instance for each endpoint
+        # Create a IPList instance for each endpoint
         for endpoint in self.endpoints:
             if not "allowedIPAddresses" in endpoint:
                 #  This feature is not supported by the current aikido server version
@@ -63,7 +62,7 @@ def set_endpoints(self, endpoints):
                 #  Skip empty allowlist
                 continue
 
-            endpoint["allowedIPAddresses"] = IPMatcher(endpoint["allowedIPAddresses"])
+            endpoint["allowedIPAddresses"] = IPList(endpoint["allowedIPAddresses"])
 
     def get_endpoints(self, route_metadata):
         """
@@ -73,14 +72,14 @@ def get_endpoints(self, route_metadata):
         return match_endpoints(route_metadata, self.endpoints)
 
     def set_bypassed_ips(self, bypassed_ips):
-        """Creates an IPMatcher from the given bypassed ip set"""
-        self.bypassed_ips = IPMatcher()
+        """Creates an IPList from the given bypassed ip set"""
+        self.bypassed_ips = IPList()
         for ip in bypassed_ips:
             self.bypassed_ips.add(ip)
 
     def is_bypassed_ip(self, ip):
         """Checks if the IP is on the bypass list"""
-        return self.bypassed_ips.has(ip)
+        return self.bypassed_ips.matches(ip)
 
     def set_blocked_ips(self, blocked_ip_entries):
         self.blocked_ips = list(map(parse_ip_entry, blocked_ip_entries))
@@ -90,7 +89,7 @@ def set_allowed_ips(self, allowed_ip_entries):
 
     def is_blocked_ip(self, ip):
         for entry in self.blocked_ips:
-            if entry["iplist"].has(ip):
+            if entry["iplist"].matches(ip):
                 return entry["description"]
         return False
 
@@ -112,8 +111,12 @@ def parse_ip_entry(entry):
     """
     Converts ip entry: {"source": "example", "description": "Example description", "ips": []}
     """
+    iplist = IPList()
+    for ip in entry["ips"]:
+        iplist.add(ip)
+
     return {
         "source": entry["source"],
         "description": entry["description"],
-        "iplist": IPMatcher(entry["ips"]),
+        "iplist": iplist,
     }

aikido_zen/background_process/service_config_test.py

@@ -1,6 +1,6 @@
 import pytest
 from .service_config import ServiceConfig
-from aikido_zen.helpers.ip_matcher import IPMatcher
+from aikido_zen.helpers.iplist import IPList
 
 
 def test_service_config_initialization():
@@ -66,10 +66,10 @@ def test_service_config_initialization():
     assert service_config.endpoints[1]["route"] == "/v3"
     assert service_config.endpoints[2]["route"] == "/admin"
     assert service_config.last_updated_at == last_updated_at
-    assert isinstance(service_config.bypassed_ips, IPMatcher)
-    assert service_config.bypassed_ips.has("127.0.0.1")
-    assert service_config.bypassed_ips.has("123.1.2.2")
-    assert not service_config.bypassed_ips.has("1.1.1.1")
+    assert isinstance(service_config.bypassed_ips, IPList)
+    assert service_config.bypassed_ips.matches("127.0.0.1")
+    assert service_config.bypassed_ips.matches("123.1.2.2")
+    assert not service_config.bypassed_ips.matches("1.1.1.1")
     assert service_config.blocked_uids == set(["1", "0", "5"])
 
     v1_endpoint = service_config.get_endpoints(
@@ -90,10 +90,10 @@ def test_service_config_initialization():
         }
     )[0]
     assert admin_endpoint["route"] == "/admin"
-    assert isinstance(admin_endpoint["allowedIPAddresses"], IPMatcher)
-    assert admin_endpoint["allowedIPAddresses"].has("192.168.2.1")
-    assert admin_endpoint["allowedIPAddresses"].has("1.2.3.4")
-    assert not admin_endpoint["allowedIPAddresses"].has("192.168.0.1")
+    assert isinstance(admin_endpoint["allowedIPAddresses"], IPList)
+    assert admin_endpoint["allowedIPAddresses"].matches("192.168.2.1")
+    assert admin_endpoint["allowedIPAddresses"].matches("1.2.3.4")
+    assert not admin_endpoint["allowedIPAddresses"].matches("192.168.0.1")
 
 
 # Sample data for testing
@@ -118,7 +118,7 @@ def service_config():
 def test_initialization(service_config):
     assert len(service_config.endpoints) == 2  # Only non-graphql endpoints
     assert service_config.last_updated_at == "2023-10-01T00:00:00Z"
-    assert isinstance(service_config.bypassed_ips, IPMatcher)
+    assert isinstance(service_config.bypassed_ips, IPList)
     assert service_config.blocked_uids == {"user1", "user2"}
 
 
@@ -194,10 +194,10 @@ def test_service_config_with_empty_allowlist():
     assert len(service_config.endpoints) == 1
     assert service_config.endpoints[0]["route"] == "/admin"
     assert service_config.last_updated_at == last_updated_at
-    assert isinstance(service_config.bypassed_ips, IPMatcher)
-    assert service_config.bypassed_ips.has("127.0.0.1")
-    assert service_config.bypassed_ips.has("123.1.2.2")
-    assert not service_config.bypassed_ips.has("1.1.1.1")
+    assert isinstance(service_config.bypassed_ips, IPList)
+    assert service_config.bypassed_ips.matches("127.0.0.1")
+    assert service_config.bypassed_ips.matches("123.1.2.2")
+    assert not service_config.bypassed_ips.matches("1.1.1.1")
     assert service_config.blocked_uids == set(["1", "0", "5"])
 
     admin_endpoint = service_config.get_endpoints(

aikido_zen/helpers/iplist.py

@@ -0,0 +1,71 @@
+"""Exports blocklist class and get_ip_address_type"""
+
+import ipaddress
+
+
+def get_ip_address_type(ip):
+    """Returns type, ipv4 or ipv6"""
+    try:
+        if ipaddress.ip_address(ip).version == 4:
+            return "ipv4"
+        elif ipaddress.ip_address(ip).version == 6:
+            return "ipv6"
+    except ValueError:
+        return None
+
+
+class IPList:
+    """A blocklist, where you can add subnets and addresses"""
+
+    def __init__(self, ip_list=None):
+        """Initializes the IPList, optionally with a list of IPs or CIDRs"""
+        self.blocked_addresses = set()
+        self.blocked_subnets = []
+
+        if ip_list:
+            for ip_or_cidr in ip_list:
+                self.add(ip_or_cidr)
+
+    def add_address(self, ip: str, ip_type: str):
+        self.blocked_addresses.add((ip, ip_type))
+
+    def add_subnet(self, plain_ip: str, ip_range: int, ip_type: str):
+        self.blocked_subnets.append((plain_ip, ip_range, ip_type))
+
+    def add(self, ip_or_cidr):
+        """
+        Checks whether ip_or_cidr is an IP address or is a subnet, and decides the correct IP type (IPv4 or IPv6)
+        """
+        if "/" not in ip_or_cidr:  # IP Address
+            ip_type = get_ip_address_type(ip_or_cidr)
+            if ip_type:
+                self.add_address(ip_or_cidr, ip_type)
+        else:  # Subnet
+            plain_ip, range_str = ip_or_cidr.split("/")
+            try:
+                ip_range = int(range_str)
+            except ValueError:
+                return
+            ip_type = get_ip_address_type(plain_ip)
+            if not ip_type:
+                return
+
+            self.add_subnet(plain_ip, ip_range, ip_type)
+
+    def matches(self, ip: str) -> bool:
+        ip_type = get_ip_address_type(ip)
+        if not ip_type:
+            return False
+
+        # Check if the IP address is in the blocked addresses
+        if (ip, ip_type) in self.blocked_addresses:
+            return True
+
+        # Check if the IP address is in any of the blocked subnets
+        ip_addr = ipaddress.ip_address(ip)
+        for plain_ip, ip_range, _ in self.blocked_subnets:
+            subnet = ipaddress.ip_network(f"{plain_ip}/{ip_range}", strict=False)
+            if ip_addr in subnet:
+                return True
+
+        return False

aikido_zen/helpers/iplist_test.py

@@ -0,0 +1,162 @@
+import pytest
+from .iplist import IPList, get_ip_address_type
+
+
+def test_get_ip_address_type():
+    # Valid IPv4 addresses
+    assert get_ip_address_type("192.168.1.1") == "ipv4"
+    assert get_ip_address_type("255.255.255.255") == "ipv4"
+    assert get_ip_address_type("0.0.0.0") == "ipv4"
+
+    # Valid IPv6 addresses
+    assert get_ip_address_type("::1") == "ipv6"
+    assert get_ip_address_type("2001:0db8:85a3:0000:0000:8a2e:0370:7334") == "ipv6"
+    assert get_ip_address_type("::") == "ipv6"
+
+    # Invalid IP addresses
+    assert get_ip_address_type("invalid_ip") is None
+    assert get_ip_address_type("") is None
+    assert get_ip_address_type("192.168.1.256") is None  # Invalid octet
+    assert (
+        get_ip_address_type("2001:0db8:85a3:0000:0000:8a2e:0370:7334:1234") is None
+    )  # Too many segments
+
+
+def test_blocklist_add_address():
+    blocklist = IPList()
+    assert blocklist.matches("192.168.1.1") is False
+    blocklist.add_address("192.168.1.1", "ipv4")
+    assert blocklist.matches("192.168.1.1") is True
+
+    # Test adding the same address again
+    blocklist.add_address("192.168.1.1", "ipv4")
+    assert len(blocklist.blocked_addresses) == 1  # Should still be one entry
+
+
+def test_blocklist_add_subnet():
+    blocklist = IPList()
+    assert blocklist.matches("10.0.0.1") is False
+    blocklist.add_subnet("10.0.0.0", 8, "ipv4")
+    assert blocklist.matches("10.0.0.1") is True
+    assert blocklist.matches("10.1.1.1") is True
+    assert blocklist.matches("192.168.1.1") is False
+
+
+def test_blocklist_multiple_addresses_and_subnets():
+    blocklist = IPList()
+    blocklist.add_address("192.168.1.1", "ipv4")
+    blocklist.add_subnet("10.0.0.0", 8, "ipv4")
+
+    assert blocklist.matches("192.168.1.1") is True
+    assert blocklist.matches("10.0.0.1") is True
+    assert blocklist.matches("10.1.1.1") is True
+    assert blocklist.matches("172.16.0.1") is False
+
+
+def test_blocklist_invalid_ip():
+    blocklist = IPList()
+    blocklist.add_address("192.168.1.1", "ipv4")
+    assert blocklist.matches("invalid_ip") is False
+
+
+def test_blocklist_subnet_edge_cases():
+    blocklist = IPList()
+    blocklist.add_subnet("192.168.1.0", 24, "ipv4")
+
+    assert blocklist.matches("192.168.1.255") is True  # Last address in the subnet
+    assert blocklist.matches("192.168.1.0") is True  # First address in the subnet
+    assert blocklist.matches("192.168.2.1") is False  # Outside the subnet
+    assert blocklist.matches("192.168.1.128") is True  # Middle of the subnet
+
+
+def test_blocklist_ipv6():
+    blocklist = IPList()
+    blocklist.add_address("2001:0db8:85a3:0000:0000:8a2e:0370:7334", "ipv6")
+    assert blocklist.matches("2001:0db8:85a3:0000:0000:8a2e:0370:7334") is True
+    assert blocklist.matches("::1") is False  # Not in the blocklist
+
+    blocklist.add_subnet("2001:0db8::", 32, "ipv6")
+    assert blocklist.matches("2001:0db8:abcd:0012:0000:0000:0000:0001") is True
+    assert blocklist.matches("2001:0db9::") is False  # Outside the subnet
+
+
+def test_blocklist_overlapping_subnets():
+    blocklist = IPList()
+    blocklist.add_subnet(
+        "192.168.1.0", 24, "ipv4"
+    )  # Covers 192.168.1.0 to 192.168.1.255
+    blocklist.add_subnet(
+        "192.168.1.128", 25, "ipv4"
+    )  # Covers 192.168.1.128 to 192.168.1.255
+
+    assert blocklist.matches("192.168.1.130") is True  # Inside both subnets
+    assert blocklist.matches("192.168.1.0") is True  # Inside first subnet
+    assert (
+        blocklist.matches("192.168.1.127") is True
+    )  # Inside first subnet, outside second
+    assert blocklist.matches("192.168.1.255") is True  # Last address in both subnets
+    assert blocklist.matches("192.168.2.1") is False  # Outside both subnets
+
+
+def test_blocklist_mixed_address_types():
+    blocklist = IPList()
+    blocklist.add_address("192.168.1.1", "ipv4")
+    blocklist.add_address("2001:0db8:85a3:0000:0000:8a2e:0370:7334", "ipv6")
+
+    assert blocklist.matches("192.168.1.1") is True
+    assert blocklist.matches("2001:0db8:85a3:0000:0000:8a2e:0370:7334") is True
+    assert blocklist.matches("192.168.1.2") is False  # Different IPv4
+    assert (
+        blocklist.matches("2001:0db8:85a3:0000:0000:8a2e:0370:7335") is False
+    )  # Different IPv6
+
+
+def test_blocklist_subnet_with_single_ip():
+    blocklist = IPList()
+    blocklist.add_subnet("192.168.1.1", 32, "ipv4")  # Single IP subnet
+
+    assert blocklist.matches("192.168.1.1") is True
+    assert blocklist.matches("192.168.1.2") is False  # Outside the subnet
+
+
+# Tests add function
+def test_valid_ips():
+    blocklist = IPList()
+
+    blocklist.add("1.2.3.4")
+    assert blocklist.matches("1.2.3.4") is True
+
+    blocklist.add("fd00:ec2::254")
+    assert blocklist.matches("fd00:ec2::254") is True
+
+    blocklist.add("192.168.2.1/24")
+    assert blocklist.matches("192.168.2.1") is True
+    assert blocklist.matches("192.168.2.240") is True
+
+    blocklist.add("fd00:124::1/64")
+    assert blocklist.matches("fd00:124::1") is True
+    assert blocklist.matches("fd00:124::f") is True
+    assert blocklist.matches("fd00:124::ff13") is True
+
+    blocklist.add("fd00:f123::1/128")
+    assert blocklist.matches("fd00:f123::1") is True
+
+    assert blocklist.matches("2.3.4.5") is False
+    assert blocklist.matches("fd00:125::ff13") is False
+    assert blocklist.matches("fd00:f123::2") is False
+
+
+def test_invalid_ips():
+    blocklist = IPList()
+
+    blocklist.add("192.168.2.2.1/24")
+    assert blocklist.matches("192.168.2.2.1") is False
+
+    blocklist.add("test")
+    assert blocklist.matches("test") is False
+
+    blocklist.add("")
+    assert blocklist.matches("") is False
+
+    blocklist.add("fd00:124::1/test")
+    assert blocklist.matches("fd00:124::1") is False

aikido_zen/helpers/is_ip_allowed_by_allowlist.py

@@ -13,7 +13,7 @@ def is_ip_allowed_by_allowlist(service_config, ip):
         return True
 
     for entry in service_config.allowed_ips:
-        if entry["iplist"].has(ip):
+        if entry["iplist"].matches(ip):
             # If the IP matches one of the lists the IP is allowed :
             return True
     return False