update_geoip.py (3772B)
1 #!/usr/bin/env python3 2 """Build the bundled IP-to-country database from PDDL user-country CSVs.""" 3 4 from __future__ import annotations 5 6 import argparse 7 import csv 8 import ipaddress 9 import struct 10 import tempfile 11 import time 12 import urllib.request 13 from pathlib import Path 14 15 16 IPV4_URL = "https://github.com/sapics/ip-location-db/releases/download/latest/user-country-ipv4-cidr.csv" 17 IPV6_URL = "https://github.com/sapics/ip-location-db/releases/download/latest/user-country-ipv6-cidr.csv" 18 MAGIC = b"IUNAGEO2" 19 20 21 def download(url: str, destination: Path) -> None: 22 request = urllib.request.Request(url, headers={"User-Agent": "iuna-geoip-updater/1"}) 23 with urllib.request.urlopen(request, timeout=60) as response: 24 destination.write_bytes(response.read()) 25 26 27 def records_from_csv(path: Path, version: int) -> list[tuple[int, int, bytes, bytes]]: 28 records: list[tuple[int, int, bytes, bytes]] = [] 29 with path.open(newline="", encoding="ascii") as source: 30 for line_number, row in enumerate(csv.reader(source), 1): 31 if len(row) != 2: 32 raise ValueError(f"{path}:{line_number}: expected CIDR,country") 33 cidr, country_text = row 34 country = country_text.upper().encode("ascii") 35 if len(country) != 2 or not country.isalpha(): 36 raise ValueError(f"{path}:{line_number}: invalid country code") 37 network = ipaddress.ip_network(cidr, strict=True) 38 if network.version != version: 39 raise ValueError(f"{path}:{line_number}: unexpected IP version") 40 records.append((version, network.prefixlen, country, network.network_address.packed)) 41 return records 42 43 44 def write_database(records: list[tuple[int, int, bytes, bytes]], output: Path) -> None: 45 groups: dict[tuple[int, int], dict[bytes, bytes]] = {} 46 for version, prefix, country, address in records: 47 group = groups.setdefault((version, prefix), {}) 48 previous = group.setdefault(address, country) 49 if previous != country: 50 raise ValueError(f"conflicting countries for {address.hex()}/{prefix}") 51 52 output.parent.mkdir(parents=True, exist_ok=True) 53 with output.open("wb") as database: 54 database.write(MAGIC) 55 for version, max_prefix in ((4, 32), (6, 128)): 56 for prefix in range(max_prefix + 1): 57 database.write(struct.pack("<I", len(groups.get((version, prefix), ())))) 58 for version, max_prefix in ((4, 32), (6, 128)): 59 for prefix in range(max_prefix + 1): 60 for address, country in sorted(groups.get((version, prefix), {}).items()): 61 database.write(address) 62 database.write(country) 63 64 65 def main() -> None: 66 parser = argparse.ArgumentParser(description=__doc__) 67 parser.add_argument("--ipv4-csv", type=Path) 68 parser.add_argument("--ipv6-csv", type=Path) 69 parser.add_argument( 70 "--output", 71 type=Path, 72 default=Path("src/ip_geolocation/embedded/ip-country.bin"), 73 ) 74 args = parser.parse_args() 75 76 with tempfile.TemporaryDirectory(prefix="iuna-geoip-") as temporary_directory: 77 temporary = Path(temporary_directory) 78 ipv4_csv = args.ipv4_csv or temporary / "user-country-ipv4-cidr.csv" 79 ipv6_csv = args.ipv6_csv or temporary / "user-country-ipv6-cidr.csv" 80 if args.ipv4_csv is None: 81 download(IPV4_URL, ipv4_csv) 82 time.sleep(1) 83 if args.ipv6_csv is None: 84 download(IPV6_URL, ipv6_csv) 85 86 records = records_from_csv(ipv4_csv, 4) 87 records.extend(records_from_csv(ipv6_csv, 6)) 88 write_database(records, args.output) 89 print(f"wrote {len(records)} prefixes to {args.output}") 90 91 92 if __name__ == "__main__": 93 main()