#include "CertPinner.h"
#include <wincrypt.h>
#include <cstdio>
#include <cstring>
#include <vector>

#pragma comment(lib, "crypt32.lib")

// ─── XOR-obfuscated pinned SPKI SHA-256 hash ──────────
// For development/testing, set to all-zeros to skip pinning.
// In production, replace with the actual SPKI hash of your server's public key.
//
// To generate: openssl x509 -in cert.pem -pubkey -noout | openssl pkey -pubin -outform der | openssl dgst -sha256

static const unsigned char OBFUSCATED_PIN[] = {
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
};
static const unsigned char XOR_KEY = 0x00;  // dev mode: XOR with 0 = no change

std::string CertPinner::GetPinnedHash()
{
    std::string result(sizeof(OBFUSCATED_PIN), '\0');
    for (size_t i = 0; i < sizeof(OBFUSCATED_PIN); ++i)
        result[i] = OBFUSCATED_PIN[i] ^ XOR_KEY;
    return result;
}

// ─── Public Key Hash Extraction ───────────────────────

std::string CertPinner::GetSpkiHash(PCCERT_CONTEXT pCertContext)
{
    if (!pCertContext)
        return {};

    // Get the SubjectPublicKeyInfo from the certificate
    PCERT_PUBLIC_KEY_INFO pKeyInfo = &pCertContext->pCertInfo->SubjectPublicKeyInfo;

    // Hash the raw public key bytes with SHA-256
    // Note: For proper SPKI pinning you'd encode the full SPKI (algorithm + key) as DER.
    // For production, use: CryptEncodeObjectEx(X509_ASN_ENCODING, X509_PUBLIC_KEY_INFO, ...)
    // then hash the DER output. See OpenSSL command in the header comment.
    HCRYPTPROV hProv = 0;
    HCRYPTHASH hHash = 0;
    std::string result;

    if (CryptAcquireContextA(&hProv, nullptr, nullptr, PROV_RSA_AES, CRYPT_VERIFYCONTEXT))
    {
        if (CryptCreateHash(hProv, CALG_SHA_256, 0, 0, &hHash))
        {
            if (CryptHashData(hHash, pKeyInfo->PublicKey.pbData,
                              pKeyInfo->PublicKey.cbData, 0))
            {
                BYTE hash[32];
                DWORD hashLen = sizeof(hash);
                if (CryptGetHashParam(hHash, HP_HASHVAL, hash, &hashLen, 0))
                {
                    result.assign(reinterpret_cast<char*>(hash), hashLen);
                }
            }
            CryptDestroyHash(hHash);
        }
        CryptReleaseContext(hProv, 0);
    }

    return result;
}

// ─── Public API ───────────────────────────────────────

bool CertPinner::VerifyServerCert(HINTERNET hRequest)
{
    const std::string pinned = GetPinnedHash();

    // Check if all zeros — development mode, skip pinning
    bool devMode = true;
    for (char c : pinned) {
        if (c != '\0') { devMode = false; break; }
    }
    if (devMode) {
        printf("[CertPinner] DEV MODE — pinning skipped (all-zeros hash)\n");
        return true;
    }

    // Retrieve server certificate from WinHTTP
    PCCERT_CONTEXT pCertContext = nullptr;
    DWORD certSize = sizeof(pCertContext);
    if (!WinHttpQueryOption(
            hRequest,
            WINHTTP_OPTION_SERVER_CERT_CONTEXT,
            &pCertContext,
            &certSize))
    {
        printf("[CertPinner] WinHttpQueryOption(SERVER_CERT_CONTEXT) failed: 0x%lx\n", GetLastError());
        return false;
    }

    // Verify certificate chain
    CERT_CHAIN_PARA chainPara = {};
    chainPara.cbSize = sizeof(chainPara);
    PCCERT_CHAIN_CONTEXT pChainContext = nullptr;
    if (!CertGetCertificateChain(
            nullptr,
            pCertContext,
            nullptr,
            pCertContext->hCertStore,
            &chainPara,
            0,
            nullptr,
            &pChainContext))
    {
        printf("[CertPinner] CertGetCertificateChain failed: 0x%lx\n", GetLastError());
        CertFreeCertificateContext(pCertContext);
        return false;
    }

    // Hash the server's public key
    std::string serverHash = GetSpkiHash(pCertContext);

    CertFreeCertificateChain(pChainContext);
    CertFreeCertificateContext(pCertContext);

    if (serverHash.empty())
    {
        printf("[CertPinner] ERROR: failed to hash server public key\n");
        return false;
    }

    // Compare
    if (serverHash == pinned)
    {
        printf("[CertPinner] SPKI hash MATCH — connection trusted\n");
        return true;
    }

    printf("[CertPinner] SPKI hash MISMATCH — possible MITM detected!\n");
    return false;
}
