from mythic_container.MythicCommandBase import *
from mythic_container.MythicRPC import (
    SendMythicRPCProxyStartCommand,
    MythicRPCProxyStartMessage,
    SendMythicRPCProxyStopCommand,
    MythicRPCProxyStopMessage,
    CALLBACK_PORT_TYPE_SOCKS
)
import json
import logging

logging.basicConfig(level=logging.INFO)


class SocksArguments(TaskArguments):
    def __init__(self, command_line, **kwargs):
        super().__init__(command_line, **kwargs)
        self.args = [
            CommandParameter(
                name="action",
                type=ParameterType.ChooseOne,
                description="Acción a realizar (start/stop)",
                choices=["start", "stop"],
                default_value="start",
                parameter_group_info=[ParameterGroupInfo(required=True, ui_position=1)]
            ),
            CommandParameter(
                name="port",
                type=ParameterType.Number,
                description="Puerto local a exponer (solo para start)",
                parameter_group_info=[ParameterGroupInfo(required=False, ui_position=2)]
            ),
        ]

    async def parse_arguments(self):
        if len(self.command_line) == 0:
            raise ValueError("Debe especificar 'start' o 'stop'.")
        parts = self.command_line.strip().split()
        action = parts[0].lower()
        
        if action not in ["start", "stop"]:
            raise ValueError("Acción debe ser 'start' o 'stop'.")
        
        # Parsear puerto (opcional para ambos)
        # Si no se especifica, usar 7002 como default (puerto común para SOCKS en Mythic)
        port = int(parts[1]) if len(parts) > 1 else 7002
        self.add_arg("port", port)
        self.add_arg("action", action)

    async def parse_dictionary(self, dictionary_arguments):
        action = dictionary_arguments.get("action", "start")
        # Si no hay puerto, usar 7002 como default
        if "port" not in dictionary_arguments:
            dictionary_arguments["port"] = 7002
        self.load_args_from_dictionary(dictionary_arguments)


