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


class Rev2SelfArguments(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 Rev2SelfCommand(CommandBase):
    cmd = "rev2self"
    needs_admin = False
    help_cmd = "rev2self"
    description = "Revert token impersonation back to the original process token. This command stops impersonating any stolen or created token and returns to the agent's original security context. The impersonation context in the Mythic UI is automatically cleared. Use this command after completing operations that required elevated privileges to return to a less-privileged context and reduce detection risk."
    version = 1
    author = "@KaseyaOFSTeam"
    attackmapping = ["T1134.001"]
    argument_class = Rev2SelfArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about token reversion risks.
        """
        message = "⚠️ OPSEC WARNING - Token Reversion Detected\n\n"
        message += "Reverting token impersonation is monitored by:\n"
        message += "  • EDR/XDR solutions (behavioral detection)\n"
        message += "  • Security auditing systems\n"
        message += "  • Token manipulation monitoring\n\n"
        
        message += "🔍 DETECTION MECHANISMS:\n"
        message += "  • Token reversion API calls\n"
        message += "  • Context change detection\n"
        message += "  • Privilege level changes\n"
        message += "  • Security event logging\n\n"
        
        message += "💡 CONSIDERATIONS:\n"
        message += "  • Reverting from elevated token may reduce privileges\n"
        message += "  • Context tracking will update in Mythic UI\n"
        message += "  • Lower detection risk than token theft/creation\n\n"
        
        message += "✅ This is a WARNING only - command will proceed.\n"
        message += "Ensure you understand the privilege implications."
        
        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.
        """
        message = "⚠️ OPSEC POST-CHECK - Token Reversion Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Token Reversion: RevertToSelf API call\n"
        message += "  • Context Change: Privilege level change\n"
        message += "  • Context Tracking: Impersonation context cleared in Mythic UI\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR behavioral detection (low-moderate confidence)\n"
        message += "  • Security event logging (if enabled)\n"
        message += "  • Token manipulation monitoring\n\n"
        message += "⚠️ DETECTION PROBABILITY:\n"
        message += "  • Lower than token theft/creation operations\n"
        message += "  • Context changes may be logged\n"
        message += "  • Privilege changes may trigger alerts\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,
        )
        response.DisplayParams = "Reverting to original token..."
        return response

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

