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


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

    async def parse_arguments(self):
        pass

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


class WhoamiCommand(CommandBase):
    cmd = "whoami"
    needs_admin = False
    help_cmd = "whoami"
    description = "Display the current security context of the agent. Shows the username and domain of the active security token. If token impersonation is active (via 'steal_token' or 'make_token'), displays the impersonated user context. Otherwise, displays the process token context. Useful for verifying that token operations succeeded and for confirming the current privilege level before executing commands that require specific permissions."
    version = 1
    author = "@KaseyaOFSTeam"
    attackmapping = ["T1083"]
    argument_class = WhoamiArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def create_go_tasking(self, taskData: PTTaskMessageAllData) -> PTTaskCreateTaskingMessageResponse:
        response = PTTaskCreateTaskingMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
        )
        response.DisplayParams = "Getting current user context..."
        return response

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

