Optimized IP matching

filter_subdomains.py

@@ -15,14 +15,20 @@ import ipaddress
 
 # DomainRule = typing.Union[bool, typing.Dict[str, 'DomainRule']]
 DomainRule = typing.Union[bool, typing.Dict]
+# IpRule = typing.Union[bool, typing.Dict[int, 'DomainRule']]
+IpRule = typing.Union[bool, typing.Dict]
 
 RULES_DICT: DomainRule = dict()
-RULES_IP: typing.Set[ipaddress.IPv4Network] = set()
+RULES_IP_DICT: IpRule = dict()
+
+
+def get_bits(address: ipaddress.IPv4Address) -> typing.Iterator[int]:
+    for char in address.packed:
+        for i in range(7, -1, -1):
+            yield (char >> i) & 0b1
 
 
 def subdomain_matching(subdomain: str) -> bool:
-    if not RULES_DICT:
-        return False
     parts = subdomain.split('.')
     parts.reverse()
     dic = RULES_DICT
@@ -36,12 +42,16 @@ def subdomain_matching(subdomain: str) -> bool:
 
 
 def ip_matching(ip_str: str) -> bool:
-    if not RULES_IP:
-        return False
     ip = ipaddress.ip_address(ip_str)
-    for net in RULES_IP:
-        if ip in net:
-            return True
+    dic = RULES_IP_DICT
+    i = 0
+    for bit in get_bits(ip):
+        i += 1
+        if isinstance(dic, bool) or bit not in dic:
+            break
+        dic = dic[bit]
+    if isinstance(dic, bool):
+        return dic
     return False
 
 
@@ -78,9 +88,17 @@ def register_rule(subdomain: str) -> None:
 
 def register_rule_ip(network: str) -> None:
     net = ipaddress.ip_network(network)
-    RULES_IP.add(net)
-    # If RULES_IP start becoming bigger,
-    # we might implement a binary tree for performance
+    ip = net.network_address
+    dic = RULES_IP_DICT
+    last_bit = net.prefixlen - 1
+    for b, bit in enumerate(get_bits(ip)):
+        if isinstance(dic, bool):
+            return
+        if b == last_bit:
+            dic[bit] = True
+        else:
+            dic.setdefault(bit, dict())
+            dic = dic[bit]
 
 
 if __name__ == '__main__':