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


class MakeTokenArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="domain",
                type=ParameterType.String,
                description="Domain of the account credentials (e.g., acme.corp or empty for local)",
                parameter_group_info=[ParameterGroupInfo(
                    required=True,
                    ui_position=1
                )]
            ),
            CommandParameter(
                name="username",
                type=ParameterType.String,
                description="Username of the account",
                parameter_group_info=[ParameterGroupInfo(
                    required=True,
                    ui_position=2
                )]
            ),
            CommandParameter(
                name="password",
                type=ParameterType.String,
                description="The plaintext password for the account",
                parameter_group_info=[ParameterGroupInfo(
                    required=True,
                    ui_position=3
                )]
            ),
            CommandParameter(
                name="logon_type",
                type=ParameterType.Number,
                default_value=9,
                description="Logon type (default: 9 = LOGON32_LOGON_NEW_CREDENTIALS)",
                parameter_group_info=[ParameterGroupInfo(
                    required=False,
                    ui_position=4
                )]
            ),
        ]

    async def parse_arguments(self):
        parts = self.command_line.split(" ", 2)
        if len(parts) < 3:
            raise ValueError("Must supply domain, username, and password: make_token <domain> <username> <password> [logon_type]")
        
        self.add_arg("domain", parts[0])
        self.add_arg("username", parts[1])
        self.add_arg("password", parts[2])
        
        if len(parts) > 3:
            try:
                self.add_arg("logon_type", int(parts[3]))
            except ValueError:
                self.add_arg("logon_type", 9)

    async def parse_dictionary(self, dictionary_arguments):
        self.load_args_from_dictionary(dictionary_arguments)
        if not self.get_arg("logon_type"):
            self.add_arg("logon_type", 9)


class MakeTokenCommand(CommandBase):
    cmd = "make_token"
    needs_admin = False
    help_cmd = "make_token <domain> <username> <password> [logon_type]"
    description = "Create a new security token using plaintext credentials and impersonate it. Authenticates with the provided domain (or local machine if empty), username, and password using LogonUser API. The logon_type parameter (default: 9 - NewCredentials) determines the authentication type. The impersonation context is automatically tracked and displayed in the Mythic UI. Less detectable than 'steal_token' as it doesn't require accessing LSASS memory, but may still generate authentication logs. Use 'rev2self' to revert to the original token."
    version = 1
    author = "@KaseyaOFSTeam"
    attackmapping = ["T1134.003"]
    argument_class = MakeTokenArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about token creation risks.
        """
        domain = taskData.args.get_arg("domain") or ""
        username = taskData.args.get_arg("username") or ""
        logon_type = taskData.args.get_arg("logon_type") or 9
        
        message = "⚠️ OPSEC WARNING - Token Creation Detected\n\n"
        message += "Token creation operations are monitored by:\n"
        message += "  • EDR/XDR solutions (behavioral detection)\n"
        message += "  • Authentication monitoring tools\n"
        message += "  • Security auditing systems\n"
        message += "  • Windows Event Logging (if enabled)\n\n"
        
        message += f"Configuration:\n"
        message += f"  • Domain: {domain or '(local)'}\n"
        message += f"  • Username: {username}\n"
        message += f"  • Logon Type: {logon_type}\n\n"
        
        message += "🔍 DETECTION MECHANISMS:\n"
        message += "  • LogonUser API calls (authentication events)\n"
        message += "  • Token impersonation API calls\n"
        message += "  • Privilege escalation detection\n"
        message += "  • Suspicious authentication patterns\n"
        message += "  • Windows Event Logs (Security events)\n\n"
        
        message += "💡 COMPARED TO steal_token:\n"
        message += "  • Less detectable than steal_token (no LSASS access)\n"
        message += "  • Requires valid credentials (domain/local)\n"
        message += "  • May generate authentication logs\n"
        message += "  • Network authentication may be logged\n\n"
        
        message += "⚠️ WARNING: Authentication events may be logged.\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.
        """
        domain = taskData.args.get_arg("domain") or ""
        username = taskData.args.get_arg("username") or ""
        
        message = "⚠️ OPSEC POST-CHECK - Token Creation Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Authentication Event: LogonUser API call\n"
        message += "  • Token Creation: New token created\n"
        message += "  • Token Impersonation: Token impersonation API calls\n"
        message += "  • Context Tracking: Impersonation context updated in Mythic UI\n\n"
        
        message += f"Configuration:\n"
        message += f"  • Domain: {domain or '(local)'}\n"
        message += f"  • Username: {username}\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR behavioral detection (moderate confidence)\n"
        message += "  • Windows Event Logging (Security events)\n"
        message += "  • Authentication monitoring tools\n"
        message += "  • Network authentication logs (if domain)\n\n"
        message += "⚠️ DETECTION PROBABILITY:\n"
        message += "  • Lower than steal_token (no LSASS access)\n"
        message += "  • Authentication events may be logged\n"
        message += "  • Network authentication (domain) is logged\n"
        message += "  • Behavioral 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,
        )
        domain = taskData.args.get_arg("domain") or ""
        username = taskData.args.get_arg("username")
        logon_type = taskData.args.get_arg("logon_type") or 9
        response.DisplayParams = f"Creating token for {domain}\\{username} (logon_type: {logon_type})"
        return response

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

