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

logging.basicConfig(level=logging.INFO)


class RpfwdArguments(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 del agente donde escuchar (solo para start)",
                parameter_group_info=[ParameterGroupInfo(required=False, ui_position=2)]
            ),
            CommandParameter(
                name="remote_host",
                type=ParameterType.String,
                description="IP o hostname remoto al que Mythic debe conectar (solo para start)",
                parameter_group_info=[ParameterGroupInfo(required=False, ui_position=3)]
            ),
            CommandParameter(
                name="remote_port",
                type=ParameterType.Number,
                description="Puerto remoto al que Mythic debe conectar (solo para start)",
                parameter_group_info=[ParameterGroupInfo(required=False, ui_position=4)]
            ),
        ]

    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 local (opcional)
        local_port = int(parts[1]) if len(parts) > 1 else 8080
        self.add_arg("port", local_port)
        
        # Parsear host remoto y puerto (opcionales)
        remote_host = parts[2] if len(parts) > 2 else "127.0.0.1"
        remote_port = int(parts[3]) if len(parts) > 3 else 80
        self.add_arg("remote_host", remote_host)
        self.add_arg("remote_port", remote_port)
        
        self.add_arg("action", action)

    async def parse_dictionary(self, dictionary_arguments):
        action = dictionary_arguments.get("action", "start")
        # Valores por defecto
        if "port" not in dictionary_arguments:
            dictionary_arguments["port"] = 8080
        if "remote_host" not in dictionary_arguments:
            dictionary_arguments["remote_host"] = "127.0.0.1"
        if "remote_port" not in dictionary_arguments:
            dictionary_arguments["remote_port"] = 80
        self.load_args_from_dictionary(dictionary_arguments)


