From 6530e51700accd43545a1c2115d187d7c78029a42866ea8ec332184b10efd9a0 Mon Sep 17 00:00:00 2001 From: larssand Date: Thu, 18 Jun 2026 23:03:21 +0200 Subject: [PATCH] add ip rep --- README.md | 12 +++++++ src/fgai/threat_intel.py | 63 ++++++++++++++++++++++++++++++++--- tests/test_recommendations.py | 16 +++++++++ 3 files changed, 86 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index bb91746..712b22e 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,15 @@ Optional external reputation enrichment is disabled by default. To use VirusTota ```bash export FGAI_THREAT_INTEL=1 +export ABUSEIPDB_API_KEY='...' +fgai recommend --logs logs/fg_syslog.jsonl --min-score 35 --threat-intel +``` + +VirusTotal is also supported: + +```bash +export FGAI_THREAT_INTEL=1 +export FGAI_THREAT_INTEL_PROVIDER=virustotal export VIRUSTOTAL_API_KEY='...' fgai recommend --logs logs/fg_syslog.jsonl --min-score 35 --threat-intel ``` @@ -174,6 +183,9 @@ end - `OLLAMA_MODEL`: defaults to `llama3.3`. - `OLLAMA_TIMEOUT`: Ollama request timeout in seconds, defaults to `180`. - `FGAI_THREAT_INTEL`: set to `1` to enable external threat intelligence lookups. +- `ABUSEIPDB_API_KEY`: AbuseIPDB API key for public IP reputation enrichment. +- `ABUSEIPDB_MAX_AGE_DAYS`: report age window for AbuseIPDB, defaults to `90`. +- `FGAI_THREAT_INTEL_PROVIDER`: `auto`, `abuseipdb`, or `virustotal`. - `VIRUSTOTAL_API_KEY`: VirusTotal API key for public IP reputation enrichment. ## Safety Model diff --git a/src/fgai/threat_intel.py b/src/fgai/threat_intel.py index a744d97..e2a0eae 100644 --- a/src/fgai/threat_intel.py +++ b/src/fgai/threat_intel.py @@ -6,6 +6,7 @@ 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: @@ -20,7 +21,10 @@ def is_public_ip(value: str | None) -> bool: 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.abuseipdb_key = os.getenv("ABUSEIPDB_API_KEY") self.virustotal_key = os.getenv("VIRUSTOTAL_API_KEY") + self.provider = os.getenv("FGAI_THREAT_INTEL_PROVIDER", "auto").lower() + self.max_age_days = int(os.getenv("ABUSEIPDB_MAX_AGE_DAYS", "90")) self.cache_path = Path(cache_file) self.ttl_seconds = ttl_seconds self.cache = self._read_cache() @@ -40,21 +44,70 @@ class ThreatIntelClient: 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) + provider = self._select_provider() + cache_key = f"{provider}:{ip}" + cached = self.cache.get(cache_key) 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} + 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} - result = self._lookup_virustotal_ip(ip) result["cached_at"] = now - self.cache[ip] = result + 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}", diff --git a/tests/test_recommendations.py b/tests/test_recommendations.py index acc8410..e8eda3e 100644 --- a/tests/test_recommendations.py +++ b/tests/test_recommendations.py @@ -33,6 +33,22 @@ class RecommendationTests(unittest.TestCase): self.assertEqual(result["status"], "disabled") + def test_abuseipdb_is_preferred_when_key_exists(self): + with patch.dict(os.environ, {"FGAI_THREAT_INTEL": "1", "ABUSEIPDB_API_KEY": "test"}, clear=True): + client = ThreatIntelClient(cache_file="/tmp/fgai-test-threat-cache.json") + + self.assertEqual(client._select_provider(), "abuseipdb") + + def test_virustotal_can_be_forced(self): + with patch.dict( + os.environ, + {"FGAI_THREAT_INTEL": "1", "ABUSEIPDB_API_KEY": "test", "VIRUSTOTAL_API_KEY": "test", "FGAI_THREAT_INTEL_PROVIDER": "virustotal"}, + clear=True, + ): + client = ThreatIntelClient(cache_file="/tmp/fgai-test-threat-cache.json") + + self.assertEqual(client._select_provider(), "virustotal") + if __name__ == "__main__": unittest.main()