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 diff --git a/dzgui/config/xdg.py b/dzgui/config/xdg.py index 1d7f5b3..b72ea7b 100644 --- a/dzgui/config/xdg.py +++ b/dzgui/config/xdg.py @@ -22,7 +22,7 @@ class Xdg: version: Path system: Path debug: Path - ips: Path + db: Path filters: Path shortcut: Path @@ -88,7 +88,7 @@ def parse_filepaths(xdg: dict) -> Xdg: resolution = state / "dzg.res.json" history = state / "dzg.history" versions = state / "dzg.versions" - ips = state / "ips.csv" + db = state / "ips.db" filters = state / "dzg.filters.json" shortcut = share / APP_NAME_LOWER @@ -102,7 +102,7 @@ def parse_filepaths(xdg: dict) -> Xdg: versions, system, debug, - ips, + db, filters, shortcut, ) diff --git a/dzgui/init/coords.py b/dzgui/init/coords.py index 009b386..9381efb 100644 --- a/dzgui/init/coords.py +++ b/dzgui/init/coords.py @@ -10,10 +10,10 @@ if TYPE_CHECKING: logger = logging.getLogger(APP_NAME) -def get_local_coords(path: "Path") -> Coords | None: +def get_local_coords(db_path: "Path") -> Coords | None: try: my_ip = get_local_ip() - return get_coords(path, my_ip) + return get_coords(db_path, my_ip) except Exception as e: logger.warn(e) return None diff --git a/dzgui/util/dist.py b/dzgui/util/dist.py index 0bc6f2b..46c2d44 100644 --- a/dzgui/util/dist.py +++ b/dzgui/util/dist.py @@ -80,8 +80,9 @@ class CalcDist: if local is None: return None try: - remote = get_coords(prefs.paths.ips, remote_ip) - except Exception: + remote = get_coords(prefs.paths.db, remote_ip) + except Exception as e: + logger.critical(e) return None haversine = Haversine(local.lat, local.lon, remote.lat, remote.lon) diff --git a/dzgui/util/ip.py b/dzgui/util/ip.py index 2901986..d9a2c12 100644 --- a/dzgui/util/ip.py +++ b/dzgui/util/ip.py @@ -1,4 +1,5 @@ import a2s +from ipaddress import ip_address import requests import shutil import subprocess @@ -7,6 +8,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING from dzgui.api.servers import Record +from dzgui.config.ipdb import read_from_table from dzgui.const.constants import REQUEST_TIMEOUT from dzgui.const.endpoints import IP_ECHO @@ -33,54 +35,21 @@ def ip2list(ip: str) -> list: return [int(digit) for digit in ip.split(".")] -# TODO: tests -def is_reserved(digits: list) -> bool: - reserved = [0, 10, 127] - - if digits[0] in reserved: - return True - if digits[0] == 192 and digits[1] == 168: - return True - return False +def is_usable(ip: str) -> bool: + address = ip_address(ip) + if address.is_reserved or address.is_private: + return False + return True -def get_coords(ips: "Path", ip: str) -> Coords: - split = ip2list(ip) - - if is_reserved(split): +def get_coords(db_path: "Path", ip: str) -> Coords: + if not is_usable(ip): raise ReservedIPError(f"{ip} is a reserved IP address.") - - prefix = f"^{split[0]}.{split[1]}." - - if shutil.which("rg") is not None: - args = ["/usr/bin/rg", prefix, str(ips)] - else: - args = ["/usr/bin/grep", "-E", prefix, str(ips)] - proc = subprocess.run(args, capture_output=True, text=True) - - if proc.returncode != 0: - raise GeolocationError("Failed to split records") - - ip_list = proc.stdout.splitlines() - - for address in ip_list: - fields = address.split(",") - upper = fields[1] - - upper_digits = ip2list(upper) - if ip == upper: - return Coords(float(fields[-2]), float(fields[-1])) - - if split[2] > upper_digits[2]: - continue - - if split[2] == upper_digits[2]: - if split[3] > upper_digits[3]: - continue - - if split[2] <= upper_digits[2]: - return Coords(float(fields[-2]), float(fields[-1])) - raise GeolocationError("No matching records found") + try: + lat, lon = read_from_table(db_path, ip) + return Coords(lat, lon) + except Exception: + raise GeolocationError("No matching records found") def get_local_ip() -> str: diff --git a/dzgui/views/dialogs/boot.py b/dzgui/views/dialogs/boot.py index e16792c..d4aa557 100644 --- a/dzgui/views/dialogs/boot.py +++ b/dzgui/views/dialogs/boot.py @@ -132,9 +132,9 @@ class BootDialog(ColorAwareApp, Gtk.Dialog): # type: ignore preboot.signatures, False, ), - (StoredFunc(get_ipdb, self.xdg.ips), preboot.geo, False), + (StoredFunc(get_ipdb, self.xdg.db), preboot.geo, False), # TODO: handle IP DB failure and use coords fallback - (StoredFunc(get_local_coords, self.xdg.ips), preboot.coords, True), + (StoredFunc(get_local_coords, self.xdg.db), preboot.coords, True), (StoredFunc(check_updates, self.version), preboot.updates, True), ] self.results: list[Any] = [] diff --git a/tests/test_geo_file.py b/tests/test_geo_file.py index d142edc..17c195f 100644 --- a/tests/test_geo_file.py +++ b/tests/test_geo_file.py @@ -1,16 +1,40 @@ -from dzgui.config.ipdb import get_ipdb -from pathlib import Path import os +import sqlite3 import pytest import tempfile +from pathlib import Path + +from dzgui.const.endpoints import DB_IP +from dzgui.config.ipdb import ( + create_table, + find_download_url, + serialize, + strip_ipv6, + unzip, +) + + +def get_db_len(path: Path) -> int: + conn = sqlite3.connect(path) + c = conn.cursor() + c.execute("SELECT COUNT(*) from ips") + db_len = c.fetchone()[0] + return db_len + + +@pytest.fixture +def csv() -> Path: + csv = Path(tempfile.NamedTemporaryFile().name) + url = find_download_url(DB_IP) + tmp = serialize(url) + unzip(tmp, csv) + strip_ipv6(csv) + return csv @pytest.mark.slow -def test_stripping(): - d = tempfile.TemporaryDirectory() - path = Path(d.name).joinpath("ips.csv") - get_ipdb(path) - with open(path, "rb") as f: +def test_stripping(csv): + with open(csv, "rb") as f: first = f.readline().decode().rstrip() try: f.seek(-2, os.SEEK_END) @@ -25,3 +49,16 @@ def test_stripping(): assert first_els[0] == "0.0.0.0" assert last_els[0] == "224.0.0.0" + + +@pytest.mark.slow +def test_row_parity(csv) -> None: + rows = csv.read_text().splitlines() + + db = Path(tempfile.NamedTemporaryFile().name) + create_table(db, csv) + + csv_len = len(rows) + db_len = get_db_len(db) + + assert csv_len == db_len