""" 离线 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()