from __future__ import annotations import ipaddress import json import os import time from pathlib import Path from urllib import error, request from urllib.parse import urlencode def is_public_ip(value: str | None) -> bool: if not value: return False try: return ipaddress.ip_address(value).is_global except ValueError: return False class ThreatIntelClient: def __init__( self, *, cache_file: str = "state/threat-intel-cache.json", ttl_seconds: int | None = None, enabled: bool | None = None, provider: str | None = None, abuseipdb_key: str | None = None, virustotal_key: str | None = None, daily_limit: int | None = None, error_ttl_seconds: int | None = None, abuseipdb_max_age_days: int | None = None, ) -> None: self.enabled = os.getenv("FGAI_THREAT_INTEL", "").lower() in {"1", "true", "yes", "on"} if enabled is None else enabled self.abuseipdb_key = abuseipdb_key if abuseipdb_key is not None else os.getenv("ABUSEIPDB_API_KEY") self.virustotal_key = virustotal_key if virustotal_key is not None else os.getenv("VIRUSTOTAL_API_KEY") self.provider = (provider if provider is not None else os.getenv("FGAI_THREAT_INTEL_PROVIDER", "auto")).lower() self.max_age_days = abuseipdb_max_age_days if abuseipdb_max_age_days is not None else int(os.getenv("ABUSEIPDB_MAX_AGE_DAYS", "90")) self.cache_path = Path(cache_file) self.ttl_seconds = ttl_seconds if ttl_seconds is not None else int(os.getenv("FGAI_THREAT_INTEL_TTL_SECONDS", "604800")) self.error_ttl_seconds = error_ttl_seconds if error_ttl_seconds is not None else int(os.getenv("FGAI_THREAT_INTEL_ERROR_TTL_SECONDS", "3600")) self.daily_limit = daily_limit if daily_limit is not None else int(os.getenv("FGAI_THREAT_INTEL_DAILY_LIMIT", "100")) self.cache = self._read_cache() def status(self) -> dict[str, object]: provider = self._select_provider() has_key = bool(self.abuseipdb_key if provider == "abuseipdb" else self.virustotal_key) return { "enabled": self.enabled, "provider": provider, "configured": has_key, "cache_ttl_seconds": self.ttl_seconds, "daily_limit": self.daily_limit, "requests_today": self._requests_today(provider), } def _read_cache(self) -> dict[str, dict[str, object]]: if not self.cache_path.exists(): return {} try: return json.loads(self.cache_path.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): return {} def _write_cache(self) -> None: self.cache_path.parent.mkdir(parents=True, exist_ok=True) self.cache_path.write_text(json.dumps(self.cache, indent=2, sort_keys=True), encoding="utf-8") def _cache_fresh(self, cached: dict[str, object], now: int) -> bool: ttl = self.ttl_seconds if cached.get("status") == "ok" else self.error_ttl_seconds return now - int(cached.get("cached_at", 0)) < ttl def _today(self) -> str: return time.strftime("%Y-%m-%d", time.gmtime()) def _requests_today(self, provider: str) -> int: meta = self.cache.get("_meta", {}) if not isinstance(meta, dict): return 0 requests_by_day = meta.get("requests_by_day", {}) if not isinstance(requests_by_day, dict): return 0 provider_counts = requests_by_day.get(provider, {}) return int(provider_counts.get(self._today(), 0)) if isinstance(provider_counts, dict) else 0 def _record_request(self, provider: str) -> None: meta = self.cache.setdefault("_meta", {}) if not isinstance(meta, dict): meta = {} self.cache["_meta"] = meta requests_by_day = meta.setdefault("requests_by_day", {}) if not isinstance(requests_by_day, dict): requests_by_day = {} meta["requests_by_day"] = requests_by_day provider_counts = requests_by_day.setdefault(provider, {}) if not isinstance(provider_counts, dict): provider_counts = {} requests_by_day[provider] = provider_counts today = self._today() provider_counts[today] = int(provider_counts.get(today, 0)) + 1 def lookup_ip(self, ip: str) -> dict[str, object]: if not is_public_ip(ip): return {"ip": ip, "provider": "local", "status": "skipped", "reason": "not_public_ip", "score": 0} provider = self._select_provider() cache_key = f"{provider}:{ip}" cached = self.cache.get(cache_key) now = int(time.time()) if isinstance(cached, dict) and self._cache_fresh(cached, now): return cached if not self.enabled: return {"ip": ip, "provider": "none", "status": "disabled", "score": 0} if self.daily_limit > 0 and self._requests_today(provider) >= self.daily_limit: return { "ip": ip, "provider": provider, "status": "daily_limit_reached", "score": 0, "reason": "external lookup budget reached; cached results remain available", } if provider == "abuseipdb": if not self.abuseipdb_key: return {"ip": ip, "provider": "abuseipdb", "status": "missing_api_key", "score": 0} result = self._lookup_abuseipdb_ip(ip) elif provider == "virustotal": if not self.virustotal_key: return {"ip": ip, "provider": "virustotal", "status": "missing_api_key", "score": 0} result = self._lookup_virustotal_ip(ip) else: return {"ip": ip, "provider": provider, "status": "unsupported_provider", "score": 0} self._record_request(provider) result["cached_at"] = now self.cache[cache_key] = result self._write_cache() return result def _select_provider(self) -> str: if self.provider in {"abuseipdb", "virustotal"}: return self.provider if self.abuseipdb_key: return "abuseipdb" if self.virustotal_key: return "virustotal" return "abuseipdb" def _lookup_abuseipdb_ip(self, ip: str) -> dict[str, object]: query = urlencode({"ipAddress": ip, "maxAgeInDays": str(self.max_age_days)}) req = request.Request( f"https://api.abuseipdb.com/api/v2/check?{query}", headers={"Key": self.abuseipdb_key or "", "Accept": "application/json"}, ) try: with request.urlopen(req, timeout=20) as response: payload = json.loads(response.read().decode("utf-8")) except error.HTTPError as exc: return {"ip": ip, "provider": "abuseipdb", "status": f"http_{exc.code}", "score": 0} except Exception as exc: return {"ip": ip, "provider": "abuseipdb", "status": "error", "error": str(exc), "score": 0} data = payload.get("data", {}) score = int(data.get("abuseConfidenceScore", 0) or 0) return { "ip": ip, "provider": "abuseipdb", "status": "ok", "score": score, "abuse_confidence_score": score, "total_reports": int(data.get("totalReports", 0) or 0), "country_code": data.get("countryCode"), "usage_type": data.get("usageType"), "isp": data.get("isp"), "domain": data.get("domain"), "is_tor": bool(data.get("isTor", False)), "last_reported_at": data.get("lastReportedAt"), } def _lookup_virustotal_ip(self, ip: str) -> dict[str, object]: req = request.Request( f"https://www.virustotal.com/api/v3/ip_addresses/{ip}", headers={"x-apikey": self.virustotal_key or "", "accept": "application/json"}, ) try: with request.urlopen(req, timeout=20) as response: data = json.loads(response.read().decode("utf-8")) except error.HTTPError as exc: return {"ip": ip, "provider": "virustotal", "status": f"http_{exc.code}", "score": 0} except Exception as exc: return {"ip": ip, "provider": "virustotal", "status": "error", "error": str(exc), "score": 0} stats = data.get("data", {}).get("attributes", {}).get("last_analysis_stats", {}) malicious = int(stats.get("malicious", 0) or 0) suspicious = int(stats.get("suspicious", 0) or 0) harmless = int(stats.get("harmless", 0) or 0) undetected = int(stats.get("undetected", 0) or 0) score = min(100, malicious * 20 + suspicious * 10) return { "ip": ip, "provider": "virustotal", "status": "ok", "score": score, "malicious": malicious, "suspicious": suspicious, "harmless": harmless, "undetected": undetected, } def enrich_ips( ips: list[str], *, cache_file: str = "state/threat-intel-cache.json", limit: int = 25, enabled: bool | None = None, config: dict[str, object] | None = None ) -> dict[str, dict[str, object]]: config = config or {} client = ThreatIntelClient( cache_file=cache_file, enabled=enabled, provider=str(config.get("threat_intel_provider", "auto")), abuseipdb_key=str(config.get("abuseipdb_api_key", "") or "") or None, virustotal_key=str(config.get("virustotal_api_key", "") or "") or None, daily_limit=int(config.get("threat_intel_daily_limit", 100) or 100), ttl_seconds=int(config.get("threat_intel_ttl_seconds", 604800) or 604800), error_ttl_seconds=int(config.get("threat_intel_error_ttl_seconds", 3600) or 3600), abuseipdb_max_age_days=int(config.get("abuseipdb_max_age_days", 90) or 90), ) enriched: dict[str, dict[str, object]] = {} for ip in ips[:limit]: enriched[ip] = client.lookup_ip(ip) return enriched