chore: update tests

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

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,
     )

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

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)

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:

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] = []

tests/test_config.py

@@ -36,6 +36,7 @@ def keys():
 
 # TODO: use a static fixture instead of system config
 @pytest.fixture
+@pytest.mark.realconfig
 def config():
     paths = get_xdg_paths()
     xdg = parse_filepaths(paths)

tests/test_flatpak_override.py

@@ -1,7 +1,7 @@
 import configparser
 import pytest
 
-from dzgui.const.constants import FLATPAK_APPID, GTK_APP_ID, TMP_PATH
+from dzgui.const.constants import FLATPAK_APPID, GTK_APP_ID
 from dzgui.config.xdg import get_xdg_paths
 from dzgui.util.flatpak import apply_overrides
 
@@ -38,8 +38,6 @@ def test_session_bus(override, config) -> None:
 
 
 def test_paths(override, xdg_paths, config) -> None:
-    expect: list[str] = []
-    expect.append(TMP_PATH)
-    expect.extend([*xdg_paths.values()])
+    expect = xdg_paths
     paths = config[PATHS_KEY][PATHS_SUBKEY].split(";")
-    assert all(str(path) in paths for path in expect)
+    assert all(str(path) in paths for path in expect.values())

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