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


class ExitArguments(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("Exit command takes no parameters.")


class ExitCommand(CommandBase):
    cmd = "exit"
    needs_admin = False
    help_cmd = "exit"
    description = "Instruct the agent to exit and terminate. The agent will perform cleanup operations and then terminate its process. This command is permanent - the agent will stop running and will not check in again. Use this command when you want to remove the agent from the target system or when operational requirements dictate agent termination. Compatible with Mythic's callback table exit functionality."
    version = 1
    is_exit = True  # Important: marks this as an exit command (like Athena)
    supported_ui_features = ["callback_table:exit"]
    author = "Kaseya OFSTeam"
    argument_class = ExitArguments
    attackmapping = []
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows],
        builtin=True
    )

    async def create_go_tasking(self, taskData: PTTaskMessageAllData) -> PTTaskCreateTaskingMessageResponse:
        response = PTTaskCreateTaskingMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
        )
        # Don't mark as completed immediately - let the agent send the response
        # This follows Xenon and Athena's pattern
        return response

    async def process_response(self, task: PTTaskMessageAllData, response: any) -> PTTaskProcessResponseMessageResponse:
        resp = PTTaskProcessResponseMessageResponse(TaskID=task.Task.ID, Success=True)
        # Mark task as completed explicitly when we receive the response
        resp.Completed = True
        return resp