add ip rep
This commit is contained in:
12
README.md
12
README.md
@@ -68,6 +68,15 @@ Optional external reputation enrichment is disabled by default. To use VirusTota
|
|||||||
|
|
||||||
```bash
|
```bash
|
||||||
export FGAI_THREAT_INTEL=1
|
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='...'
|
export VIRUSTOTAL_API_KEY='...'
|
||||||
fgai recommend --logs logs/fg_syslog.jsonl --min-score 35 --threat-intel
|
fgai recommend --logs logs/fg_syslog.jsonl --min-score 35 --threat-intel
|
||||||
```
|
```
|
||||||
@@ -174,6 +183,9 @@ end
|
|||||||
- `OLLAMA_MODEL`: defaults to `llama3.3`.
|
- `OLLAMA_MODEL`: defaults to `llama3.3`.
|
||||||
- `OLLAMA_TIMEOUT`: Ollama request timeout in seconds, defaults to `180`.
|
- `OLLAMA_TIMEOUT`: Ollama request timeout in seconds, defaults to `180`.
|
||||||
- `FGAI_THREAT_INTEL`: set to `1` to enable external threat intelligence lookups.
|
- `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.
|
- `VIRUSTOTAL_API_KEY`: VirusTotal API key for public IP reputation enrichment.
|
||||||
|
|
||||||
## Safety Model
|
## Safety Model
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ import os
|
|||||||
import time
|
import time
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from urllib import error, request
|
from urllib import error, request
|
||||||
|
from urllib.parse import urlencode
|
||||||
|
|
||||||
|
|
||||||
def is_public_ip(value: str | None) -> bool:
|
def is_public_ip(value: str | None) -> bool:
|
||||||
@@ -20,7 +21,10 @@ def is_public_ip(value: str | None) -> bool:
|
|||||||
class ThreatIntelClient:
|
class ThreatIntelClient:
|
||||||
def __init__(self, *, cache_file: str = "state/threat-intel-cache.json", ttl_seconds: int = 86400) -> None:
|
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.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.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.cache_path = Path(cache_file)
|
||||||
self.ttl_seconds = ttl_seconds
|
self.ttl_seconds = ttl_seconds
|
||||||
self.cache = self._read_cache()
|
self.cache = self._read_cache()
|
||||||
@@ -40,21 +44,70 @@ class ThreatIntelClient:
|
|||||||
def lookup_ip(self, ip: str) -> dict[str, object]:
|
def lookup_ip(self, ip: str) -> dict[str, object]:
|
||||||
if not is_public_ip(ip):
|
if not is_public_ip(ip):
|
||||||
return {"ip": ip, "provider": "local", "status": "skipped", "reason": "not_public_ip", "score": 0}
|
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())
|
now = int(time.time())
|
||||||
if cached and now - int(cached.get("cached_at", 0)) < self.ttl_seconds:
|
if cached and now - int(cached.get("cached_at", 0)) < self.ttl_seconds:
|
||||||
return cached
|
return cached
|
||||||
if not self.enabled:
|
if not self.enabled:
|
||||||
return {"ip": ip, "provider": "none", "status": "disabled", "score": 0}
|
return {"ip": ip, "provider": "none", "status": "disabled", "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:
|
if not self.virustotal_key:
|
||||||
return {"ip": ip, "provider": "virustotal", "status": "missing_api_key", "score": 0}
|
return {"ip": ip, "provider": "virustotal", "status": "missing_api_key", "score": 0}
|
||||||
|
|
||||||
result = self._lookup_virustotal_ip(ip)
|
result = self._lookup_virustotal_ip(ip)
|
||||||
|
else:
|
||||||
|
return {"ip": ip, "provider": provider, "status": "unsupported_provider", "score": 0}
|
||||||
|
|
||||||
result["cached_at"] = now
|
result["cached_at"] = now
|
||||||
self.cache[ip] = result
|
self.cache[cache_key] = result
|
||||||
self._write_cache()
|
self._write_cache()
|
||||||
return result
|
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]:
|
def _lookup_virustotal_ip(self, ip: str) -> dict[str, object]:
|
||||||
req = request.Request(
|
req = request.Request(
|
||||||
f"https://www.virustotal.com/api/v3/ip_addresses/{ip}",
|
f"https://www.virustotal.com/api/v3/ip_addresses/{ip}",
|
||||||
|
|||||||
@@ -33,6 +33,22 @@ class RecommendationTests(unittest.TestCase):
|
|||||||
|
|
||||||
self.assertEqual(result["status"], "disabled")
|
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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user