class RpfwdCommand(CommandBase):
    cmd = "rpfwd"
    needs_admin = False
    help_cmd = "rpfwd start [local_port] [remote_host] [remote_port] / rpfwd stop"
    description = "Start or stop reverse port forwarding (RPFWD). When started, the agent listens on a local port on the target system. Connections to this local port are forwarded through the agent to Mythic, which then connects to a remote destination (remote_host:remote_port). Default local port: 8080. Useful for accessing services on the target network that are not directly reachable from Mythic (e.g., internal services, databases, management interfaces). WARNING: Reverse port forwarding creates persistent network connections and may be detected by network monitoring, firewall logging, NIDS, and EDR/XDR network monitoring. Avoid using privileged ports (<1024) which require elevated privileges."
    version = 1
    author = "OFSTeam"
    argument_class = RpfwdArguments
    attackmapping = []
    attributes = CommandAttributes(
        supported_os=[SupportedOS.Windows],
    )

    async def opsec_pre(self, taskData: PTTaskMessageAllData) -> PTTTaskOPSECPreTaskMessageResponse:
        """
        OPSEC check before creating the task.
        Warns about reverse port forwarding risks.
        """
        action = taskData.args.get_arg("action") or "start"
        local_port = taskData.args.get_arg("port") or 8080
        remote_host = taskData.args.get_arg("remote_host") or "127.0.0.1"
        remote_port = taskData.args.get_arg("remote_port") or 80
        
        if action == "stop":
            # Stop action is low risk
            return PTTTaskOPSECPreTaskMessageResponse(
                TaskID=taskData.Task.ID,
                Success=True,
                OpsecPreBlocked=False,
                OpsecPreBypassRole="operator",
                OpsecPreMessage="Stopping reverse port forwarding - low OPSEC risk."
            )
        
        message = "⚠️ OPSEC WARNING - Reverse Port Forwarding Detected\n\n"
        message += "Reverse port forwarding 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 += "  • Network flow analysis tools\n\n"
        
        message += f"Configuration:\n"
        message += f"  • Local Port (Agent): {local_port}\n"
        message += f"  • Remote Host: {remote_host}\n"
        message += f"  • Remote Port: {remote_port}\n\n"
        
        # Check for privileged ports
        if local_port < 1024:
            message += "⚠️ WARNING: Privileged port (<1024) requires administrator privileges!\n"
            message += "This may trigger additional security alerts.\n\n"
        
        message += "🔍 DETECTION MECHANISMS:\n"
        message += "  • Network connection monitoring\n"
        message += "  • Port binding detection\n"
        message += "  • Suspicious network patterns\n"
        message += "  • Connection to non-standard ports\n\n"
        
        message += "💡 RECOMMENDATIONS:\n"
        message += "  • Use non-standard ports (avoid 80, 443, 22, 3389)\n"
        message += "  • Use ports commonly used by legitimate software\n"
        message += "  • Network traffic may be logged\n"
        message += "  • Firewall rules may block the connection\n\n"
        
        message += "⚠️ DETECTION RISK: MODERATE\n"
        message += "Network activity is visible but may blend in with normal traffic if configured properly.\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 reverse port forwarding - no additional artifacts."
            )
        
        local_port = taskData.args.get_arg("port") or 8080
        remote_host = taskData.args.get_arg("remote_host") or "127.0.0.1"
        remote_port = taskData.args.get_arg("remote_port") or 80
        
        message = "⚠️ OPSEC POST-CHECK - Reverse Port Forwarding Artifacts\n\n"
        message += "This task will create the following artifacts:\n"
        message += "  • Network Connection: Listener on agent port\n"
        message += "  • Port Binding: Socket bind on local port\n"
        message += "  • Network Activity: Connection forwarding\n"
        message += "  • Process Network Activity: Agent network connections\n\n"
        message += f"Configuration:\n"
        message += f"  • Local Port: {local_port}\n"
        message += f"  • Remote: {remote_host}:{remote_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\n"
        
        message += "⚠️ HIGH DETECTION PROBABILITY IF:\n"
        message += "  • Using well-known ports (80, 443, 22, etc.)\n"
        message += "  • Network traffic is heavily monitored\n"
        message += "  • Firewall rules are strict\n"
        message += "  • Using privileged ports (<1024)\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")
        if action == "start":
            local_port = task.args.get_arg("port")
            remote_host = task.args.get_arg("remote_host")
            remote_port = task.args.get_arg("remote_port")
            task.display_params = f"{action.upper()} (local:{local_port} -> {remote_host}:{remote_port})"
        else:
            task.display_params = "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")
        local_port = taskData.args.get_arg("port") if taskData.args.get_arg("port") else 8080

        try:
            if action == "start":
                remote_host = taskData.args.get_arg("remote_host") or "127.0.0.1"
                remote_port = taskData.args.get_arg("remote_port") or 80
                
                # Warn if using a privileged port on Windows (may require admin)
                if local_port < 1024:
                    logging.warning(f"Warning: Port {local_port} is privileged (<1024) and may require administrator privileges on Windows")
                    logging.warning(f"Consider using a non-privileged port (>=1024) like 8080, 4444, or 5555")
                
                # Register RPFWD with Mythic - Mythic will automatically send start_rpfwd command to agent
                resp = await SendMythicRPCProxyStartCommand(MythicRPCProxyStartMessage(
                    TaskID=taskData.Task.ID,
                    LocalPort=local_port,
                    RemoteIP=remote_host,
                    RemotePort=remote_port,
                    PortType=CALLBACK_PORT_TYPE_RPORTFWD
                ))
                logging.info(f"RPFWD registered with Mythic: local port {local_port} on agent -> {remote_host}:{remote_port}")
                logging.info(f"To test: Connect to port {local_port} on the AGENT host (not Mythic server)")
                if not resp.Success:
                    response.Success = False
                    response.Error = resp.Error
            elif action == "stop":
                # Unregister RPFWD with Mythic - Mythic will automatically send stop_rpfwd command to agent
                resp = await SendMythicRPCProxyStopCommand(MythicRPCProxyStopMessage(
                    TaskID=taskData.Task.ID,
                    Port=local_port,
                    PortType=CALLBACK_PORT_TYPE_RPORTFWD
                ))
                logging.info(f"RPFWD unregistered from Mythic: {resp}")
                if not resp.Success:
                    response.Success = False
                    response.Error = resp.Error
        except Exception as e:
            logging.error(f"Error en RPFWD 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