class SocksCommand(CommandBase):
    cmd = "socks"
    needs_admin = False
    help_cmd = "socks start [port] / socks stop [port]"
    description = "Start or stop a SOCKS5 proxy server on the Mythic server. When started, the SOCKS proxy allows routing network traffic through the agent to access internal network resources. The default port is 7002. The proxy runs on the Mythic server and uses the agent as a tunnel for network connections. WARNING: SOCKS proxies generate continuous network traffic and may be detected by network monitoring tools, firewall logging, EDR/XDR network monitoring, and network flow analysis. Use non-standard ports and be aware of high bandwidth usage."
    version = 1
    author = "OFSTeam"
    argument_class = SocksArguments
    attackmapping = []
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows],
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about SOCKS proxy risks.
        """
        action = taskData.args.get_arg("action") or "start"
        port = taskData.args.get_arg("port") or 7002
        
        if action == "stop":
            # Stop action is low risk
            return PTTTaskOPSECPreTaskMessageResponse(
                TaskID=taskData.Task.ID,
                Success=True,
                OpsecPreBlocked=False,
                OpsecPreBypassRole="operator",
                OpsecPreMessage="Stopping SOCKS proxy - low OPSEC risk."
            )
        
        message = "⚠️ OPSEC WARNING - SOCKS Proxy Detected\n\n"
        message += "SOCKS proxy is monitored by:\n"
        message += "  • Network monitoring tools\n"
        message += "  • Firewall logging\n"
        message += "  • Network intrusion detection systems (NIDS)\n"
        message += "  • EDR/XDR network monitoring\n"
        message += "  • Proxy detection mechanisms\n"
        message += "  • Network flow analysis tools\n\n"
        
        message += f"Configuration:\n"
        message += f"  • SOCKS Port: {port}\n\n"
        
        message += "🔍 DETECTION MECHANISMS:\n"
        message += "  • Network connection monitoring\n"
        message += "  • Proxy protocol detection\n"
        message += "  • High-volume network traffic patterns\n"
        message += "  • Suspicious network behavior\n"
        message += "  • Connection to non-standard ports\n\n"
        
        message += "⚠️ HIGH TRAFFIC WARNING:\n"
        message += "  • SOCKS proxies generate continuous network activity\n"
        message += "  • High bandwidth usage may trigger alerts\n"
        message += "  • Network patterns may be analyzed for anomalies\n"
        message += "  • All traffic through proxy is visible\n\n"
        
        message += "💡 CONSIDERATIONS:\n"
        message += "  • SOCKS proxies generate significant network traffic\n"
        message += "  • Using common proxy ports (1080, 8080) may blend better\n"
        message += "  • Firewall rules may block the connection\n"
        message += "  • Network patterns may be analyzed\n\n"
        
        message += "⚠️ DETECTION RISK: MODERATE-HIGH\n"
        message += "Network activity is highly visible and proxy traffic patterns may be detected by network analysis.\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,
            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.
        """
        action = taskData.args.get_arg("action") or "start"
        
        if action == "stop":
            return PTTTaskOPSECPostTaskMessageResponse(
                TaskID=taskData.Task.ID,
                Success=True,
                OpsecPostBlocked=False,
                OpsecPostBypassRole="operator",
                OpsecPostMessage="Stopping SOCKS proxy - no additional artifacts."
            )
        
        port = taskData.args.get_arg("port") or 7002
        
        message = "⚠️ OPSEC POST-CHECK - SOCKS Proxy Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Network Connection: SOCKS proxy listener\n"
        message += "  • Port Binding: Socket bind on SOCKS port\n"
        message += "  • Network Activity: Continuous proxy traffic\n"
        message += "  • Process Network Activity: Agent network connections\n"
        message += "  • High Bandwidth Usage: Proxy data transfer\n\n"
        message += f"Configuration:\n"
        message += f"  • SOCKS Port: {port}\n\n"
        
        message += "🔍 DETECTION RISKS:\n"
        message += "  • Network monitoring tools\n"
        message += "  • Firewall logging\n"
        message += "  • NIDS/NIPS systems\n"
        message += "  • Network flow analysis\n"
        message += "  • Bandwidth monitoring systems\n\n"
        
        message += "⚠️ HIGH DETECTION PROBABILITY IF:\n"
        message += "  • High volume of traffic is generated\n"
        message += "  • Network traffic patterns are analyzed\n"
        message += "  • SOCKS protocol is detected\n"
        message += "  • Firewall rules are strict\n"
        message += "  • Network monitoring is aggressive\n\n"
        
        message += "⚠️ WARNING: SOCKS proxies generate significant network activity\n"
        message += "that may be detected by network security tools.\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_tasking(self, task: MythicTask) -> MythicTask:
        action = task.args.get_arg("action")
        port = task.args.get_arg("port")
        task.display_params = f"{action.upper()} (puerto {port})" if action == "start" else "STOP"
        return task

    async def create_go_tasking(self, taskData: PTTaskMessageAllData) -> PTTaskCreateTaskingMessageResponse:
        response = PTTaskCreateTaskingMessageResponse(TaskID=taskData.Task.ID, Success=True)

        action = taskData.args.get_arg("action")
        port = taskData.args.get_arg("port") if taskData.args.get_arg("port") else 7002

        try:
            if action == "start":
                resp = await SendMythicRPCProxyStartCommand(MythicRPCProxyStartMessage(
                    TaskID=taskData.Task.ID,
                    LocalPort=port,
                    PortType=CALLBACK_PORT_TYPE_SOCKS
                ))
                logging.info(f"SOCKS proxy started on port {port}: {resp}")
                if not resp.Success:
                    response.Success = False
                    response.Error = resp.Error
            elif action == "stop":
                # MythicRPCProxyStopMessage también requiere Port y PortType
                resp = await SendMythicRPCProxyStopCommand(MythicRPCProxyStopMessage(
                    TaskID=taskData.Task.ID,
                    Port=port,
                    PortType=CALLBACK_PORT_TYPE_SOCKS
                ))
                logging.info(f"SOCKS proxy stopped: {resp}")
                if not resp.Success:
                    response.Success = False
                    response.Error = resp.Error
        except Exception as e:
            logging.error(f"Error en SOCKS RPC: {e}")
            response.Success = False
            response.Error = str(e)
        return response

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