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


class StealTokenArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="pid",
                type=ParameterType.Number,
                description="Process ID to steal token from",
                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 StealTokenCommand(CommandBase):
    cmd = "steal_token"
    needs_admin = False
    help_cmd = "steal_token <pid>"
    description = "Steal and impersonate the security token from a target process. Opens the target process, extracts its token, and impersonates it to assume the security context (user, privileges, groups) of that process. Compatible with Mythic Process Browser for interactive token theft. The impersonation context is automatically tracked and displayed in the Mythic UI. CRITICAL OPSEC RISK: Token theft is heavily monitored by EDR/XDR solutions, especially when targeting LSASS or other critical processes. Always requires approval from another operator. Prefer 'make_token' for domain credentials when possible."
    version = 1
    author = "@KaseyaOFSTeam"
    attackmapping = ["T1134.001"]
    supported_ui_features = ["process_browser:steal_token"]
    argument_class = StealTokenArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about token theft risks, especially for high-value processes.
        """
        # Debug: Print to see if this function is being called
        print(f"[DEBUG opsec_pre] steal_token opsec_pre called")
        print(f"[DEBUG opsec_pre] taskData.args: {taskData.args}")
        if hasattr(taskData, 'Task'):
            print(f"[DEBUG opsec_pre] taskData.Task: {taskData.Task}")
            if hasattr(taskData.Task, 'Params'):
                print(f"[DEBUG opsec_pre] taskData.Task.Params: {taskData.Task.Params}")
        
        # Try to get pid from parsed args first, then from raw parameters
        pid = taskData.args.get_arg("pid") or taskData.args.get_arg("process_id")
        print(f"[DEBUG opsec_pre] pid from args: {pid}")
        
        # If not found in parsed args, try to get from raw task parameters
        # (needed because opsec_pre runs before parse_dictionary)
        if not pid:
            # Try multiple ways to access the parameters
            if hasattr(taskData.Task, 'Params') and taskData.Task.Params:
                import json
                try:
                    raw_params = json.loads(taskData.Task.Params) if isinstance(taskData.Task.Params, str) else taskData.Task.Params
                    pid = raw_params.get("pid") or raw_params.get("process_id")
                except:
                    pass
            
            # Also try from taskData.Task directly if it has a params attribute
            if not pid and hasattr(taskData.Task, 'params') and taskData.Task.params:
                import json
                try:
                    raw_params = json.loads(taskData.Task.params) if isinstance(taskData.Task.params, str) else taskData.Task.params
                    pid = raw_params.get("pid") or raw_params.get("process_id")
                except:
                    pass
            
            # Try from taskData directly
            if not pid and hasattr(taskData, 'params') and taskData.params:
                import json
                try:
                    raw_params = json.loads(taskData.params) if isinstance(taskData.params, str) else taskData.params
                    pid = raw_params.get("pid") or raw_params.get("process_id")
                except:
                    pass
        
        # High-value processes that are heavily monitored
        high_value_pids = {
            4: "System (PID 4) - Critical system process, heavily monitored",
        }
        
        # Common high-value process names (if we could check them)
        high_value_processes = {
            "lsass.exe": "LSASS - Credential extraction is heavily monitored by EDR/XDR",
            "winlogon.exe": "Winlogon - Authentication process, monitored",
            "csrss.exe": "CSRSS - Critical system process, monitored",
            "services.exe": "Services - System service manager, monitored",
        }
        
        blocked = False
        critical_warning = ""
        
        if pid:
            pid_int = int(pid) if pid and str(pid).isdigit() else None
            if pid_int in high_value_pids:
                critical_warning = f"🚨 CRITICAL: {high_value_pids[pid_int]}\n⚠️ WARNING: Stealing from this process may cause system instability!\n\n"
                blocked = True
        
        # Always block steal_token to ensure OPSEC warning is shown
        # (similar to how shell always blocks for cmd.exe spawn)
        blocked = True
        
        message = "🚨 HIGH OPSEC RISK - Token Theft Operation\n\n"
        message += "Token theft operations are EXTREMELY monitored by:\n"
        message += "  • EDR/XDR solutions (Defender, CrowdStrike, SentinelOne)\n"
        message += "  • Credential protection (LSASS protection)\n"
        message += "  • Process monitoring tools\n"
        message += "  • Security auditing systems\n\n"
        
        if pid:
            message += f"Target Process ID: {pid}\n"
            if critical_warning:
                message += critical_warning
            message += "\n"
        
        message += "🔍 DETECTION MECHANISMS:\n"
        message += "  • OpenProcess(PROCESS_QUERY_INFORMATION) on sensitive processes\n"
        message += "  • Token impersonation API calls (behavioral detection)\n"
        message += "  • Privilege escalation detection\n"
        message += "  • Suspicious process access patterns\n"
        message += "  • LSASS memory access (if targeting LSASS)\n\n"
        
        message += "💡 ALTERNATIVES:\n"
        message += "  • Use 'make_token' for domain credentials (less detectable)\n"
        message += "  • Ensure proper evasion techniques are in place\n"
        message += "  • Verify LSASS protection status before targeting LSASS\n\n"
        
        if critical_warning:
            message += "🚨 CRITICAL WARNING: Targeting a critical system process!\n"
            message += "⚠️ TO PROCEED: You need approval from another operator.\n"
            message += "This command is blocked due to targeting a critical system process."
        else:
            message += "⚠️ TO PROCEED: You need approval from another operator.\n"
            message += "This operation is highly detectable and requires explicit approval."
        
        # Always return a message, even if pid is not found
        # This ensures opsec_pre always triggers
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=blocked,
            OpsecPreBypassRole="other_operator",  # Requires another operator to approve
            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 - Token Theft Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Process Access: OpenProcess on target process\n"
        message += "  • Token Manipulation: Token theft API calls\n"
        message += "  • Privilege Escalation: Token impersonation\n"
        message += "  • Context Tracking: Impersonation context updated in Mythic UI\n\n"
        
        if pid:
            message += f"Target Process ID: {pid}\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR behavioral detection (high confidence)\n"
        message += "  • Windows Event Logging (if enabled)\n"
        message += "  • Process monitoring tools\n"
        message += "  • Credential protection mechanisms\n\n"
        message += "⚠️ HIGH DETECTION PROBABILITY IF:\n"
        message += "  • Target is LSASS (PID varies by system)\n"
        message += "  • LSASS protection is enabled\n"
        message += "  • EDR has behavioral monitoring enabled\n"
        message += "  • Process is a critical system process\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,  # Don't block, just warn
            OpsecPostBypassRole="operator",  # Any operator can bypass
            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")
        response.DisplayParams = f"Stealing token from PID: {pid}"
        return response

    async def process_response(self, task: PTTaskMessageAllData, response: any) -> PTTaskProcessResponseMessageResponse:
        resp = PTTaskProcessResponseMessageResponse(TaskID=task.Task.ID, Success=True)
        return resp

