dzgui/dzgui/config/ipdb.py

import gzip
import logging
import re
import requests
import shutil
import socket
import sqlite3
import time

from pathlib import Path

from dzgui.const.constants import APP_NAME, REQUEST_TIMEOUT
from dzgui.const.endpoints import DB_IP

logger = logging.getLogger(APP_NAME)


def find_download_url(url: str) -> str | None:
    res = requests.get(url, timeout=REQUEST_TIMEOUT)
    res.raise_for_status()
    reg = r'ref=[\'"]([^\'">]+.csv.gz)'
    urls = re.findall(reg, res.text)
    if len(urls) < 1:
        return None
    return str(urls[0])


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)
        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}")
                return

        # TODO: log additional output
        logger.info(f"Fetching IPDB for {date} from {url}")
        tmp = serialize(url)
        logger.info(f"Extracting {tmp}")
        unzip(tmp, ips_path)

        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}")
    except Exception as e:
        # NOTE: in the event of failure, geolocation calc is simply not performed
        logger.critical(e)
        return


def find_date(url: str) -> str:
    date = re.sub(r"(dbip-city-lite-)(.*)(.csv.gz)", r"\2", url)
    date = date.split("/")[-1]
    return date


def unzip(tmp: Path, ips_path: Path) -> None:
    with gzip.open(tmp, "rb") as f:
        with open(ips_path, "wb") as outfile:
            shutil.copyfileobj(f, outfile)
    tmp.unlink()


def serialize(url: str) -> Path:
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        filename = "ips.csv.gz"
        tmp = Path("/tmp") / filename
        with open(tmp, "wb") as f:
            shutil.copyfileobj(r.raw, f)
    return tmp


def strip_ipv6(path: Path) -> None:
    """Can be IO intensive and cause visual lag on UI frames
    running in the main thread even when run in its own thread;
    lines are batched into memory-manageable chunks to reduce
    disk writes. Raw file can be 8M+ records long, so it is
    not read into memory at once.

    Relative size is reduced by ~100MB by pruning unwanted columns
    """
    # NOTE: "^::," is the boundary line between IPv4 and IPv6
    # NOTE: deprecated regex matching (slower by 5s)
    # reg = r"^\d{1,3}\..*"
    alt_path = path.parent.joinpath("ips_stripped.csv")
    merged = ""
    its = 0
    with open(path, "r") as f, open(alt_path, "w") as out:
        for line in f:
            els = line.split(",")
            if "." not in els[0]:
                break
            final = ",".join([els[0], els[1], els[-2], els[-1]])
            merged += final
            its += 1
            if its == 500:
                out.write(merged)
                its = 0
                merged = ""
        if its > 0:
            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