import json
import os
import tempfile
import unittest

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa

from hawkSoar.security.hybrid_catalog import HybridToolCatalog
from hawkSoar.security.hybrid_execution import HybridTaskExecutor


def _base_task():
    return {
        "task_id": "11111111-1111-1111-1111-111111111111",
        "group_id": "22222222-2222-2222-2222-222222222222",
        "case_id": "33333333-3333-3333-3333-333333333333",
        "category": "NETWORK",
        "action": "ping",
        "parameters": {"target": "127.0.0.1", "count": 1},
        "requires_approval": True,
        "timeout": 5,
        "priority": "LOW",
    }


class HybridExecutionTest(unittest.TestCase):
    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        public_key = self.private_key.public_key().public_bytes(
            serialization.Encoding.PEM,
            serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        self.pub_path = os.path.join(self.temp_dir.name, "public.pem")
        with open(self.pub_path, "wb") as handle:
            handle.write(public_key)

    def tearDown(self):
        self.temp_dir.cleanup()

    def _sign(self, task):
        canonical = json.dumps(task, sort_keys=True, separators=(",", ":")).encode("utf-8")
        sig = self.private_key.sign(canonical, padding.PKCS1v15(), hashes.SHA256())
        return sig.hex()

    def _executor(self, expected_group_getter=None, **overrides):
        config = {
            "task_signing_public_key": self.pub_path,
            "enforce_signed_hybrid_tasks": True,
            "enforce_group_id": True,
            "max_task_timeout": 60,
            "plugin_data_path": self.temp_dir.name,
            "hybrid_dynamic_tools_file": os.path.join(self.temp_dir.name, "hybrid_dynamic_tools.json"),
            "enable_command_execution": True,
        }
        config.update(overrides)

        def fake_runner(cmd, timeout):
            del timeout
            return 0, "ok: %s" % " ".join(cmd), ""

        return HybridTaskExecutor(
            config=config,
            expected_group_id_getter=expected_group_getter
            or (lambda: "22222222-2222-2222-2222-222222222222"),
            command_runner=fake_runner,
        )

    def test_blocks_missing_signature(self):
        executor = self._executor()
        response = executor.process(_base_task(), signature=None)
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("signature", response["error"].lower())

    def test_blocks_non_object_task_payload(self):
        executor = self._executor()
        response = executor.process(None, signature=None)
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("json object", response["error"].lower())

    def test_blocks_group_mismatch(self):
        executor = self._executor()
        task = _base_task()
        task["group_id"] = "aaaaaaaa-2222-2222-2222-222222222222"
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("group_id", response["error"])

    def test_blocks_disallowed_wmic(self):
        executor = self._executor()
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "wmic_query"
        task["parameters"] = {"query": "process_create"}
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("allowlisted", response["error"])

    def test_blocks_containment_without_approval(self):
        executor = self._executor()
        task = _base_task()
        task["category"] = "CONTAINMENT"
        task["action"] = "kill_process"
        task["requires_approval"] = False
        task["parameters"] = {"pid": 100}
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("require approval", response["error"].lower())

    def test_successful_signed_task(self):
        executor = self._executor(enforce_group_id=False)
        task = _base_task()
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "SUCCESS")
        self.assertIn("stdout", response["output"])

    def test_blocks_remote_process_snapshot_when_policy_disabled(self):
        executor = self._executor(enforce_group_id=False)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "remote_process_snapshot"
        task["requires_approval"] = True
        task["parameters"] = {
            "target_host": "host01",
            "pid": 1234,
            "username": "admin",
            "password": "pw",
        }
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("disabled by policy", response["error"].lower())

    def test_blocks_remote_process_snapshot_by_name_without_approval(self):
        executor = self._executor(enforce_group_id=False)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "remote_process_snapshot_by_name"
        task["requires_approval"] = False
        task["parameters"] = {
            "target_host": "host01",
            "process_name": "lsass",
            "username": "admin",
            "password": "pw",
        }
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("requires approval", response["error"].lower())

    def test_blocks_remote_file_collect_when_policy_disabled(self):
        executor = self._executor(enforce_group_id=False)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "remote_file_collect"
        task["requires_approval"] = True
        task["parameters"] = {
            "target_host": "host01",
            "remote_path": "C:\\Windows\\Temp\\a.bin",
            "username": "admin",
            "password": "pw",
        }
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("disabled by policy", response["error"].lower())

    def test_blocks_run_command_without_approval(self):
        executor = self._executor(enforce_group_id=False)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "run_command"
        task["requires_approval"] = False
        task["parameters"] = {"argv": ["whoami"]}
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("requires approval", response["error"].lower())

    def test_register_and_execute_dynamic_tool(self):
        executor = self._executor(enforce_group_id=False)
        register = _base_task()
        register["category"] = "ENDPOINT"
        register["action"] = "register_dynamic_tool"
        register["requires_approval"] = True
        register["parameters"] = {
            "tool_id": "custom.echo",
            "name": "Custom Echo",
            "category": "ENDPOINT",
            "action": "custom_echo",
            "description": "Echo a message.",
            "command": ["echo", "{message}"],
            "parameters_schema": {"message": "hello"},
            "requires_approval": True,
            "risk": "HIGH",
        }
        register_response = executor.process(register, signature=self._sign(register))
        self.assertEqual(register_response["status"], "SUCCESS")

        run_custom = _base_task()
        run_custom["category"] = "ENDPOINT"
        run_custom["action"] = "custom_echo"
        run_custom["requires_approval"] = True
        run_custom["parameters"] = {"message": "hello"}
        run_response = executor.process(run_custom, signature=self._sign(run_custom))
        self.assertEqual(run_response["status"], "SUCCESS")
        self.assertIn("stdout", run_response["output"])

    def test_blocks_remote_command_execute_when_policy_disabled(self):
        executor = self._executor(enforce_group_id=False, enable_remote_command_execution=False)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "remote_command_execute"
        task["requires_approval"] = True
        task["parameters"] = {
            "target_host": "host01",
            "username": "admin",
            "password": "pw",
            "argv": ["whoami"],
        }
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")

    def test_blocks_remote_registry_query_when_policy_disabled(self):
        executor = self._executor(enforce_group_id=False, enable_remote_registry_read=False)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "remote_registry_query"
        task["requires_approval"] = False
        task["parameters"] = {
            "target_host": "host01",
            "path": "HKLM\\SOFTWARE",
        }
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")

    def test_blocks_register_windows_service_without_approval(self):
        executor = self._executor(enforce_group_id=False)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "register_windows_service"
        task["requires_approval"] = False
        task["parameters"] = {
            "service_name": "HawkSoard",
            "executable_path": "C:\\does-not-exist\\hawk-soard.exe",
        }
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("requires approval", response["error"].lower())

    def test_blocks_restart_windows_service_without_approval(self):
        executor = self._executor(enforce_group_id=False)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "restart_windows_service"
        task["requires_approval"] = False
        task["parameters"] = {"service_name": "HawkSoard"}
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("requires approval", response["error"].lower())

    def test_catalog_executable_actions_are_wired(self):
        executor = self._executor(enforce_group_id=False)
        catalog = HybridToolCatalog(config=executor.config).export(group_id="g-1")
        executable_actions = set(
            tool["action"]
            for tool in catalog["tools"]
            if tool.get("execution_mode") == "executable" and isinstance(tool.get("action"), str)
        )
        allowlisted = set()
        for actions in executor.ALLOWED_ACTIONS.values():
            allowlisted.update(actions)
        missing_allowlist = sorted(x for x in executable_actions if x not in allowlisted)
        missing_handlers = sorted(
            x
            for x in executable_actions
            if x not in executor._action_handlers and x not in executor.dynamic_tool_actions
        )
        self.assertEqual([], missing_allowlist)
        self.assertEqual([], missing_handlers)

    def test_blocks_remote_config_update_without_approval(self):
        executor = self._executor(enforce_group_id=False, enable_remote_config_management=True)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "remote_config_update"
        task["requires_approval"] = False
        task["parameters"] = {
            "target_host": "host01",
            "remote_path": "C:\\ProgramData\\Hawk\\hawk-soard.cfg",
            "content": "[SETTINGS]\n",
            "username": "admin",
            "password": "pw",
        }
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("requires approval", response["error"].lower())

    def test_blocks_remote_service_restart_when_policy_disabled(self):
        executor = self._executor(enforce_group_id=False, enable_remote_config_management=False)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "remote_service_restart"
        task["requires_approval"] = True
        task["parameters"] = {
            "target_host": "host01",
            "service_name": "HawkSoard",
            "username": "admin",
            "password": "pw",
        }
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")

    def test_blocks_remote_process_restart_when_policy_disabled(self):
        executor = self._executor(enforce_group_id=False, enable_remote_config_management=False)
        task = _base_task()
        task["category"] = "ENDPOINT"
        task["action"] = "remote_process_restart"
        task["requires_approval"] = True
        task["parameters"] = {
            "target_host": "host01",
            "process_name": "hawk-soard",
            "executable_path": "C:\\Program Files\\Hawk\\hawk-soard.exe",
            "arguments": ["-c", "C:\\ProgramData\\Hawk\\hawk-soard.cfg"],
            "username": "admin",
            "password": "pw",
        }
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")

    def test_blocks_when_local_group_unavailable(self):
        executor = self._executor(expected_group_getter=lambda: None)
        task = _base_task()
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("local node group is unavailable", response["error"].lower())

    def test_blocks_when_group_getter_missing(self):
        executor = HybridTaskExecutor(
            config={
                "task_signing_public_key": self.pub_path,
                "enforce_signed_hybrid_tasks": True,
                "enforce_group_id": True,
                "max_task_timeout": 60,
            },
            expected_group_id_getter=None,
            command_runner=lambda cmd, timeout: (0, "ok", ""),
        )
        task = _base_task()
        response = executor.process(task, signature=self._sign(task))
        self.assertEqual(response["status"], "BLOCKED")
        self.assertIn("group resolver", response["error"].lower())


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