352 lines
14 KiB
Python
352 lines
14 KiB
Python
"""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 hashlib
|
|
import json
|
|
import tempfile
|
|
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
|
|
|
|
|
|
class IPtoASNPrefixGeoCollector(BaseCollector):
|
|
name = "iptoasn_prefix_geo"
|
|
priority = "P1"
|
|
module = "L3"
|
|
frequency_hours = 24
|
|
data_type = "prefix_geography"
|
|
fail_on_empty = True
|
|
_cache_dir = Path(tempfile.gettempdir()) / "planet-download-cache" / "iptoasn"
|
|
|
|
@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]
|
|
|
|
@staticmethod
|
|
def _cache_key(url: str) -> str:
|
|
return hashlib.sha1(url.encode("utf-8")).hexdigest()[:16]
|
|
|
|
@classmethod
|
|
def _cache_paths(cls, url: str) -> tuple[Path, Path, Path]:
|
|
key = cls._cache_key(url)
|
|
gz_path = cls._cache_dir / f"{key}.tsv.gz"
|
|
part_path = cls._cache_dir / f"{key}.tsv.gz.part"
|
|
meta_path = cls._cache_dir / f"{key}.meta.json"
|
|
return gz_path, part_path, meta_path
|
|
|
|
@staticmethod
|
|
def _load_meta(meta_path: Path) -> dict[str, Any]:
|
|
if not meta_path.exists():
|
|
return {}
|
|
try:
|
|
return json.loads(meta_path.read_text(encoding="utf-8"))
|
|
except (json.JSONDecodeError, OSError):
|
|
return {}
|
|
|
|
@staticmethod
|
|
def _save_meta(meta_path: Path, payload: dict[str, Any]) -> None:
|
|
meta_path.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")
|
|
|
|
@staticmethod
|
|
def _validators_match(meta: dict[str, Any], remote: dict[str, Any]) -> bool:
|
|
etag = str(remote.get("etag") or "").strip()
|
|
last_modified = str(remote.get("last_modified") or "").strip()
|
|
if etag:
|
|
return etag == str(meta.get("etag") or "").strip()
|
|
if last_modified:
|
|
return last_modified == str(meta.get("last_modified") or "").strip()
|
|
return True
|
|
|
|
async def _fetch_remote_info(self, client: httpx.AsyncClient, url: str) -> dict[str, Any]:
|
|
try:
|
|
response = await client.head(url)
|
|
if response.status_code >= 400:
|
|
return {}
|
|
content_length_raw = response.headers.get("content-length")
|
|
content_length = int(content_length_raw) if content_length_raw else None
|
|
return {
|
|
"etag": response.headers.get("etag"),
|
|
"last_modified": response.headers.get("last-modified"),
|
|
"content_length": content_length,
|
|
"accept_ranges": (response.headers.get("accept-ranges") or "").lower(),
|
|
}
|
|
except (httpx.HTTPError, ValueError):
|
|
return {}
|
|
|
|
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 _download_dataset_with_resume(
|
|
self,
|
|
client: httpx.AsyncClient,
|
|
url: str,
|
|
*,
|
|
progress_callback=None,
|
|
) -> Path:
|
|
self._cache_dir.mkdir(parents=True, exist_ok=True)
|
|
gz_path, part_path, meta_path = self._cache_paths(url)
|
|
meta = self._load_meta(meta_path)
|
|
remote = await self._fetch_remote_info(client, url)
|
|
|
|
# Reuse existing complete file when validators and size still match.
|
|
if gz_path.exists():
|
|
local_size = gz_path.stat().st_size
|
|
remote_size = remote.get("content_length")
|
|
if (
|
|
self._validators_match(meta, remote)
|
|
and (remote_size is None or local_size == remote_size)
|
|
):
|
|
try:
|
|
self._parse_rows_from_gzip_file(gz_path)
|
|
if progress_callback and local_size > 0:
|
|
await progress_callback(local_size)
|
|
return gz_path
|
|
except Exception:
|
|
gz_path.unlink(missing_ok=True)
|
|
|
|
expected_size = remote.get("content_length")
|
|
can_resume = (remote.get("accept_ranges") or "") == "bytes"
|
|
resume_from = part_path.stat().st_size if part_path.exists() else 0
|
|
if expected_size is not None and resume_from > expected_size:
|
|
part_path.unlink(missing_ok=True)
|
|
resume_from = 0
|
|
if not self._validators_match(meta, remote):
|
|
part_path.unlink(missing_ok=True)
|
|
resume_from = 0
|
|
|
|
headers = {
|
|
"User-Agent": "Planet-Intelligence-System/1.0 (Python/collector)",
|
|
"Accept": "application/gzip,application/octet-stream,*/*",
|
|
}
|
|
if can_resume and resume_from > 0:
|
|
headers["Range"] = f"bytes={resume_from}-"
|
|
if remote.get("etag"):
|
|
headers["If-Range"] = str(remote.get("etag"))
|
|
elif remote.get("last_modified"):
|
|
headers["If-Range"] = str(remote.get("last_modified"))
|
|
|
|
async with client.stream("GET", url, headers=headers) as response:
|
|
response.raise_for_status()
|
|
|
|
if response.status_code == 206 and resume_from > 0:
|
|
mode = "ab"
|
|
else:
|
|
mode = "wb"
|
|
resume_from = 0
|
|
|
|
with part_path.open(mode) as f:
|
|
downloaded = resume_from
|
|
if progress_callback and downloaded > 0:
|
|
await progress_callback(downloaded)
|
|
async for chunk in response.aiter_bytes():
|
|
if chunk:
|
|
f.write(chunk)
|
|
downloaded += len(chunk)
|
|
if progress_callback:
|
|
await progress_callback(downloaded)
|
|
|
|
part_size = part_path.stat().st_size if part_path.exists() else 0
|
|
if expected_size is not None and part_size != expected_size:
|
|
raise RuntimeError(
|
|
f"IPtoASN download incomplete for {url}: expected={expected_size}, got={part_size}"
|
|
)
|
|
|
|
part_path.replace(gz_path)
|
|
self._save_meta(
|
|
meta_path,
|
|
{
|
|
"url": url,
|
|
"etag": remote.get("etag"),
|
|
"last_modified": remote.get("last_modified"),
|
|
"content_length": expected_size,
|
|
"updated_at": datetime.now(UTC).isoformat(),
|
|
},
|
|
)
|
|
return gz_path
|
|
|
|
async def _fetch_dataset_rows(
|
|
self,
|
|
client: httpx.AsyncClient,
|
|
url: str,
|
|
*,
|
|
progress_callback=None,
|
|
) -> list[dict[str, Any]]:
|
|
try:
|
|
gz_path = await self._download_dataset_with_resume(
|
|
client,
|
|
url,
|
|
progress_callback=progress_callback,
|
|
)
|
|
return self._parse_rows_from_gzip_file(gz_path)
|
|
except Exception:
|
|
# Bad partial/complete cache or interrupted stream: clear local files
|
|
# and retry once from scratch.
|
|
gz_path, part_path, meta_path = self._cache_paths(url)
|
|
gz_path.unlink(missing_ok=True)
|
|
part_path.unlink(missing_ok=True)
|
|
meta_path.unlink(missing_ok=True)
|
|
|
|
gz_path = await self._download_dataset_with_resume(
|
|
client,
|
|
url,
|
|
progress_callback=progress_callback,
|
|
)
|
|
return self._parse_rows_from_gzip_file(gz_path)
|
|
|
|
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._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
|
|
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) -> 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
|
|
await self.update_progress(min(aggregated, total_expected), commit=True)
|
|
|
|
batches = await asyncio.gather(
|
|
*(
|
|
self._fetch_dataset_rows(
|
|
client,
|
|
url,
|
|
progress_callback=lambda downloaded, u=url: on_url_progress(u, downloaded),
|
|
)
|
|
for url in dataset_urls
|
|
)
|
|
)
|
|
if total_expected > 0:
|
|
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
|