diff --git a/dzgui/config/ipdb.py b/dzgui/config/ipdb.py index a25fbbe..f853ec3 100644 --- a/dzgui/config/ipdb.py +++ b/dzgui/config/ipdb.py @@ -3,6 +3,9 @@ import logging import re import requests import shutil +import socket +import sqlite3 +import time from pathlib import Path @@ -22,15 +25,16 @@ def find_download_url(url: str) -> str | None: return str(urls[0]) -def get_ipdb(ips_path: Path) -> None: +def get_ipdb(db_path: Path) -> None: try: url = find_download_url(DB_IP) if url is None: return + ips_path = db_path.with_suffix(".csv") + month_file = db_path.parent / ".month" date = find_date(url) - month_file = ips_path.parent / ".month" - if ips_path.exists() and month_file.exists(): + if db_path.exists() and month_file.exists(): old_date = month_file.read_text().rstrip("\n") if old_date == date: logger.info(f"IP DB date matches: {date}") @@ -45,6 +49,9 @@ def get_ipdb(ips_path: Path) -> None: logger.info("Stripping IPv6 records") strip_ipv6(ips_path) + create_table(db_path, ips_path) + ips_path.unlink() + with open(month_file, "w") as f: f.write(date) logger.info(f"Wrote {date} to {month_file}") @@ -108,3 +115,95 @@ def strip_ipv6(path: Path) -> None: out.write(merged) path.unlink() alt_path.rename(path) + + +def ip_to_int(ip: str) -> int: + return int.from_bytes(socket.inet_aton(ip), "big") + + +def marshal_rows(ip_file: Path) -> list[tuple[int, int, str, str]]: + YIELD = 0.003 + SLEEP = 0.001 + with open(ip_file, "r") as f: + rows: list[tuple[int, int, str, str]] = [] + # NOTE: yield to GIL periodically to reduce load on spinner + last_yield = time.monotonic() + for row in f: + start, end, lat, lon = row.rstrip("\n").split(",") + int_start = ip_to_int(start) + int_end = ip_to_int(end) + rows.append((int_start, int_end, lat, lon)) + + now = time.monotonic() + if (now - last_yield) >= YIELD: + time.sleep(SLEEP) + last_yield = time.monotonic() + return rows + + +def create_table(db_file: Path, ip_file: Path) -> None: + create = """ + CREATE TABLE IF NOT EXISTS ips ( + start_range INTEGER PRIMARY KEY, + end_range INTEGER, + latitude REAL, + longitude REAL + ) WITHOUT ROWID; + """ + insert = """ + INSERT INTO ips + (start_range, end_range, latitude, longitude) + values(?, ?, ?, ?) + """ + + if db_file.exists(): + db_file.unlink() + + conn = sqlite3.connect(":memory:") + c = conn.cursor() + + try: + c.execute(create) + rows = marshal_rows(ip_file) + c.executemany(insert, rows) + logger.info(f"Inserted {len(rows)} IP records into database") + conn.commit() + with sqlite3.connect(db_file) as dest: + conn.backup(dest) + logger.info(f"Serialized IP records to database: '{db_file}'") + conn.close() + except Exception as e: + logger.critical(e) + + +def read_from_table(db_path: Path, ip: str) -> tuple[float, float]: + if not db_path.exists(): + raise FileNotFoundError(f"No file found at the location '{db_path}") + + conn = sqlite3.connect(db_path) + c = conn.cursor() + + int_ip = ip_to_int(ip) + + query = """ + SELECT end_range, latitude, longitude + FROM ips + WHERE start_range <= ? + ORDER BY start_range DESC + LIMIT 1 + """ + index = """ + CREATE INDEX IF NOT EXISTS start_index + ON ips(start_range); + """ + try: + c.execute(index) + conn.commit() + c.execute(query, (int_ip,)) + row = c.fetchone() + end, lat, lon = row + conn.close() + return lat, lon + except Exception as e: + logger.critical(e) + raise e