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


class ShellArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="command", 
                type=ParameterType.String, 
                description="Comando a ejecutar",
                parameter_group_info=[ParameterGroupInfo(
                    required=True
                )]
            ),
        ]

    async def parse_arguments(self):
        if len(self.command_line) == 0:
            raise ValueError("Se debe indicar un comando")
        self.add_arg("command", self.command_line)

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


class ShellCommand(CommandBase):
    cmd = "shell"
    needs_admin = False
    help_cmd = "shell <command>"
    description = "Execute arbitrary shell commands on the target system by spawning cmd.exe. The command is executed in a non-interactive shell and the output is captured and returned. HIGH OPSEC RISK: This command spawns cmd.exe which is heavily monitored by EDR/XDR solutions. Command execution and command-line arguments are logged by security tools. Consider using built-in Angerona commands (ps, ls, cat, etc.) instead when possible, as they provide better OPSEC. Always requires approval from another operator before execution."
    version = 1
    author = "Kaseya OFSTeam"
    attackmapping = ["T1059"]
    argument_class = ShellArguments
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows]
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about dangerous commands that are highly detectable.
        """
        command = (taskData.args.get_arg("command") or "").lower()
        
        # Highly detectable commands
        dangerous_commands = {
            "powershell": "PowerShell execution is heavily monitored by EDR/XDR",
            "pwsh": "PowerShell Core execution is heavily monitored",
            "wmic": "WMI queries are logged and monitored",
            "net.exe": "Network enumeration commands trigger alerts",
            "nslookup": "DNS queries may be monitored",
            "systeminfo": "System information gathering is logged",
            "whoami": "Identity queries are logged",
            "tasklist": "Process enumeration is monitored",
            "netstat": "Network connection enumeration triggers alerts",
            "ipconfig": "Network configuration queries are logged",
        }
        
        blocked = False
        warnings = []
        
        for dangerous_cmd, reason in dangerous_commands.items():
            if dangerous_cmd in command:
                warnings.append(f"⚠️ {dangerous_cmd}: {reason}")
                blocked = True
        
        # Check for cmd.exe spawning (always present in shell command)
        warnings.append("⚠️ cmd.exe spawn: This command spawns cmd.exe which is highly detectable by EDR/XDR")
        
        message = "⚠️ HIGH OPSEC RISK - Shell Command Detected\n\n"
        message += "This command will spawn cmd.exe which is heavily monitored.\n\n"
        
        if warnings:
            message += "Additional risks detected:\n"
            for warning in warnings:
                message += f"  • {warning}\n"
            message += "\n"
        
        message += "Consider using built-in Angerona commands instead:\n"
        message += "  • Use 'ps' instead of 'tasklist'\n"
        message += "  • Use 'ls' instead of 'dir'\n"
        message += "  • Use 'cat' instead of 'type'\n"
        message += "  • Use 'whoami' command instead of 'whoami.exe'\n\n"
        message += "If you must proceed, ensure you understand the detection risks."
        
        return PTTTaskOPSECPreTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPreBlocked=blocked,
            OpsecPreBypassRole="other_operator",  # Requires another operator to approve
            OpsecPreMessage=message
        )

    async def opsec_post(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPostTaskMessageResponse:
        """
        OPSEC check after creating the task.
        Warns about artifacts that will be created.
        """
        command = taskData.args.get_arg("command") or ""
        
        message = "⚠️ OPSEC POST-CHECK - Shell Command Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Process Create: cmd.exe spawned\n"
        message += "  • Command Line: The executed command will be logged\n"
        message += "  • Network Activity: If command makes network connections\n\n"
        message += "These artifacts may be detected by:\n"
        message += "  • EDR/XDR solutions\n"
        message += "  • SIEM systems\n"
        message += "  • Process monitoring tools\n"
        message += "  • Command line auditing\n\n"
        message += "Ensure you understand the detection risks before proceeding."
        
        return PTTTaskOPSECPostTaskMessageResponse(
            TaskID=taskData.Task.ID,
            Success=True,
            OpsecPostBlocked=False,  # Don't block, just warn
            OpsecPostBypassRole="operator",  # Any operator can bypass
            OpsecPostMessage=message
        )

    async def create_tasking(self, task: MythicTask) -> MythicTask:
        task.display_params = task.args.get_arg("command")
        return task

    async def process_response(self, task: PTTaskMessageAllData, response: any) -> PTTaskProcessResponseMessageResponse:
        """
        Process response from agent and add artifacts for Process Create.
        
        According to Mythic documentation:
        https://docs.mythic-c2.net/customizing/hooking-features/artifacts
        
        Artifacts should be added to the response JSON with an 'artifacts' array.
        However, the current Mythic API doesn't allow commands to directly modify
        the translator's JSON response. Instead, we could:
        
        1. Use MythicRPC to create artifacts after the response (current limitation)
        2. Modify the agent to send command context in responses
        3. Have the translator detect shell commands automatically
        
        For now, artifacts are prepared in the translator but need command context.
        """
        resp = PTTaskProcessResponseMessageResponse(TaskID=task.Task.ID, Success=True)
        
        # TODO: Add artifacts via MythicRPC or modify translator to detect shell commands
        # command = task.args.get_arg("command")
        # if command:
        #     # Could use MythicRPC CreateArtifact here if available
        
        return resp