from mythic_container.MythicCommandBase import *
import json
import logging


class PsArguments(TaskArguments):

    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = []

    async def parse_arguments(self):
        if len(self.command_line) > 0:
            raise Exception("ps no debe lanzarse con parametros")


class PsCommand(CommandBase):
    cmd = "ps"
    needs_admin = False
    help_cmd = "ps"
    description = "List all running processes on the target system. Displays detailed process information including PID, process name, parent PID, architecture (x64/x86), username, and session ID. Automatically updates the Mythic Process Browser, allowing interactive process management (kill, steal_token, list_tokens) directly from the UI. The Process Browser cache is automatically cleared to show only currently running processes."
    version = 1
    supported_ui_features = ["process_browser:list"]
    author = "Kaseya OFSTeam"
    argument_class = PsArguments
    browser_script = BrowserScript(
        script_name="ps", author="Kaseya OFSTeam", for_new_ui=True
    )
    attributes = CommandAttributes(
        builtin=False,
        supported_os=[ SupportedOS.Windows ],
        suggested_command=False
    )
    attackmapping = []

    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:
        """
        Process response from agent. The translator automatically converts 
        tab-separated process list to Process Browser JSON format.
        We also ensure update_deleted=True is set for all processes to clear the cache.
        """
        resp = PTTaskProcessResponseMessageResponse(TaskID=task.Task.ID, Success=True)
        
        # Merge with existing response if present
        if response and isinstance(response, dict):
            resp.ProcessResponse = response
            
            # Ensure processes have update_deleted=True to clear cache
            if "processes" in resp.ProcessResponse and isinstance(resp.ProcessResponse["processes"], list):
                for process in resp.ProcessResponse["processes"]:
                    if isinstance(process, dict):
                        process["update_deleted"] = True
                logging.info(f"[PROCESS_BROWSER] Set update_deleted=True for {len(resp.ProcessResponse['processes'])} processes in process_response")
        
        return resp