import unittest

from fgai.logs import parse_log_line
from fgai.mitigation import is_blockable_public_ip, suggest_block_candidates


class MitigationTests(unittest.TestCase):
    def test_private_ips_are_not_blockable(self):
        self.assertFalse(is_blockable_public_ip("192.168.1.10"))
        self.assertFalse(is_blockable_public_ip("10.0.0.5"))
        self.assertTrue(is_blockable_public_ip("8.8.8.8"))

    def test_suggests_repeated_public_utm_offender(self):
        events = [
            parse_log_line('type=utm subtype=ips srcip=8.8.8.8 action=blocked severity=high attack="scan"'),
            parse_log_line('type=utm subtype=ips srcip=8.8.8.8 action=blocked severity=high attack="scan"'),
            parse_log_line('type=utm subtype=ips srcip=8.8.8.8 action=blocked severity=high attack="scan"'),
            parse_log_line('type=utm subtype=ips srcip=192.168.1.5 action=blocked severity=critical attack="scan"'),
        ]

        candidates = suggest_block_candidates(events)

        self.assertEqual([candidate.src_ip for candidate in candidates], ["8.8.8.8"])


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