from mythic_container.MythicCommandBase import *
from mythic_container.MythicRPC import *


class KillArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="pid",
                type=ParameterType.Number,
                description="Process ID to kill",
                parameter_group_info=[ParameterGroupInfo(
                    required=True,
                    ui_position=1
                )]
            ),
            # Process Browser parameters (optional, passed automatically by Mythic)
            CommandParameter(
                name="host",
                type=ParameterType.String,
                description="Host name (from Process Browser)",
                parameter_group_info=[ParameterGroupInfo(
                    required=False,
                    ui_position=2
                )]
            ),
            CommandParameter(
                name="process_id",
                type=ParameterType.Number,
                description="Process ID (from Process Browser, takes precedence over pid)",
                parameter_group_info=[ParameterGroupInfo(
                    required=False,
                    ui_position=3
                )]
            ),
            CommandParameter(
                name="architecture",
                type=ParameterType.String,
                description="Process architecture (from Process Browser)",
                parameter_group_info=[ParameterGroupInfo(
                    required=False,
                    ui_position=4
                )]
            ),
        ]

    async def parse_arguments(self):
        if len(self.command_line) == 0:
            # Check if process_id is provided (from Process Browser)
            if not self.get_arg("process_id") and not self.get_arg("pid"):
                raise ValueError("Must supply a PID")
        else:
            try:
                self.add_arg("pid", int(self.command_line))
            except ValueError:
                raise ValueError("PID must be a number")

    async def parse_dictionary(self, dictionary_arguments):
        self.load_args_from_dictionary(dictionary_arguments)
        # If process_id is provided (from Process Browser), use it instead of pid
        if self.get_arg("process_id") and not self.get_arg("pid"):
            self.add_arg("pid", self.get_arg("process_id"))


class KillCommand(CommandBase):
    cmd = "kill"
    needs_admin = False
    help_cmd = "kill <pid>"
    description = "Terminate a process by its Process ID (PID). The process is immediately terminated using TerminateProcess API. Compatible with Mythic Process Browser for interactive process termination. Automatically removes the killed process from the Process Browser view. WARNING: Critical system processes (PID 0, 4, 8) are blocked from termination to prevent system instability. Creates a 'Process Termination' artifact for tracking."
    version = 1
    author = "@KaseyaOFSTeam"
    attackmapping = ["T1489"]
    supported_ui_features = ["process_browser:kill"]
    argument_class = KillArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about process termination risks, especially for critical processes.
        """
        pid = taskData.args.get_arg("pid") or taskData.args.get_arg("process_id")
        
        # Critical system processes that should NEVER be killed
        critical_pids = {
            0: "Idle Process - Cannot be killed",
            4: "System Process - Critical system process, NEVER kill",
            8: "Memory Compression - Critical system process",
        }
        
        blocked = False
        critical_warning = ""
        
        if pid:
            pid_int = int(pid) if pid and str(pid).isdigit() else None
            if pid_int in critical_pids:
                critical_warning = f"🚨 CRITICAL: {critical_pids[pid_int]}\n⚠️ WARNING: Killing this process may cause system instability or crash!\n\n"
                blocked = True
        
        message = "⚠️ OPSEC RISK - Process Termination Detected\n\n"
        message += "Process termination is monitored by:\n"
        message += "  • EDR/XDR solutions\n"
        message += "  • Process monitoring tools\n"
        message += "  • Security auditing systems\n"
        message += "  • Endpoint protection platforms\n\n"
        
        if pid:
            message += f"Target Process ID: {pid}\n"
            if critical_warning:
                message += critical_warning
            message += "\n"
        
        message += "🔍 DETECTION MECHANISMS:\n"
        message += "  • TerminateProcess API calls\n"
        message += "  • Process termination events\n"
        message += "  • Suspicious process killing patterns\n"
        message += "  • Termination of security processes\n\n"
        message += "💡 RECOMMENDATIONS:\n"
        message += "  • Verify the process is not critical using Process Browser\n"
        message += "  • Ensure you understand the impact before proceeding\n"
        message += "  • Check if process is a security tool (EDR, AV) - may trigger immediate alerts\n\n"
        
        if blocked:
            message += "⚠️ TO PROCEED: You need approval from another operator.\n"
            message += "This command is blocked due to targeting a critical system process."
        else:
            message += "⚠️ WARNING: Killing security processes may trigger immediate alerts!\n"
            message += "This is a WARNING only - command will proceed."
        
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=blocked,
            OpsecPreBypassRole="other_operator" if blocked else "operator",
            OpsecPreMessage=message
        )

    async def opsec_post(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPostTaskMessageResponse:
        """
        OPSEC check after creating the task.
        Warns about artifacts that will be created.
        """
        pid = taskData.args.get_arg("pid") or taskData.args.get_arg("process_id")
        
        message = "⚠️ OPSEC POST-CHECK - Process Termination Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Process Termination: TerminateProcess API call\n"
        message += "  • Process Access: OpenProcess with PROCESS_TERMINATE\n"
        message += "  • Event Log: Process termination event (if logging enabled)\n\n"
        
        if pid:
            message += f"Target Process ID: {pid}\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR behavioral detection\n"
        message += "  • Windows Event Logging (if enabled)\n"
        message += "  • Process monitoring tools\n"
        message += "  • Security information systems\n\n"
        message += "⚠️ HIGH DETECTION PROBABILITY IF:\n"
        message += "  • Target is a security process (EDR, AV)\n"
        message += "  • Target is a critical system process\n"
        message += "  • Multiple processes are terminated rapidly\n"
        message += "  • Process is a protected service\n\n"
        message += "✅ This is a WARNING only - task will proceed.\n"
        message += "Ensure you understand the detection risks before continuing."
        
        return PTTTaskOPSECPostTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPostBlocked=False,
            OpsecPostBypassRole="operator",
            OpsecPostMessage=message
        )

    async def create_go_tasking(self, taskData: PTTaskMessageAllData) -> PTTaskCreateTaskingMessageResponse:
        response = PTTaskCreateTaskingMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
        )
        pid = taskData.args.get_arg("pid")
        if pid:
            response.DisplayParams = f"Killing process PID: {pid}"
        else:
            response.DisplayParams = "Killing process (no PID specified)"
        return response

    async def process_response(self, task: PTTaskMessageAllData, response: any) -> PTTaskProcessResponseMessageResponse:
        resp = PTTaskProcessResponseMessageResponse(TaskID=task.Task.ID, Success=True)
        
        # Check if kill was successful by checking user_output
        if response:
            user_output = response.get("user_output", "")
            # Check if output indicates success (doesn't contain error indicators)
            if user_output and not any(err in user_output.lower() for err in ["error", "no se pudo", "falló", "failed"]):
                # Get the PID that was killed
                pid = task.args.get_arg("pid") or task.args.get_arg("process_id")
                if pid:
                    # Get hostname
                    host = task.Callback.Host if task.Callback else ""
                    
                    # Add removed_processes array to response (similar to removed_files for file browser)
                    # This tells Mythic Process Browser to remove this process from the cache
                    resp.ProcessResponse = response if isinstance(response, dict) else {}
                    if "removed_processes" not in resp.ProcessResponse:
                        resp.ProcessResponse["removed_processes"] = []
                    resp.ProcessResponse["removed_processes"].append({
                        "host": host,
                        "process_id": int(pid)
                    })
                    logging.info(f"[PROCESS_BROWSER] Marked process PID {pid} as removed in Process Browser")
        
        return resp

