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


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

    async def parse_arguments(self):
        # Screenshot command takes no arguments
        pass

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


class ScreenshotCommand(CommandBase):
    cmd = "screenshot"
    needs_admin = False
    help_cmd = "screenshot"
    description = "Capture a screenshot of the primary desktop display and upload it to the Mythic server. Uses Windows GDI API (BitBlt) to capture the screen content. The screenshot is automatically uploaded as a file to Mythic for viewing. Requires an active interactive desktop session. May be detected by screen capture detection mechanisms, behavioral analysis, and data loss prevention (DLP) systems. Large screenshots may trigger network monitoring alerts."
    version = 1
    author = "@KaseyaOFSTeam"
    argument_class = ScreenshotArguments
    attackmapping = ["T1113"]
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about screenshot capture risks.
        """
        message = "⚠️ OPSEC WARNING - Screenshot Capture Detected\n\n"
        message += "Screenshot capture is monitored by:\n"
        message += "  • EDR/XDR solutions (behavioral detection)\n"
        message += "  • Screen capture detection mechanisms\n"
        message += "  • Data loss prevention (DLP) systems\n"
        message += "  • User activity monitoring tools\n\n"
        
        message += "🔍 DETECTION MECHANISMS:\n"
        message += "  • BitBlt/GDI screen capture API calls\n"
        message += "  • PrintWindow API calls\n"
        message += "  • Suspicious screen capture patterns\n"
        message += "  • Large data uploads (screenshot files)\n\n"
        
        message += "💡 CONSIDERATIONS:\n"
        message += "  • Screenshots may be large and trigger network alerts\n"
        message += "  • Multiple screenshots may be suspicious\n"
        message += "  • Screen capture events may be logged\n"
        message += "  • Sensitive data in screenshots may trigger DLP alerts\n\n"
        
        message += "⚠️ DETECTION RISK: MODERATE\n"
        message += "Usually not immediately detected, but may be flagged by behavioral analysis over time.\n\n"
        message += "✅ This is a WARNING only - command will proceed.\n"
        message += "Ensure you understand the detection risks before continuing."
        
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=False,  # Moderate risk, don't block
            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 - Screenshot Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Screen Capture: BitBlt/PrintWindow API calls\n"
        message += "  • File Upload: Screenshot image file transmission\n"
        message += "  • Network Activity: Image data transfer\n"
        message += "  • Memory Artifacts: Screenshot data in memory\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • EDR/XDR behavioral detection\n"
        message += "  • Network monitoring tools (large file uploads)\n"
        message += "  • DLP systems (sensitive data capture)\n"
        message += "  • User activity monitoring\n\n"
        
        message += "⚠️ HIGH DETECTION PROBABILITY IF:\n"
        message += "  • Multiple screenshots are taken rapidly\n"
        message += "  • Screenshots contain sensitive data\n"
        message += "  • Network traffic is monitored\n"
        message += "  • Behavioral analysis is aggressive\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 = "Capture screenshot"
        return response

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

