add ipv4-mapped ipv6 addresses to imds check

Author: bitterpanda63

aikido_zen/vulnerabilities/ssrf/imds.py

@@ -3,41 +3,24 @@
 is_imds_ip_address, is_trusted_hostname, resolves_to_imds_ip
 """
 
+from aikido_zen.helpers.ip_matcher import IPMatcher
+from aikido_zen.helpers.ip_matcher.map_ipv4_to_ipv6 import map_ipv4_to_ipv6
 
-class BlockList:
-    """A list of IP's that shouldn't be accessed"""
-
-    def __init__(self):
-        self.blocked_addresses = {"ipv4": set(), "ipv6": set()}
-
-    def add_address(self, address, address_type):
-        """Add an address to this list"""
-        if address_type in self.blocked_addresses:
-            self.blocked_addresses[address_type].add(address)
-
-    def check(self, address, address_type=None):
-        """Check if the IP is on the list"""
-        if address_type:
-            return address in self.blocked_addresses.get(address_type, set())
-        return any(
-            address in addresses for addresses in self.blocked_addresses.values()
-        )
-
-
-# Create an instance of BlockList
-imds_addresses = BlockList()
+imds_addresses = IPMatcher()
 
 # Block the IP addresses used by AWS EC2 instances for IMDS
-imds_addresses.add_address("169.254.169.254", "ipv4")
-imds_addresses.add_address("fd00:ec2::254", "ipv6")
+imds_addresses.add("169.254.169.254")
+imds_addresses.add("fd00:ec2::254")
+imds_addresses.add(map_ipv4_to_ipv6("169.254.169.254"))
 
 # Block the IP address used by Alibaba Cloud
-imds_addresses.add_address("100.100.100.200", "ipv4")
+imds_addresses.add("100.100.100.200")
+imds_addresses.add(map_ipv4_to_ipv6("100.100.100.200"))
 
 
 def is_imds_ip_address(ip):
     """Checks if the IP is an imds ip"""
-    return imds_addresses.check(ip) or imds_addresses.check(ip, "ipv6")
+    return imds_addresses.has(ip)
 
 
 # Trusted hostnames for Google Cloud

aikido_zen/vulnerabilities/ssrf/imds_test.py

@@ -6,12 +6,16 @@ def test_returns_true_for_imds_ip_addresses():
     assert is_imds_ip_address("169.254.169.254") is True
     assert is_imds_ip_address("fd00:ec2::254") is True
 
+
 def test_returns_false_for_non_imds_ip_addresses():
     assert is_imds_ip_address("1.2.3.4") is False
     assert is_imds_ip_address("example.com") is False
 
+
 def test_is_imds_ip_address_ipv6_mapped():
     assert is_imds_ip_address("::ffff:169.254.169.254") is True
+    assert is_imds_ip_address("::ffff:100.100.100.200") is True
+
 
 # --- Tests ---
 def test_trusted_hostname_returns_none():