""" 离线 GeoIP 查询 — 支持 MMDB (GeoLite2-City / GeoLite2-ASN) 和文本格式回退。 优先级： 1. GeoLite2-City.mmdb (MaxMind MMDB 格式，城市级精度) 2. GeoLite2-ASN.mmdb (可选，提供 ISP / 组织信息) 3. geoip_data.txt (文本格式回退，803 段覆盖 30 城市) MMDB 查询使用 ``maxminddb`` 库，文本格式使用二分查找 O(log N)。 """ import bisect import logging from pathlib import Path from typing import Optional logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Data structures # --------------------------------------------------------------------------- # MMDB readers (lazy-loaded, None if unavailable) _mmdb_city: Optional[object] = None # maxminddb.Reader for City DB _mmdb_asn: Optional[object] = None # maxminddb.Reader for ASN DB # Text-format data (fallback) # [(start_int, end_int, lat, lon, city, country), ...] _GEO_DATA: list[tuple[int, int, float, float, str, str]] = [] _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. """ 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 # --------------------------------------------------------------------------- # MMDB loading # --------------------------------------------------------------------------- def _load_mmdb() -> None: """Try to load MaxMind MMDB databases from ``data/`` directory. Looks for: - ``data/GeoLite2-City.mmdb`` (primary — city, country, lat/lon) - ``data/GeoLite2-ASN.mmdb`` (optional — ISP, organization, ASN) """ global _mmdb_city, _mmdb_asn data_dir = Path(__file__).parent.parent / 'data' city_path = data_dir / 'GeoLite2-City.mmdb' asn_path = data_dir / 'GeoLite2-ASN.mmdb' if not city_path.exists(): logger.info('[GEOIP] GeoLite2-City.mmdb not found at %s — will use text fallback', city_path) return try: import maxminddb _mmdb_city = maxminddb.open_database(str(city_path)) logger.info('[GEOIP] loaded GeoLite2-City.mmdb (MMDB format)') if asn_path.exists(): _mmdb_asn = maxminddb.open_database(str(asn_path)) logger.info('[GEOIP] loaded GeoLite2-ASN.mmdb (ISP/ORG)') except Exception as e: logger.warning('[GEOIP] failed to open MMDB: %s', e) _mmdb_city = None _mmdb_asn = None # --------------------------------------------------------------------------- # Text-format loading (fallback) # --------------------------------------------------------------------------- def _build_index() -> tuple[list[int], list[tuple[int, int, float, float, str, str]]]: 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`` (text fallback).""" data_path = Path(__file__).parent.parent / 'data' / 'geoip_data.txt' if not data_path.exists(): logger.warning('[GEOIP] text 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 raw in f: 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)) global _GEO_DATA, _STARTS, _RANGES _GEO_DATA.clear() _GEO_DATA.extend(loaded) _STARTS, _RANGES = _build_index() logger.info('[GEOIP] text fallback: loaded %d ranges from %s (%d skipped)', len(loaded), data_path, skipped) # --------------------------------------------------------------------------- # MMDB lookup helpers # --------------------------------------------------------------------------- def _lookup_mmdb(ip_str: str) -> Optional[dict]: """Query MMDB database for an IP. Returns rich result with ISP/ORG when ASN database is also available.""" if _mmdb_city is None: return None try: city_result = _mmdb_city.get(ip_str) if city_result is None: return None # Extract location location = city_result.get('location', {}) lat = location.get('latitude') lon = location.get('longitude') if lat is None or lon is None: return None # Extract names (English preferred) city = None city_names = city_result.get('city', {}).get('names', {}) if city_names: city = city_names.get('en') or next(iter(city_names.values()), None) country = None country_info = city_result.get('country', {}) country = country_info.get('iso_code') or '' country_names = country_info.get('names', {}) country_name = country_names.get('en') or next(iter(country_names.values()), None) if country_names else None # Subdivisions subdivisions = city_result.get('subdivisions', []) region = None if subdivisions: sub_names = subdivisions[0].get('names', {}) region = sub_names.get('en') or next(iter(sub_names.values()), None) if sub_names else None result = { 'lat': round(float(lat), 4), 'lon': round(float(lon), 4), 'city': city or region or '', 'country': country, 'country_name': country_name or '', 'region': region or '', } # Try ASN database for ISP/ORG if _mmdb_asn is not None: try: asn_result = _mmdb_asn.get(ip_str) if asn_result: result['isp'] = asn_result.get('autonomous_system_organization', '') result['asn'] = asn_result.get('autonomous_system_number', '') except Exception: pass return result except Exception as e: logger.debug('[GEOIP] MMDB lookup error for %s: %s', ip_str, e) return None def _lookup_text(ip_str: str) -> Optional[dict]: """Query text-format database (fallback).""" 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, 'country_name': '', 'region': '', } return None # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- def lookup(ip_str: str) -> Optional[dict]: """查询 IP 对应的地理位置。 Tries MMDB first (if available), then falls back to text format. Parameters ---------- ip_str: IPv4 address as a dotted string (e.g. ``'8.8.8.8'``). Returns ------- dict or None Keys: ``lat``, ``lon``, ``city``, ``country`` (ISO code), ``country_name``, ``region``. When ASN database is available, also includes: ``isp``, ``asn``. """ if not ip_str or ip_str in ('None', '', '+', '-'): return None # MMDB first if _mmdb_city is not None: result = _lookup_mmdb(ip_str) if result is not None: return result # MMDB miss — fall through to text # Text fallback return _lookup_text(ip_str) def is_available() -> bool: """Return True if at least one GeoIP source is loaded.""" return _mmdb_city is not None or bool(_STARTS) def source_name() -> str: """Return human-readable name of the active GeoIP source.""" if _mmdb_city is not None: name = 'GeoLite2-City.mmdb' if _mmdb_asn is not None: name += ' + GeoLite2-ASN.mmdb' return name if _STARTS: return 'geoip_data.txt (text fallback)' return '(none)' # --------------------------------------------------------------------------- # Module-level initialisation # --------------------------------------------------------------------------- _load_mmdb() if _mmdb_city is None: _load_data_file() logger.info('[GEOIP] active source: %s', source_name())