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

class CpArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="fichero existente", 
                type=ParameterType.String, 
                description="Ruta del fichero a copiar",
                parameter_group_info=[ParameterGroupInfo(
                    required=True,
                    ui_position=1
                )]
            ),
            CommandParameter(
                name="fichero nuevo", 
                type=ParameterType.String, 
                description="Ruta del nuevo fichero",
                parameter_group_info=[ParameterGroupInfo(
                    required=True,
                    ui_position=2
                )]
            )
        ]

    async def parse_arguments(self):
        if len(self.command_line) == 0:
            raise ValueError("Este comando necesita parametros")
        self.add_arg("command", self.command_line)

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

class CpCommand(CommandBase):
    cmd = "cp"
    needs_admin = False
    help_cmd = "cp <source_path> <destination_path>"
    description = "Copy a file from a source location to a destination path on the target system. Supports both absolute and relative paths (relative paths are resolved using the current working directory). The destination file will be overwritten if it already exists. Creates a 'File Write' artifact for the destination file. Useful for creating backups, moving files, or preparing payloads in staging directories."
    version = 1
    author = "Kaseya OFSTeam"
    attackmapping = []
    argument_class = CpArguments
    attributes = CommandAttributes(
        builtin=False,
        supported_os=[ SupportedOS.Windows ],
        suggested_command=False
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about file copy risks, especially for sensitive files.
        """
        source_path = taskData.args.get_arg("fichero existente") or ""
        dest_path = taskData.args.get_arg("fichero nuevo") 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",
        }
        
        # Sensitive locations
        sensitive_locations = {
            "system32": "System32 directory - Highly monitored",
            "syswow64": "SysWOW64 directory - Highly monitored",
            "windows": "Windows directory - Highly monitored",
        }
        
        warnings = []
        source_lower = source_path.lower()
        dest_lower = dest_path.lower()
        
        for pattern, reason in sensitive_patterns.items():
            if pattern in source_lower:
                warnings.append(f"🚨 Source contains {pattern}: {reason}")
            if pattern in dest_lower:
                warnings.append(f"🚨 Destination contains {pattern}: {reason}")
        
        for location, reason in sensitive_locations.items():
            if location in source_lower:
                warnings.append(f"⚠️ Source: {location} - {reason}")
            if location in dest_lower:
                warnings.append(f"⚠️ Destination: {location} - {reason}")
        
        message = "⚠️ OPSEC WARNING - File Copy Detected\n\n"
        message += "File copying is monitored by:\n"
        message += "  • EDR/XDR solutions (file operations monitoring)\n"
        message += "  • Data loss prevention (DLP) systems\n"
        message += "  • File access auditing (if enabled)\n"
        message += "  • File integrity monitoring\n\n"
        
        if source_path:
            message += f"Source File: {source_path}\n"
        if dest_path:
            message += f"Destination: {dest_path}\n"
        if source_path or dest_path:
            message += "\n"
        
        if warnings:
            message += "🚨 RISKS DETECTED:\n"
            for warning in warnings:
                message += f"  {warning}\n"
            message += "\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • File read/write events (if auditing enabled)\n"
        message += "  • Sensitive file patterns (DLP detection)\n"
        message += "  • File operations monitoring\n"
        message += "  • System file modification detection\n\n"
        message += "💡 CONSIDERATIONS:\n"
        message += "  • Copying sensitive files (credentials, databases) is monitored\n"
        message += "  • System directories are highly monitored\n"
        message += "  • File operations may be logged\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.
        """
        source_path = taskData.args.get_arg("fichero existente") or ""
        dest_path = taskData.args.get_arg("fichero nuevo") or ""
        
        message = "⚠️ OPSEC POST-CHECK - File Copy Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • File Read: Read operations on source file\n"
        message += "  • File Write: Write operations to destination\n"
        message += "  • File Copy: CopyFile API calls\n"
        message += "  • File Access Logs: Event logs (if auditing enabled)\n\n"
        
        if source_path:
            message += f"Source File: {source_path}\n"
        if dest_path:
            message += f"Destination: {dest_path}\n"
        if source_path or dest_path:
            message += "\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR file operation monitoring\n"
        message += "  • DLP systems (sensitive file detection)\n"
        message += "  • File access auditing (if enabled)\n"
        message += "  • File integrity monitoring\n\n"
        message += "⚠️ HIGH DETECTION PROBABILITY IF:\n"
        message += "  • Source or destination contains sensitive files\n"
        message += "  • Copying to/from system directories\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