import unittest

from hawkSoar.security.hybrid_catalog import HybridToolCatalog


class HybridCatalogTest(unittest.TestCase):
    def test_export_contains_tools_and_trigger_map(self):
        catalog = HybridToolCatalog()
        data = catalog.export(group_id="g-1")
        self.assertIn("tools", data)
        self.assertIn("trigger_map", data)
        self.assertEqual(data["group_id"], "g-1")
        self.assertGreater(len(data["tools"]), 0)

    def test_includes_innovative_tools(self):
        catalog = HybridToolCatalog().export(group_id="g-1")
        tool_ids = {tool["tool_id"] for tool in catalog["tools"]}
        self.assertIn("dns.tunnel_score", tool_ids)
        self.assertIn("ad.replication_drift_probe", tool_ids)
        self.assertIn("gpo.startup_script_view", tool_ids)
        self.assertIn("tier0.admin_path_mapper", tool_ids)
        self.assertIn("legacy_protocol_exposure_map", tool_ids)
        self.assertIn("deception.honeypot.early_warning", tool_ids)
        self.assertIn("memory.suspect_modules", tool_ids)
        self.assertIn("endpoint.process_list", tool_ids)
        self.assertIn("endpoint.user.list", tool_ids)
        self.assertIn("endpoint.logged_in_users.login_time", tool_ids)
        self.assertIn("endpoint.registry.read", tool_ids)
        self.assertIn("endpoint.disk.read", tool_ids)
        self.assertIn("endpoint.disk.drives", tool_ids)
        self.assertIn("endpoint.task.terminate", tool_ids)
        self.assertIn("endpoint.process_snapshot_to_disk", tool_ids)
        self.assertIn("endpoint.command.run", tool_ids)
        self.assertIn("endpoint.dynamic_tool.register", tool_ids)
        self.assertIn("endpoint.eventlog.enriched", tool_ids)
        self.assertIn("endpoint.eventlog.channels", tool_ids)
        self.assertIn("net.http.client", tool_ids)
        self.assertIn("net.smb.client", tool_ids)
        self.assertIn("endpoint.remote_wmic.query", tool_ids)
        self.assertIn("endpoint.remote_command.execute", tool_ids)
        self.assertIn("endpoint.remote_registry.query", tool_ids)
        self.assertIn("endpoint.remote_config.update", tool_ids)
        self.assertIn("endpoint.remote_service.restart", tool_ids)
        self.assertIn("endpoint.windows_service.register", tool_ids)
        self.assertIn("endpoint.windows_service.unregister", tool_ids)
        self.assertIn("endpoint.windows_service.status", tool_ids)
        self.assertIn("endpoint.windows_service.restart", tool_ids)
        self.assertIn("endpoint.remote_process.restart", tool_ids)
        self.assertIn("living_off_land.detector", tool_ids)
        self.assertIn("deception.signal_fuser", tool_ids)
        self.assertIn("endpoint.remote_process_snapshot", tool_ids)
        self.assertIn("endpoint.remote_process_snapshot_by_name", tool_ids)
        self.assertIn("endpoint.remote_file_system_list", tool_ids)
        self.assertIn("endpoint.remote_file_collect", tool_ids)

    def test_includes_deploy_workflow_tools(self):
        catalog = HybridToolCatalog().export(group_id="g-1")
        tool_ids = {tool["tool_id"] for tool in catalog["tools"]}
        self.assertIn("deploy.hawk_soard", tool_ids)
        self.assertIn("deploy.hawk_soard.verify", tool_ids)

    def test_deploy_workflow_has_correct_shape(self):
        catalog = HybridToolCatalog()
        deploy = catalog.get_tool("deploy.hawk_soard")
        verify = catalog.get_tool("deploy.hawk_soard.verify")

        self.assertIsNotNone(deploy)
        self.assertEqual(deploy["action"], "deploy_hawk_soard")
        self.assertEqual(deploy["category"], "DEPLOY")
        self.assertTrue(deploy["requires_approval"])
        self.assertIn("target_host", deploy["parameters_schema"])
        self.assertIn("username", deploy["parameters_schema"])
        self.assertIn("password", deploy["parameters_schema"])
        self.assertEqual(deploy["workflow"]["step"], 1)
        self.assertEqual(deploy["workflow"]["next_action"], "verify_hawk_soard")

        self.assertIsNotNone(verify)
        self.assertEqual(verify["action"], "verify_hawk_soard")
        self.assertEqual(verify["category"], "DEPLOY")
        self.assertFalse(verify["requires_approval"])
        self.assertIn("target_host", verify["parameters_schema"])
        self.assertEqual(verify["workflow"]["step"], 2)
        self.assertEqual(verify["workflow"]["previous_action"], "deploy_hawk_soard")

    def test_deploy_trigger_plan_includes_both_steps(self):
        plan = HybridToolCatalog().build_trigger_plan(
            trigger="remote_agent_install_requested",
            group_id="g-1",
            case_id="c-1",
        )
        self.assertEqual(plan["trigger"], "remote_agent_install_requested")
        actions = [t["task_template"]["action"] for t in plan["tasks"]]
        self.assertIn("deploy_hawk_soard", actions)
        self.assertIn("verify_hawk_soard", actions)
        deploy_rank = next(t["rank"] for t in plan["tasks"] if t["task_template"]["action"] == "deploy_hawk_soard")
        verify_rank = next(t["rank"] for t in plan["tasks"] if t["task_template"]["action"] == "verify_hawk_soard")
        self.assertLess(deploy_rank, verify_rank)

    def test_includes_network_tools(self):
        catalog = HybridToolCatalog().export(group_id="g-1")
        tool_ids = {tool["tool_id"] for tool in catalog["tools"]}
        self.assertIn("net.ping", tool_ids)
        self.assertIn("net.dns_lookup", tool_ids)
        self.assertIn("net.traceroute", tool_ids)
        self.assertIn("net.arp_table", tool_ids)
        self.assertIn("net.route_table", tool_ids)
        self.assertIn("net.port_scan", tool_ids)
        self.assertIn("net.smb_sessions", tool_ids)
        self.assertIn("net.smb_share_query", tool_ids)

    def test_includes_ad_tools(self):
        catalog = HybridToolCatalog().export(group_id="g-1")
        tool_ids = {tool["tool_id"] for tool in catalog["tools"]}
        self.assertIn("ad.get_user_details", tool_ids)
        self.assertIn("ad.enumerate_group_membership", tool_ids)
        self.assertIn("ad.disable_account", tool_ids)
        self.assertIn("ad.enable_account", tool_ids)
        self.assertIn("ad.reset_password", tool_ids)
        self.assertIn("ad.retrieve_last_logon", tool_ids)
        self.assertIn("ad.audit_privileged_group_membership", tool_ids)

    def test_includes_containment_tools(self):
        catalog = HybridToolCatalog().export(group_id="g-1")
        tool_ids = {tool["tool_id"] for tool in catalog["tools"]}
        self.assertIn("containment.logoff_user", tool_ids)
        self.assertIn("containment.kill_process", tool_ids)
        self.assertIn("containment.stop_service", tool_ids)
        self.assertIn("containment.add_firewall_block_rule", tool_ids)
        self.assertIn("containment.host_isolation", tool_ids)
        self.assertIn("containment.disable_user", tool_ids)

    def test_all_containment_tools_require_approval(self):
        catalog = HybridToolCatalog().export(group_id="g-1")
        containment_tools = [t for t in catalog["tools"] if t["category"] == "CONTAINMENT"]
        self.assertGreater(len(containment_tools), 0)
        for tool in containment_tools:
            self.assertTrue(
                tool["requires_approval"],
                msg="Containment tool %s must require approval" % tool["tool_id"],
            )

    def test_includes_high_value_forensic_tools(self):
        catalog = HybridToolCatalog().export(group_id="g-1")
        tool_ids = {tool["tool_id"] for tool in catalog["tools"]}
        self.assertIn("endpoint.lsass_access_indicators", tool_ids)
        self.assertIn("endpoint.credential_theft_indicators", tool_ids)
        self.assertIn("endpoint.lateral_movement_artifacts", tool_ids)
        self.assertIn("endpoint.memory_ioc_extractor", tool_ids)
        self.assertIn("endpoint.ransomware_indicators", tool_ids)
        self.assertIn("endpoint.process_injection_scan", tool_ids)
        self.assertIn("endpoint.wmi_persistence_deep", tool_ids)
        self.assertIn("endpoint.ntfs_journal_analysis", tool_ids)
        self.assertIn("endpoint.mft_analyzer", tool_ids)

    def test_wmic_query_exposes_allowed_values(self):
        catalog = HybridToolCatalog()
        tool = catalog.get_tool("endpoint.wmic_query")
        self.assertIsNotNone(tool)
        schema = tool["parameters_schema"]
        self.assertIn("query", schema)
        allowed = schema["query"]
        self.assertIsInstance(allowed, list)
        self.assertIn("process_list_brief", allowed)
        self.assertIn("computersystem_get_username", allowed)
        self.assertIn("service_list", allowed)
        self.assertIn("qfe_list", allowed)

    def test_remote_actions_disclose_feature_flags(self):
        flag_requirements = {
            "remote_wmic_query":       "enable_remote_wmi_collection",
            "remote_process_snapshot": "enable_remote_process_snapshot",
            "remote_file_collect":     "enable_remote_file_collection",
            "remote_command_execute":  "enable_remote_command_execution",
            "remote_registry_query":   "enable_remote_registry_read",
            "remote_config_update":    "enable_remote_config_management",
            "run_command":             "enable_command_execution",
            "register_dynamic_tool":   "enable_command_execution",
        }
        catalog = HybridToolCatalog()
        for action, flag in flag_requirements.items():
            tool = next(
                (t for t in catalog._all_tools() if t["action"] == action), None
            )
            self.assertIsNotNone(tool, msg="Tool for action %s not found" % action)
            self.assertIn(
                flag,
                tool["description"],
                msg="Tool '%s' description must mention config flag '%s'" % (action, flag),
            )

    def test_windows_only_actions_flagged_correctly(self):
        catalog = HybridToolCatalog()
        data = catalog.export(group_id="g-1")
        windows_only_actions = {
            "deploy_hawk_soard", "verify_hawk_soard",
            "get_user_details", "reset_password", "disable_account",
            "event_log_enriched", "registry_read", "smb_sessions",
            "run_powershell", "kerberos_ticket_inspect", "lsass_access_indicators",
            "credential_theft_indicators", "ntfs_journal_analysis",
            "logoff_user", "kill_process", "host_isolation",
        }
        action_to_platform = {t["action"]: t.get("supported_on") for t in data["tools"]}
        for action in windows_only_actions:
            self.assertEqual(
                action_to_platform.get(action),
                "windows",
                msg="Action '%s' should be supported_on='windows'" % action,
            )

    def test_trigger_plan_returns_executable_templates_only_by_default(self):
        plan = HybridToolCatalog().build_trigger_plan(
            trigger="honeypot_trip",
            group_id="g-1",
            case_id="c-1",
            include_trigger_only=False,
        )
        self.assertEqual(plan["trigger"], "honeypot_trip")
        self.assertGreaterEqual(len(plan["tasks"]), 1)
        for task in plan["tasks"]:
            tt = task["task_template"]
            self.assertIn("task_id", tt)
            self.assertIn("category", tt)
            self.assertIn("action", tt)
            self.assertIn("requires_approval", tt)
            self.assertIn("timeout", tt)
            self.assertIn("priority", tt)

    def test_honeypot_has_protocol_variants(self):
        catalog = HybridToolCatalog().export(group_id="g-1")
        honeypot = None
        for tool in catalog["tools"]:
            if tool["tool_id"] == "deception.honeypot.early_warning":
                honeypot = tool
                break
        self.assertIsNotNone(honeypot)
        self.assertEqual(len(honeypot.get("supported_protocols", [])), 10)
        self.assertIn("honeypot_trip_smb", catalog["trigger_map"])
        self.assertIn("honeypot_trip_rpc", catalog["trigger_map"])

    def test_tier0_admin_path_mapper_marked_pending(self):
        catalog = HybridToolCatalog()
        tool = catalog.get_tool("tier0.admin_path_mapper")
        self.assertIsNotNone(tool)
        self.assertEqual(tool["execution_mode"], "trigger_only")
        self.assertEqual(tool.get("status"), "pending_implementation")

    def test_no_executable_tool_missing_description(self):
        catalog = HybridToolCatalog()
        for tool in catalog._all_tools():
            if tool.get("execution_mode") != "executable":
                continue
            desc = tool.get("description", "")
            self.assertGreater(
                len(desc),
                15,
                msg="Tool '%s' has an unacceptably short description: '%s'" % (tool["tool_id"], desc),
            )

    def test_catalog_tool_count(self):
        catalog = HybridToolCatalog()
        tools = catalog._all_tools()
        self.assertGreaterEqual(len(tools), 100, msg="Catalog should have at least 100 tools")


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