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

logging.basicConfig(level=logging.INFO)

class RmArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="path", 
                type=ParameterType.String, 
                description="ruta al fichero o directorio",
                parameter_group_info=[ParameterGroupInfo(
                    required=True
                )]
            ),
        ]

    async def parse_arguments(self):
        if len(self.command_line) == 0:
            raise ValueError("Debe introducir una ruta a un fichero o directorio")
        self.add_arg("path", self.command_line)

    async def parse_dictionary(self, dictionary):
        logging.info(f"[rm] Parse Dictionary - {dictionary}")
        # Load args first (will set path from dictionary)
        self.load_args_from_dictionary(dictionary)
        
        # Then override with processed values if coming from File Browser
        if "host" in dictionary:
            # File Browser format
            logging.info(f"[rm] Command came from File Browser UI - {dictionary}")
            full_path = dictionary.get('full_path', '')
            if full_path:
                self.add_arg("path", full_path)
                logging.info(f"[rm] Using full_path: {full_path}")
            else:
                path = dictionary.get("path", '')
                file = dictionary.get("file", '')
                if path and file:
                    # Combine path and file
                    if not path.endswith('\\') and not path.endswith('/'):
                        path = path + "\\"
                    combined_path = path + file
                    self.add_arg("path", combined_path)
                    logging.info(f"[rm] Combined path + file: {combined_path}")
                elif path:
                    self.add_arg("path", path)
                    logging.info(f"[rm] Using path only: {path}")
                else:
                    raise ValueError("[rm] File Browser request missing both path and file")
            self.add_arg("host", dictionary.get("host", ""))
        else:
            # Command line usage - already loaded from dictionary, just verify
            logging.info(f"[rm] Command came from CMDLINE - {dictionary}")
            arg_path = self.get_arg("path")
            if not arg_path:
                arg_path = dictionary.get("path")
                if arg_path:
                    self.add_arg("path", arg_path)
                else:
                    raise ValueError("[rm] No path specified")

class RmCommand(CommandBase):
    cmd = "rm"
    needs_admin = False
    help_cmd = "rm <path>"
    description = "Delete a file or directory from the target system. Supports both absolute and relative paths (relative paths are resolved using the current working directory). Can delete individual files or entire directories. WARNING: Deletion of critical system files or directories (System32, Windows, Program Files, etc.) will be blocked for safety. Use with caution as deleted files cannot be recovered through this command."
    version = 1
    supported_ui_features = ["file_browser:remove"]
    author = "Kaseya OFSTeam"
    attackmapping = []
    argument_class = RmArguments
    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 deletion risks, especially for system files.
        """
        path = taskData.args.get_arg("path") or ""
        
        # Critical system locations
        critical_locations = {
            "system32": "System32 directory - CRITICAL system files",
            "syswow64": "SysWOW64 directory - CRITICAL system files",
            "windows": "Windows directory - CRITICAL system files",
            "program files": "Program Files - Application files",
            "boot": "Boot directory - System boot files",
            "winnt": "WinNT directory - System files",
        }
        
        # Critical file patterns
        critical_files = {
            ".dll": "DLL files - System libraries",
            ".sys": "System driver files",
            ".exe": "Executable files - System binaries",
        }
        
        blocked = False
        warnings = []
        critical_warnings = []
        path_lower = path.lower()
        
        for location, reason in critical_locations.items():
            if location in path_lower:
                critical_warnings.append(f"🚨 {location}: {reason}")
                if location in ["system32", "syswow64", "windows", "boot"]:
                    blocked = True
        
        for file_ext, reason in critical_files.items():
            if path_lower.endswith(file_ext) and any(loc in path_lower for loc in ["system32", "syswow64", "windows"]):
                warnings.append(f"🚨 {file_ext}: {reason} - System file deletion")
                blocked = True
        
        message = "⚠️ OPSEC RISK - File/Directory Deletion Detected\n\n"
        message += "File deletion is monitored by:\n"
        message += "  • EDR/XDR solutions (file deletion monitoring)\n"
        message += "  • Anti-malware solutions\n"
        message += "  • File access auditing (if enabled)\n"
        message += "  • System integrity monitoring\n\n"
        
        if path:
            message += f"Target Path: {path}\n\n"
        
        if critical_warnings:
            message += "🚨 CRITICAL LOCATION DETECTED:\n"
            for warning in critical_warnings:
                message += f"  {warning}\n"
            message += "\n"
        
        if warnings:
            message += "🚨 CRITICAL FILE DETECTED:\n"
            for warning in warnings:
                message += f"  {warning}\n"
            message += "\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • File deletion events (if auditing enabled)\n"
        message += "  • System file deletion alerts\n"
        message += "  • File integrity monitoring\n"
        message += "  • Suspicious deletion patterns\n\n"
        message += "💡 CONSIDERATIONS:\n"
        message += "  • Deleting system files may cause system instability\n"
        message += "  • System file deletion triggers immediate alerts\n"
        message += "  • File deletion is often logged\n"
        message += "  • Use with extreme caution in system directories\n\n"
        
        if blocked:
            message += "⚠️ TO PROCEED: You need approval from another operator.\n"
            message += "This command is blocked due to targeting critical system files."
        else:
            message += "⚠️ WARNING: File deletion may be logged.\n"
            message += "✅ This is a WARNING only - command will proceed."
        
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=blocked,
            OpsecPreBypassRole="other_operator" if blocked else "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 Deletion Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • File Deletion: DeleteFile/RemoveDirectory API calls\n"
        message += "  • File Access: File system operations\n"
        message += "  • File Deletion 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 deletion monitoring\n"
        message += "  • File access auditing (if enabled)\n"
        message += "  • System integrity monitoring\n"
        message += "  • File deletion event logging\n\n"
        message += "⚠️ HIGH DETECTION PROBABILITY IF:\n"
        message += "  • Target is a system file or directory\n"
        message += "  • File deletion auditing is enabled\n"
        message += "  • System integrity monitoring is active\n"
        message += "  • Multiple files are deleted rapidly\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)
        
        # Check if removal was successful by checking user_output
        if response:
            user_output = response.get("user_output", "")
            # If output doesn't contain error indicators, assume success
            if user_output and not any(err in user_output.lower() for err in ["error", "failed", "cannot", "cannot"]):
                # Get the path that was deleted
                path = task.args.get_arg("path")
                if path:
                    # Get hostname
                    host = task.Callback.Host if task.Callback else ""
                    
                    # Add removed_files array to response (merge with existing response)
                    resp.ProcessResponse = response if isinstance(response, dict) else {}
                    resp.ProcessResponse["removed_files"] = [
                        {
                            "host": host,
                            "path": path
                        }
                    ]
        
        return resp