ad VS

src/fgai/threat_intel.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import ipaddress
+import json
+import os
+import time
+from pathlib import Path
+from urllib import error, request
+
+
+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 = 86400) -> None:
+        self.enabled = os.getenv("FGAI_THREAT_INTEL", "").lower() in {"1", "true", "yes", "on"}
+        self.virustotal_key = os.getenv("VIRUSTOTAL_API_KEY")
+        self.cache_path = Path(cache_file)
+        self.ttl_seconds = ttl_seconds
+        self.cache = self._read_cache()
+
+    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 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}
+        cached = self.cache.get(ip)
+        now = int(time.time())
+        if cached and now - int(cached.get("cached_at", 0)) < self.ttl_seconds:
+            return cached
+        if not self.enabled:
+            return {"ip": ip, "provider": "none", "status": "disabled", "score": 0}
+        if not self.virustotal_key:
+            return {"ip": ip, "provider": "virustotal", "status": "missing_api_key", "score": 0}
+
+        result = self._lookup_virustotal_ip(ip)
+        result["cached_at"] = now
+        self.cache[ip] = result
+        self._write_cache()
+        return result
+
+    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) -> dict[str, dict[str, object]]:
+    client = ThreatIntelClient(cache_file=cache_file)
+    enriched: dict[str, dict[str, object]] = {}
+    for ip in ips[:limit]:
+        enriched[ip] = client.lookup_ip(ip)
+    return enriched