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


class DownloadArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="file",
                type=ParameterType.String,
                description="Path to the file to download from the target",
                parameter_group_info=[ParameterGroupInfo(
                    required=True
                )]
            ),
        ]

    async def parse_arguments(self):
        if len(self.command_line) == 0:
            raise ValueError("Must supply a path to download")
        self.add_arg("file", self.command_line)

    async def parse_dictionary(self, dictionary_arguments):
        if "file" in dictionary_arguments:
            self.add_arg("file", dictionary_arguments["file"])


class DownloadCommand(CommandBase):
    cmd = "download"
    needs_admin = False
    help_cmd = "download <file_path>"
    description = "Download a file from the target system to the Mythic server. Supports files up to 2GB in size. Files are transferred in 512KB chunks for efficient transmission. Supports both absolute and relative paths (relative paths are resolved using the current working directory). Compatible with Mythic File Browser for interactive file selection. Large files or sensitive files may trigger detection by EDR/XDR solutions and data loss prevention (DLP) systems."
    version = 1
    supported_ui_features = ["file_browser:download"]
    author = "@KaseyaOFSTeam"
    argument_class = DownloadArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about file download risks.
        """
        file_path = taskData.args.get_arg("file") or ""
        
        # Sensitive file patterns
        sensitive_patterns = {
            "sam": "SAM database - Credential extraction heavily monitored",
            "system": "SYSTEM hive - Registry access logged",
            "security": "SECURITY hive - Security database access logged",
            "ntds.dit": "NTDS.dit - Active Directory database, highly monitored",
            "lsass": "LSASS memory - Credential extraction heavily monitored",
            "passwd": "Password files - Credential access monitored",
            "shadow": "Shadow files - Credential access monitored",
        }
        
        warnings = []
        file_lower = file_path.lower()
        
        for pattern, reason in sensitive_patterns.items():
            if pattern in file_lower:
                warnings.append(f"🚨 {pattern}: {reason}")
        
        message = "⚠️ OPSEC WARNING - File Download Detected\n\n"
        message += "File downloads are monitored by:\n"
        message += "  • EDR/XDR solutions\n"
        message += "  • Data loss prevention (DLP) systems\n"
        message += "  • File access monitoring tools\n"
        message += "  • Network monitoring (large file transfers)\n\n"
        
        if file_path:
            message += f"Target File: {file_path}\n\n"
        
        if warnings:
            message += "🚨 SENSITIVE FILE DETECTED:\n"
            for warning in warnings:
                message += f"  {warning}\n"
            message += "\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • File access events (if auditing enabled)\n"
        message += "  • Large file transfers (network monitoring)\n"
        message += "  • Sensitive file patterns (DLP detection)\n"
        message += "  • Unusual file access patterns\n\n"
        message += "💡 CONSIDERATIONS:\n"
        message += "  • Sensitive files (credentials, databases) are heavily monitored\n"
        message += "  • Large files may trigger network alerts\n"
        message += "  • Multiple downloads may be suspicious\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.
        """
        file_path = taskData.args.get_arg("file") or ""
        
        message = "⚠️ OPSEC POST-CHECK - File Download Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • File Access: Read operations on target file\n"
        message += "  • Network Activity: File data transmission to Mythic\n"
        message += "  • File Access Logs: Event logs (if auditing enabled)\n\n"
        
        if file_path:
            message += f"Target File: {file_path}\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR file access monitoring\n"
        message += "  • DLP systems (sensitive file detection)\n"
        message += "  • Network monitoring (large transfers)\n"
        message += "  • File access auditing (if enabled)\n\n"
        message += "⚠️ HIGH DETECTION PROBABILITY IF:\n"
        message += "  • File contains credentials or sensitive data\n"
        message += "  • File is large (network anomaly detection)\n"
        message += "  • File access auditing is enabled\n"
        message += "  • DLP rules match file patterns\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,
        )
        return response

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

