105 lines
4.2 KiB
Python
105 lines
4.2 KiB
Python
import os
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from fgai.anomaly import detect_source_anomalies
|
|
from fgai.logs import parse_log_line
|
|
from fgai.recommendations import build_recommendations
|
|
from fgai.threat_intel import ThreatIntelClient
|
|
|
|
|
|
class RecommendationTests(unittest.TestCase):
|
|
def test_recommends_review_for_utm_anomaly(self):
|
|
events = [
|
|
parse_log_line(
|
|
'type=utm subtype=ips srcip=8.8.8.8 dstip=10.0.0.10 policyid=4 '
|
|
'service=https action=blocked severity=critical'
|
|
)
|
|
for _ in range(5)
|
|
]
|
|
anomalies = detect_source_anomalies(events)
|
|
|
|
recommendations = build_recommendations(events, anomalies)
|
|
|
|
self.assertGreaterEqual(recommendations[0].score, 60)
|
|
self.assertIn("4", recommendations[0].related_policy_ids)
|
|
self.assertIn("https", recommendations[0].related_services)
|
|
|
|
def test_policy_zero_is_implicit_deny_not_related_policy(self):
|
|
events = [
|
|
parse_log_line(
|
|
'type=traffic srcip=203.0.113.8 dstip=10.0.0.10 policyid=0 '
|
|
'service=ssh action=deny severity=warning'
|
|
)
|
|
for _ in range(12)
|
|
]
|
|
anomalies = detect_source_anomalies(events)
|
|
|
|
recommendations = build_recommendations(events, anomalies)
|
|
|
|
self.assertEqual(recommendations[0].related_policy_ids, [])
|
|
self.assertEqual(recommendations[0].title, "Implicit deny/drop traffic observed")
|
|
self.assertIn("policyid=0", recommendations[0].recommendation)
|
|
|
|
def test_threat_intel_disabled_by_default(self):
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
client = ThreatIntelClient(cache_file="/tmp/fgai-test-threat-cache.json")
|
|
|
|
result = client.lookup_ip("8.8.8.8")
|
|
|
|
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")
|
|
|
|
def test_successful_lookup_is_reused_from_cache(self):
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
os.environ,
|
|
{"FGAI_THREAT_INTEL": "1", "VIRUSTOTAL_API_KEY": "test", "FGAI_THREAT_INTEL_PROVIDER": "virustotal"},
|
|
clear=True,
|
|
):
|
|
client = ThreatIntelClient(cache_file=f"{directory}/intel.json")
|
|
with patch.object(client, "_lookup_virustotal_ip", return_value={"ip": "8.8.8.8", "provider": "virustotal", "status": "ok", "score": 0}) as lookup:
|
|
client.lookup_ip("8.8.8.8")
|
|
client.lookup_ip("8.8.8.8")
|
|
|
|
self.assertEqual(lookup.call_count, 1)
|
|
self.assertEqual(client.status()["requests_today"], 1)
|
|
|
|
def test_daily_limit_prevents_new_external_lookups(self):
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
os.environ,
|
|
{
|
|
"FGAI_THREAT_INTEL": "1",
|
|
"VIRUSTOTAL_API_KEY": "test",
|
|
"FGAI_THREAT_INTEL_PROVIDER": "virustotal",
|
|
"FGAI_THREAT_INTEL_DAILY_LIMIT": "1",
|
|
},
|
|
clear=True,
|
|
):
|
|
client = ThreatIntelClient(cache_file=f"{directory}/intel.json")
|
|
with patch.object(client, "_lookup_virustotal_ip", return_value={"ip": "8.8.8.8", "provider": "virustotal", "status": "ok", "score": 0}) as lookup:
|
|
client.lookup_ip("8.8.8.8")
|
|
limited = client.lookup_ip("1.1.1.1")
|
|
|
|
self.assertEqual(lookup.call_count, 1)
|
|
self.assertEqual(limited["status"], "daily_limit_reached")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|