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


class ListTokensArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            # 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=1
                )]
            ),
            CommandParameter(
                name="process_id",
                type=ParameterType.Number,
                description="Process ID to list tokens from (from Process Browser)",
                parameter_group_info=[ParameterGroupInfo(
                    required=False,
                    ui_position=2
                )]
            ),
            CommandParameter(
                name="architecture",
                type=ParameterType.String,
                description="Process architecture (from Process Browser)",
                parameter_group_info=[ParameterGroupInfo(
                    required=False,
                    ui_position=3
                )]
            ),
        ]

    async def parse_arguments(self):
        pass

    async def parse_dictionary(self, dictionary_arguments):
        self.load_args_from_dictionary(dictionary_arguments)


class ListTokensCommand(CommandBase):
    cmd = "list_tokens"
    needs_admin = False
    help_cmd = "list_tokens"
    description = "Enumerate and list all available security tokens from running processes on the target system. Shows detailed token information including user context, privileges, and groups for each process. Compatible with Mythic Process Browser for interactive token enumeration. Useful for reconnaissance to identify high-value tokens before attempting token theft. Lower detection risk than token theft operations, but token enumeration may still be logged by security tools."
    version = 1
    author = "@KaseyaOFSTeam"
    supported_ui_features = ["process_browser:list_tokens"]
    argument_class = ListTokensArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about token enumeration risks.
        """
        pid = taskData.args.get_arg("process_id")
        
        message = "⚠️ OPSEC WARNING - Token Enumeration Detected\n\n"
        message += "Token enumeration is monitored by:\n"
        message += "  • EDR/XDR solutions (behavioral detection)\n"
        message += "  • Process monitoring tools\n"
        message += "  • Security auditing systems\n\n"
        
        if pid:
            message += f"Target Process ID: {pid}\n\n"
        
        message += "🔍 DETECTION MECHANISMS:\n"
        message += "  • Token enumeration API calls\n"
        message += "  • Process access for token queries\n"
        message += "  • Suspicious enumeration patterns\n"
        message += "  • Security event logging (if enabled)\n\n"
        
        message += "💡 CONSIDERATIONS:\n"
        message += "  • Enumerating tokens from processes may be logged\n"
        message += "  • Accessing sensitive processes (LSASS) is monitored\n"
        message += "  • Lower detection risk than token theft\n"
        message += "  • Useful for reconnaissance before token operations\n\n"
        
        message += "✅ This is a WARNING only - command will proceed.\n"
        message += "Ensure you understand the detection risks before continuing."
        
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=False,
            OpsecPreBypassRole="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("process_id")
        
        message = "⚠️ OPSEC POST-CHECK - Token Enumeration Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Process Access: OpenProcess for token queries\n"
        message += "  • Token Enumeration: Token enumeration API calls\n"
        message += "  • Process Information: Token details retrieved\n\n"
        
        if pid:
            message += f"Target Process ID: {pid}\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR behavioral detection (low-moderate confidence)\n"
        message += "  • Process access monitoring\n"
        message += "  • Security event logging (if enabled)\n\n"
        message += "⚠️ DETECTION PROBABILITY:\n"
        message += "  • Lower than token theft operations\n"
        message += "  • Process access may be logged\n"
        message += "  • Enumeration patterns may be detected\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,
        )
        response.DisplayParams = "Listing all available tokens..."
        return response

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

