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


class UploadArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="file",
                type=ParameterType.File,
                description="File to upload to the target",
                parameter_group_info=[ParameterGroupInfo(
                    required=True,
                    ui_position=1
                )]
            ),
            CommandParameter(
                name="path",
                type=ParameterType.String,
                description="Path where the file should be saved on the target",
                parameter_group_info=[ParameterGroupInfo(
                    required=True,
                    ui_position=2
                )]
            ),
        ]

    async def parse_arguments(self):
        if len(self.command_line) == 0:
            raise ValueError("Must supply a file and path")
        
        parts = self.command_line.split(" ", 1)
        if len(parts) != 2:
            raise ValueError("Must supply both file and path: upload <file_id> <path>")
        
        self.add_arg("file", parts[0])
        # Normalize path in parse_arguments too
        path = parts[1]
        if path:
            # Replace all sequences of double backslashes with single backslash
            while "\\\\" in path:
                path = path.replace("\\\\", "\\")
        self.add_arg("path", path)

    async def parse_dictionary(self, dictionary_arguments):
        if "file" in dictionary_arguments:
            self.add_arg("file", dictionary_arguments["file"])
        if "path" in dictionary_arguments:
            # Normalize path: replace double backslashes with single backslash
            path = dictionary_arguments["path"]
            if path:
                # Replace all sequences of double backslashes with single backslash
                while "\\\\" in path:
                    path = path.replace("\\\\", "\\")
            self.add_arg("path", path)


class UploadCommand(CommandBase):
    cmd = "upload"
    needs_admin = False
    help_cmd = "upload <file_id> <destination_path>"
    description = "Upload a file from the Mythic server to the target system. Requires a file_id (obtained from Mythic's file browser) and a destination path where the file will be written. Supports both absolute and relative paths (relative paths are resolved using the current working directory). Creates a 'File Write' artifact for tracking. Uploads to sensitive locations (System32, Windows, Program Files) or with suspicious extensions (.exe, .dll, .ps1, .bat, .vbs) may trigger detection by security tools."
    version = 1
    supported_ui_features = ["file_browser:upload"]
    author = "@KaseyaOFSTeam"
    argument_class = UploadArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about file upload risks.
        """
        path = taskData.args.get_arg("path") or ""
        
        # Sensitive locations
        sensitive_locations = {
            "system32": "System32 directory - Highly monitored",
            "syswow64": "SysWOW64 directory - Highly monitored",
            "windows": "Windows directory - Highly monitored",
            "program files": "Program Files - Application installation monitoring",
            "temp": "Temporary directories - May be scanned",
        }
        
        # Suspicious file extensions
        suspicious_extensions = {
            ".exe": "Executable files - May trigger AV/EDR",
            ".dll": "DLL files - May trigger AV/EDR",
            ".ps1": "PowerShell scripts - Heavily monitored",
            ".bat": "Batch files - May trigger behavioral detection",
            ".vbs": "VBScript files - May trigger behavioral detection",
        }
        
        warnings = []
        path_lower = path.lower()
        
        for location, reason in sensitive_locations.items():
            if location in path_lower:
                warnings.append(f"⚠️ {location}: {reason}")
        
        for ext, reason in suspicious_extensions.items():
            if path_lower.endswith(ext):
                warnings.append(f"🚨 {ext}: {reason}")
        
        message = "⚠️ OPSEC WARNING - File Upload Detected\n\n"
        message += "File uploads are monitored by:\n"
        message += "  • EDR/XDR solutions (file write monitoring)\n"
        message += "  • Anti-malware solutions (file scanning)\n"
        message += "  • Application control mechanisms\n"
        message += "  • Network monitoring (file transfers)\n\n"
        
        if path:
            message += f"Target Path: {path}\n\n"
        
        if warnings:
            message += "🚨 RISKS DETECTED:\n"
            for warning in warnings:
                message += f"  {warning}\n"
            message += "\n"
        
        message += "🔍 DETECTION MECHANISMS:\n"
        message += "  • File write events (if auditing enabled)\n"
        message += "  • File scanning (AV/EDR real-time protection)\n"
        message += "  • Suspicious file patterns (executables, scripts)\n"
        message += "  • Network transfers (if monitored)\n\n"
        message += "💡 CONSIDERATIONS:\n"
        message += "  • Executable files (.exe, .dll) are heavily scanned\n"
        message += "  • Scripts (.ps1, .bat, .vbs) are monitored\n"
        message += "  • System directories are highly monitored\n"
        message += "  • Temporary directories may be scanned\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.
        """
        path = taskData.args.get_arg("path") or ""
        
        message = "⚠️ OPSEC POST-CHECK - File Upload Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • File Creation: New file written to disk\n"
        message += "  • File Write: Data written to target path\n"
        message += "  • Network Activity: File data transmission from Mythic\n"
        message += "  • File Access Logs: Event logs (if auditing enabled)\n\n"
        
        if path:
            message += f"Target Path: {path}\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR file write monitoring\n"
        message += "  • Anti-malware real-time scanning\n"
        message += "  • Application control mechanisms\n"
        message += "  • File access auditing (if enabled)\n\n"
        message += "⚠️ HIGH DETECTION PROBABILITY IF:\n"
        message += "  • File is an executable (.exe, .dll)\n"
        message += "  • File is a script (.ps1, .bat, .vbs)\n"
        message += "  • Target is a system directory\n"
        message += "  • File matches AV/EDR signatures\n"
        message += "  • File access auditing is enabled\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,
        )
        
        file_id = taskData.args.get_arg("file")
        path = taskData.args.get_arg("path")
        
        # Get file information to extract filename
        try:
            file_search = await SendMythicRPCFileSearch(MythicRPCFileSearchMessage(
                AgentFileId=file_id,
                LimitByCallback=False,
                TaskID=taskData.Task.ID
            ))
            
            if file_search.Success and len(file_search.Files) > 0:
                file_info = file_search.Files[0]
                filename = file_info.Filename
                
                # Normalize path: replace double backslashes with single backslash
                if path:
                    while "\\\\" in path:
                        path = path.replace("\\\\", "\\")
                
                # If path is empty, use current directory + filename
                if not path or path.strip() == "":
                    path = ".\\" + filename
                # If path ends with backslash or forward slash, append filename
                elif path.endswith("\\") or path.endswith("/"):
                    path = path.rstrip("\\/") + "\\" + filename
                # If path doesn't end with a backslash, check if it's a directory
                # If it's a directory (or looks like one), append filename
                else:
                    import os
                    path_normalized = path.rstrip("\\/")
                    
                    # Get the last component of the path (filename or directory name)
                    path_basename = os.path.basename(path_normalized)
                    
                    # Heuristic: If the basename doesn't have a file extension, 
                    # or if it's just a drive letter (like "C:"), treat as directory
                    has_extension = '.' in path_basename
                    is_drive_letter = len(path_basename) == 2 and path_basename[1] == ':'
                    
                    # If no extension and not empty, likely a directory - append filename
                    # Also treat drive letters as directories
                    if is_drive_letter or (not has_extension and path_basename):
                        path = path_normalized + "\\" + filename
                    # Otherwise, assume it's already a full file path and use as-is
                
                response.DisplayParams = f"file_id: {file_id}, path: {path}, filename: {filename}"
                # Update the path argument with the final path (including filename if needed)
                taskData.args.add_arg("path", path)
            else:
                response.DisplayParams = f"file_id: {file_id}, path: {path}"
        except Exception as e:
            response.DisplayParams = f"file_id: {file_id}, path: {path} (warning: could not get filename: {str(e)})"
        
        return response

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