229 lines
9.9 KiB
Python
229 lines
9.9 KiB
Python
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
|