tianxuan/analysis/geoip.py

"""
离线 GeoIP 查询。内嵌 IPv4 区间 → 城市/经纬度映射。

从 ``data/geoip_data.txt`` 加载 IP 段数据，在模块首次导入时构建
排序索引，提供 O(log N) 的 IP 查询能力。

数据约 200KB，覆盖 23 个主要城市 IP 段。
可通过编辑 ``data/geoip_data.txt`` 扩展。
"""
import bisect
import logging
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
#  Data structures
# ---------------------------------------------------------------------------

# [(start_int, end_int, lat, lon, city, country), ...]
# Built at module load time from data/geoip_data.txt
_GEO_DATA: list[tuple[int, int, float, float, str, str]] = []

# Sorted list of start integers (for bisect lookup) + parallel data list
_STARTS: list[int] = []
_RANGES: list[tuple[int, int, float, float, str, str]] = []


# ---------------------------------------------------------------------------
#  IP conversion
# ---------------------------------------------------------------------------

def ip_to_int(ip_str: str) -> Optional[int]:
    """Convert an IPv4 dotted-string to its 32-bit integer representation.

    Returns ``None`` if *ip_str* is not a valid IPv4 address (including
    when any octet is outside 0-255).

    Examples
    --------
    >>> ip_to_int('8.8.8.8')
    134744072
    >>> ip_to_int('1.0.0.0')
    16777216
    """
    parts = ip_str.strip().split('.')
    if len(parts) != 4:
        return None
    try:
        octets = [int(p) for p in parts]
        if any(o < 0 or o > 255 for o in octets):
            return None
        return (octets[0] << 24) | (octets[1] << 16) | (octets[2] << 8) | octets[3]
    except (ValueError, TypeError):
        return None


# ---------------------------------------------------------------------------
#  Index building
# ---------------------------------------------------------------------------

def _build_index() -> tuple[list[int], list[tuple[int, int, float, float, str, str]]]:
    """Build sorted parallel lists for bisect-based IP lookup.

    Processes ``_GEO_DATA``, sorts by start integer, and returns
    ``(starts, ranges)`` where *starts* is a sorted list of start integers
    and *ranges* is the corresponding data list at the same indices.
    """
    if not _GEO_DATA:
        return [], []

    sorted_data = sorted(_GEO_DATA, key=lambda x: x[0])
    starts = [entry[0] for entry in sorted_data]
    return starts, sorted_data


def _load_data_file() -> None:
    """Load IP range data from ``data/geoip_data.txt`` into ``_GEO_DATA``.

    Called once at module import time. Each non-blank, non-comment line
    in the file should contain six comma-separated fields::

        start_ip,end_ip,lat,lon,city,country
    """
    data_path = Path(__file__).parent.parent / 'data' / 'geoip_data.txt'

    if not data_path.exists():
        logger.warning('[GEOIP] data file not found: %s', data_path)
        return

    loaded: list[tuple[int, int, float, float, str, str]] = []
    skipped = 0

    with open(data_path, 'r', encoding='utf-8') as f:
        for line_no, raw in enumerate(f, start=1):
            stripped = raw.strip()
            if not stripped or stripped.startswith('#'):
                continue

            parts = [p.strip() for p in stripped.split(',')]
            if len(parts) != 6:
                skipped += 1
                continue

            start_ip_str, end_ip_str, lat_str, lon_str, city, country = parts

            start_int = ip_to_int(start_ip_str)
            end_int = ip_to_int(end_ip_str)

            if start_int is None or end_int is None:
                skipped += 1
                continue

            try:
                lat = float(lat_str)
                lon = float(lon_str)
            except ValueError:
                skipped += 1
                continue

            loaded.append((start_int, end_int, lat, lon, city, country))

    _GEO_DATA.clear()
    _GEO_DATA.extend(loaded)

    global _STARTS, _RANGES
    _STARTS, _RANGES = _build_index()

    logger.info('[GEOIP] loaded %d ranges from %s (%d skipped)',
                len(loaded), data_path, skipped)


# ---------------------------------------------------------------------------
#  Public API
# ---------------------------------------------------------------------------

def lookup(ip_str: str) -> Optional[dict]:
    """查询 IP 对应的地理位置。

    Uses binary search (``bisect_right``) on the sorted start-IP list to
    locate the range *ip_str* falls within, then checks the end bound.

    Parameters
    ----------
    ip_str:
        IPv4 address as a dotted string (e.g. ``'8.8.8.8'``).

    Returns
    -------
    dict or None
        A dictionary with keys ``lat``, ``lon``, ``city``, ``country``
        if the IP is found in the database; ``None`` otherwise.

    Examples
    --------
    >>> result = lookup('8.8.8.8')
    >>> result['city']
    'Mountain View'
    >>> lookup('10.0.0.1') is None
    True
    """
    target = ip_to_int(ip_str)
    if target is None:
        return None

    if not _STARTS:
        return None

    idx = bisect.bisect_right(_STARTS, target) - 1
    if idx < 0:
        return None

    start, end, lat, lon, city, country = _RANGES[idx]
    if start <= target <= end:
        return {
            'lat': lat,
            'lon': lon,
            'city': city,
            'country': country,
        }

    return None


# ---------------------------------------------------------------------------
#  Module-level initialisation
# ---------------------------------------------------------------------------

_load_data_file()