import unittest

from fgai.anomaly import anomaly_summary, detect_source_anomalies
from fgai.logs import parse_log_line


class AnomalyTests(unittest.TestCase):
    def test_scores_repeated_utm_public_source_as_anomaly(self):
        events = [
            parse_log_line(f'type=traffic srcip=10.0.0.{i} dstip=1.1.1.1 service=https action=accept sentbyte=100 rcvdbyte=100')
            for i in range(1, 8)
        ]
        events.extend(
            parse_log_line(
                'type=utm subtype=ips srcip=8.8.8.8 dstip=10.0.0.10 service=https '
                'action=blocked severity=critical sentbyte=0 rcvdbyte=0'
            )
            for _ in range(5)
        )

        findings = detect_source_anomalies(events)

        self.assertEqual(findings[0].subject, "8.8.8.8")
        self.assertGreaterEqual(findings[0].score, 60)
        self.assertIn(findings[0].severity, {"high", "critical"})

    def test_summary_counts_severities(self):
        events = [
            parse_log_line(
                'type=utm subtype=ips srcip=8.8.8.8 dstip=10.0.0.10 service=https '
                'action=blocked severity=critical'
            )
            for _ in range(5)
        ]

        summary = anomaly_summary(detect_source_anomalies(events))

        self.assertEqual(summary["total"], 1)


if __name__ == "__main__":
    unittest.main()