links/nginxmon/geo.py

"""
Geo-IP enrichment using a local DB-IP City Lite MMDB database.

The MMDB file is downloaded from jsDelivr on first use and refreshed
monthly via an APScheduler job.  All lookups are in-process (no network
round-trips after the initial download), so there are no timeouts or
rate-limit concerns.  Results are still cached in IPGeoCache so the
routermon receiver and other callers only do MMDB lookups for IPs that
aren't already in the database.
"""
import gzip
import ipaddress
import logging
import os
import shutil
import threading

import requests

from .models import IPGeoCache, NginxAccessLog

logger = logging.getLogger(__name__)

# DB-IP City Lite MMDB — CC BY 4.0, ~19 MB compressed, ~125 MB on disk.
_DB_URL = 'https://cdn.jsdelivr.net/npm/dbip-city-lite/dbip-city-lite.mmdb.gz'

_PRIVATE_RANGES = (
    ipaddress.ip_network('10.0.0.0/8'),
    ipaddress.ip_network('172.16.0.0/12'),
    ipaddress.ip_network('192.168.0.0/16'),
    ipaddress.ip_network('127.0.0.0/8'),
    ipaddress.ip_network('::1/128'),
    ipaddress.ip_network('fc00::/7'),
)

# Module-level reader cache — maxminddb readers are thread-safe for reads.
_reader      = None
_reader_lock = threading.Lock()


def _db_path() -> str:
    from django.conf import settings
    return os.path.join(str(settings.BASE_DIR), 'data', 'dbip-city-lite.mmdb')


def _get_reader():
    """Return the cached MMDB reader, downloading the DB if not yet present."""
    global _reader
    if _reader is not None:
        return _reader
    with _reader_lock:
        if _reader is not None:
            return _reader
        path = _db_path()
        if not os.path.exists(path):
            logger.info('nginxmon: geo MMDB not found — downloading now')
            try:
                _download(path)
            except Exception as exc:
                logger.warning('nginxmon: geo MMDB download failed: %s', exc)
                return None
        try:
            import maxminddb
            _reader = maxminddb.open_database(path)
            logger.info('nginxmon: geo MMDB loaded (%s)', path)
        except Exception as exc:
            logger.warning('nginxmon: failed to open geo MMDB: %s', exc)
        return _reader


def _reload_reader():
    """Invalidate the cached reader so it is reopened on next use."""
    global _reader
    with _reader_lock:
        if _reader is not None:
            try:
                _reader.close()
            except Exception:
                pass
            _reader = None


def _download(dest: str):
    """Download and decompress the MMDB gzip from jsDelivr."""
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    tmp = dest + '.tmp'
    logger.info('nginxmon: fetching %s', _DB_URL)
    with requests.get(_DB_URL, stream=True, timeout=120) as resp:
        resp.raise_for_status()
        resp.raw.decode_content = False  # keep raw compressed bytes
        with gzip.GzipFile(fileobj=resp.raw) as gz, open(tmp, 'wb') as out:
            shutil.copyfileobj(gz, out)
    os.replace(tmp, dest)
    logger.info('nginxmon: geo MMDB saved to %s', dest)


def download_geo_db():
    """Public entry point: download/refresh the MMDB and reload the reader."""
    _download(_db_path())
    _reload_reader()


def _is_private(ip: str) -> bool:
    try:
        addr = ipaddress.ip_address(ip)
        return any(addr in net for net in _PRIVATE_RANGES)
    except ValueError:
        return False


def _lookup_ip(ip: str) -> dict | None:
    """Return a geo dict for a single IP, or None if unavailable."""
    reader = _get_reader()
    if reader is None:
        return None
    try:
        rec = reader.get(ip)
        if rec is None:
            return None
        country      = rec.get('country') or {}
        city_rec     = rec.get('city') or {}
        subdivisions = rec.get('subdivisions') or []
        location     = rec.get('location') or {}
        return {
            'country':     (country.get('names') or {}).get('en', ''),
            'countryCode': country.get('iso_code', ''),
            'regionName':  ((subdivisions[0].get('names') or {}).get('en', '')
                            if subdivisions else ''),
            'city':        (city_rec.get('names') or {}).get('en', ''),
            'lat':         location.get('latitude'),
            'lon':         location.get('longitude'),
            'isp':         '',
            'status':      'success',
            'query':       ip,
        }
    except Exception as exc:
        logger.debug('nginxmon: MMDB lookup error for %s: %s', ip, exc)
        return None


def _lookup_batch(ips: list[str]) -> dict[str, dict]:
    """Look up IPs in the local MMDB. Returns {ip: geo_dict} (same interface as before)."""
    results = {}
    for ip in ips:
        data = _lookup_ip(ip)
        if data:
            results[ip] = data
    return results


def enrich_geo_batch(log_objs: list[NginxAccessLog]):
    """
    Look up geo data for IPs in `log_objs`, update cache and log records.
    Uses the local cache first; only queries the API for unknown IPs.
    """
    unique_ips = {obj.remote_addr for obj in log_objs}

    # Split private vs public
    private_ips = {ip for ip in unique_ips if _is_private(ip)}
    public_ips = unique_ips - private_ips

    # Ensure private IPs are in cache
    for ip in private_ips:
        IPGeoCache.objects.get_or_create(ip=ip, defaults={'is_private': True})

    # Load cached public IPs
    cached = {
        g.ip: g
        for g in IPGeoCache.objects.filter(ip__in=public_ips)
    }
    missing = [ip for ip in public_ips if ip not in cached]

    # Fetch missing from API in batches of 100
    api_results: dict[str, dict] = {}
    for i in range(0, len(missing), 100):
        api_results.update(_lookup_batch(missing[i:i + 100]))

    # Upsert cache
    for ip, data in api_results.items():
        obj, _ = IPGeoCache.objects.get_or_create(ip=ip)
        obj.country = data.get('country', '')
        obj.country_code = data.get('countryCode', '')
        obj.region = data.get('regionName', '')
        obj.city = data.get('city', '')
        obj.lat = data.get('lat')
        obj.lon = data.get('lon')
        obj.isp = data.get('isp', '')
        obj.save()
        cached[ip] = obj

    # Apply geo back via per-IP queryset update (avoids needing PKs on in-memory objects)
    for ip, geo in cached.items():
        NginxAccessLog.objects.filter(
            remote_addr=ip, country=''
        ).update(
            country=geo.country or '',
            country_code=geo.country_code or '',
            region=geo.region or '',
            city=geo.city or '',
        )