"""IPtoASN prefix geography collector. Downloads the public combined IPv4+IPv6 TSV database and stores coarse prefix-to-country/ASN geography hints for BGP enrichment. """ from __future__ import annotations import asyncio import gzip import time from datetime import UTC, datetime from ipaddress import summarize_address_range, ip_address from pathlib import Path from typing import Any import httpx from app.services.collectors.base import BaseCollector from app.services.collectors.downloads import ResumableFileDownloader class IPtoASNPrefixGeoCollector(BaseCollector): name = "iptoasn_prefix_geo" priority = "P1" module = "L3" frequency_hours = 24 data_type = "prefix_geography" fail_on_empty = True _downloader = ResumableFileDownloader( cache_namespace="iptoasn", default_accept="application/gzip,application/octet-stream,*/*", ) @staticmethod def _build_dataset_urls(resolved_url: str) -> list[str]: if "ip2asn-combined.tsv.gz" in resolved_url: return [ resolved_url.replace("ip2asn-combined.tsv.gz", "ip2asn-v4.tsv.gz"), resolved_url.replace("ip2asn-combined.tsv.gz", "ip2asn-v6.tsv.gz"), ] return [resolved_url] def _parse_rows_from_gzip_file(self, file_path: Path) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] with gzip.open(file_path, "rt", encoding="utf-8", errors="replace") as f: for raw_line in f: line = raw_line.strip() if not line or line.startswith("#"): continue parts = line.split("\t") if len(parts) < 5: continue range_start, range_end, asn, country_code, as_name = parts[:5] rows.append( { "range_start": range_start, "range_end": range_end, "asn": asn, "country_code": country_code, "as_name": as_name, } ) if not rows: raise RuntimeError(f"IPtoASN dataset parsed empty rows: {file_path.name}") return rows async def _fetch_dataset_rows( self, client: httpx.AsyncClient, url: str, *, progress_callback=None, ) -> list[dict[str, Any]]: file_path = await self._downloader.download_file( client, url, extension=".tsv.gz", progress_callback=progress_callback, validate_existing=lambda p: self._validate_gzip_dataset(p), ) return self._parse_rows_from_gzip_file(file_path) def _validate_gzip_dataset(self, file_path: Path) -> bool: try: self._parse_rows_from_gzip_file(file_path) return True except Exception: return False async def fetch(self) -> list[dict[str, Any]]: if not self._resolved_url: raise RuntimeError("IPtoASN combined URL is not configured") dataset_urls = self._build_dataset_urls(self._resolved_url) async with httpx.AsyncClient(timeout=180.0, follow_redirects=True) as client: remote_infos = await asyncio.gather( *(self._downloader.fetch_remote_info(client, url) for url in dataset_urls) ) expected_sizes = [ info.get("content_length") for info in remote_infos if isinstance(info.get("content_length"), int) ] total_expected = sum(expected_sizes) if expected_sizes else 0 if total_expected > 0 and self._current_task and self._db_session: self._current_task.total_records = total_expected self._current_task.records_processed = 0 self._current_task.progress = 0.0 self._current_task.phase_progress = 0.0 self._current_task.phase_message = "正在下载 IPtoASN 数据" self._current_task.phase_current = 0 self._current_task.phase_total = total_expected self._current_task.phase_unit = "bytes" await self._db_session.commit() await self._publish_task_update(force=True) url_progress: dict[str, int] = {url: 0 for url in dataset_urls} progress_lock = asyncio.Lock() last_emit = {"t": 0.0, "value": 0} min_emit_bytes = max(total_expected // 200, 2 * 1024 * 1024) if total_expected > 0 else 4 * 1024 * 1024 async def on_url_progress(url: str, downloaded_bytes: int, total_bytes: int | None) -> None: if total_expected <= 0: return async with progress_lock: current = max(0, downloaded_bytes) if current < url_progress[url]: return url_progress[url] = current aggregated = sum(url_progress.values()) now = time.monotonic() should_emit = ( aggregated >= total_expected or aggregated - last_emit["value"] >= min_emit_bytes or now - last_emit["t"] >= 2.0 ) if not should_emit: return last_emit["value"] = aggregated last_emit["t"] = now current = min(aggregated, total_expected) await self.update_phase_progress( current=current, total=total_expected, unit="bytes", message="正在下载 IPtoASN 数据", ) await self.update_progress(current, commit=True) batches = await asyncio.gather( *( self._fetch_dataset_rows( client, url, progress_callback=lambda downloaded, total, u=url: on_url_progress(u, downloaded, total), ) for url in dataset_urls ) ) if total_expected > 0: await self.update_phase_progress( current=total_expected, total=total_expected, unit="bytes", message="IPtoASN 数据下载完成", ) await self.update_progress(total_expected, commit=True, force=True) rows: list[dict[str, Any]] = [] for batch in batches: rows.extend(batch) return rows def transform(self, raw_data: list[dict[str, Any]]) -> list[dict[str, Any]]: reference_date = datetime.now(UTC).isoformat() transformed: list[dict[str, Any]] = [] for item in raw_data: try: start_ip = ip_address(str(item["range_start"])) end_ip = ip_address(str(item["range_end"])) except ValueError: continue if start_ip.version != end_ip.version: continue summarized = list(summarize_address_range(start_ip, end_ip)) primary_prefix = str(summarized[0]) if summarized else f"{start_ip}/{32 if start_ip.version == 4 else 128}" family = f"ipv{start_ip.version}" asn_value = item.get("asn") try: normalized_asn = int(str(asn_value)) except (TypeError, ValueError): normalized_asn = None transformed.append( { "source_id": f"{family}:{item['range_start']}-{item['range_end']}", "name": primary_prefix, "title": f"{primary_prefix} {item.get('country_code', '').strip()}".strip(), "country": item.get("country_code"), "city": "", "latitude": None, "longitude": None, "metadata": { "family": family, "range_start": item["range_start"], "range_end": item["range_end"], "prefix": primary_prefix, "prefixes": [str(prefix) for prefix in summarized[:8]], "range_prefix_count": len(summarized), "country_code": item.get("country_code"), "asn": normalized_asn, "as_name": item.get("as_name"), "source_dataset": "iptoasn_combined", }, "reference_date": reference_date, } ) return transformed