import unittest

from fgai.policies import audit_policies, parse_policy_config, parse_policy_json


class PolicyTests(unittest.TestCase):
    def test_policy_audit_finds_broad_unprotected_allow(self):
        policies = parse_policy_config(
            """
config firewall policy
    edit 1
        set srcaddr "all"
        set dstaddr "all"
        set service "ALL"
        set action accept
        set logtraffic disable
    next
end
"""
        )

        findings = audit_policies(policies)

        titles = {finding.title for finding in findings}
        self.assertIn("Broad allow policy", titles)
        self.assertIn("Accepted traffic lacks UTM inspection", titles)

    def test_policy_audit_reads_fortios_api_json_shape(self):
        policies = parse_policy_json(
            """
{
  "results": [
    {
      "policyid": 7,
      "srcaddr": [{"name": "all"}],
      "dstaddr": [{"name": "all"}],
      "service": [{"name": "ALL"}],
      "action": "accept",
      "logtraffic": "disable"
    }
  ]
}
"""
        )

        findings = audit_policies(policies)

        self.assertEqual(policies[0].policy_id, "7")
        self.assertIn("Broad allow policy", {finding.title for finding in findings})


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