Refactor is_local_ip function

osuTitanic · Jan 3, 2024 · 4b9bd3d · 4b9bd3d
1 parent a1b9896
commit 4b9bd3d
Showing 1 changed file with 20 additions and 19 deletions.
diff --git a/helpers/external/location.py b/helpers/external/location.py
@@ -9,9 +9,8 @@
 from dataclasses import dataclass
 from functools import cache
 
+import ipaddress
 import config
-import struct
-import socket
 import app
 
 @dataclass(slots=True)
@@ -39,27 +38,29 @@ def download_database():
         f.write(response.content)
 
 def is_local_ip(ip: str) -> bool:
-    if ':' in ip:
-        # TODO: IPv6 parsing
-        return False
+    address = ipaddress.ip_address(ip)
 
-    private = (
-        [ 2130706432, 4278190080 ], # 127.0.0.0
-        [ 3232235520, 4294901760 ], # 192.168.0.0
-        [ 2886729728, 4293918720 ], # 172.16.0.0
-        [ 167772160,  4278190080 ], # 10.0.0.0
-    )
+    if address.version == 6:
+        private = [
+            ipaddress.IPv6Network('fc00::/7'),
+            ipaddress.IPv6Network('::1/128')
+        ]
 
-    f = struct.unpack(
-        '!I',
-        socket.inet_pton(
-            socket.AF_INET,
-            ip
-        )
-    )[0]
+        for net in private:
+            if address in net:
+                return True
+
+        return False
+
+    private = [
+        ipaddress.IPv4Network('127.0.0.0/8'),
+        ipaddress.IPv4Network('192.168.0.0/16'),
+        ipaddress.IPv4Network('172.16.0.0/12'),
+        ipaddress.IPv4Network('10.0.0.0/8')
+    ]
 
     for net in private:
-        if (f & net[1]) == net[0]:
+        if address in net:
             return True
 
     return False