#include <windows.h>
#include <wincrypt.h>
#include <winhttp.h>

#include <assert.h>
#include <stdbool.h>
#include <stdlib.h>

#include "hook/table.h"

#include "platform/cert.h"
#include "hook/procaddr.h"

#include "util/dprintf.h"
#include "util/str.h"

/*
 * The game uses a client cert to do mTLS. We can either bypass this if we're
 * connecting to a server that doesn't care, or splice in our own if it does.
 * For server TLS verification the client just uses whatever certs are trusted
 * on the PCBs root store, which we can't easily modify, but we can just set
 * security flags such that it ignores any issues with the server's cert. As
 * far as I can tell, no actual client-side validation of the client cert is
 * done, so ANY cert here should work.
*/

static CRITICAL_SECTION cert_lock;
static wchar_t path[MAX_PATH];
static HCERTSTORE p12store;
static PCCERT_CONTEXT client_cert_ctx = NULL;

static PCCERT_CONTEXT WINAPI my_CertFindCertificateInStore(
    HCERTSTORE     hCertStore,
    DWORD          dwCertEncodingType,
    DWORD          dwFindFlags,
    DWORD          dwFindType,
    const void     *pvFindPara,
    PCCERT_CONTEXT pPrevCertContext
);

static HCERTSTORE WINAPI my_CertOpenStore(
    LPCSTR            lpszStoreProvider,
    DWORD             dwEncodingType,
    HCRYPTPROV_LEGACY hCryptProv,
    DWORD             dwFlags,
    const void        *pvPara
);

static BOOL WINAPI my_WinHttpSetOption(
    HINTERNET hInternet,
    DWORD     dwOption,
    LPVOID    lpBuffer,
    DWORD     dwBufferLength
);

static PCCERT_CONTEXT (WINAPI *next_CertFindCertificateInStore)(
    HCERTSTORE     hCertStore,
    DWORD          dwCertEncodingType,
    DWORD          dwFindFlags,
    DWORD          dwFindType,
    const void     *pvFindPara,
    PCCERT_CONTEXT pPrevCertContext
);

static HCERTSTORE (WINAPI *next_CertOpenStore)(
    LPCSTR            lpszStoreProvider,
    DWORD             dwEncodingType,
    HCRYPTPROV_LEGACY hCryptProv,
    DWORD             dwFlags,
    const void        *pvPara
);

static BOOL (WINAPI *next_WinHttpSetOption)(
    HINTERNET hInternet,
    DWORD     dwOption,
    LPVOID    lpBuffer,
    DWORD     dwBufferLength
);

static void ca_error_cb(HINTERNET hInternet, DWORD_PTR dwContext, DWORD dwInternetStatus, LPVOID lpvStatusInformation, DWORD dwStatusInformationLength);

static const struct hook_symbol cert_syms[] = {
    {
        .name   = "CertFindCertificateInStore",
        .patch  = (void *) my_CertFindCertificateInStore,
        .link   = (void **) &next_CertFindCertificateInStore,
    },
    {
        .name   = "CertOpenStore",
        .patch  = my_CertOpenStore,
        .link   = (void **) &next_CertOpenStore,
    },
};

static const struct hook_symbol winhttp_syms[] = {
    {
        .name   = "WinHttpSetOption",
        .patch  = my_WinHttpSetOption,
        .link   = (void **) &next_WinHttpSetOption,
    },
};

HRESULT cert_hook_init(const struct cert_config *cfg)
{
    assert(cfg != NULL);
    
    if (!cfg->enable) {
        return S_FALSE;
    }

    dprintf("Cert hook init\n");
    wcscpy_s(path, MAX_PATH, cfg->path);
    InitializeCriticalSection(&cert_lock);
    cert_hook_insert_hooks(NULL);

    proc_addr_table_push(
        NULL,
        "crypt32.dll",
        cert_syms,
        _countof(cert_syms));

    proc_addr_table_push(
        NULL,
        "Winhttp.dll",
        winhttp_syms,
        _countof(winhttp_syms));
    
    if (cfg->path[0] != L'\0') {
        char client_cert_bfr[4096] = {0};
        BYTE client_cert_decode_bfr[4096] = {0};
        DWORD pcbBinary = 4096;
        DWORD cert_size = 0;
        HRESULT cert_hr;
        
        HANDLE client_cert_file = CreateFileW(cfg->path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
        if (client_cert_file == INVALID_HANDLE_VALUE) {
            cert_hr = HRESULT_FROM_WIN32(GetLastError());
            dprintf("Cert: CreateFileW Failed to open cert %ls because %X, using WINHTTP_NO_CLIENT_CERT_CONTEXT\n", cfg->path, (int)cert_hr);
            //SetLastError(0);
            return S_OK;
        }

        if (!ReadFile(client_cert_file, client_cert_bfr, 4096, &cert_size, NULL)) {
            cert_hr = HRESULT_FROM_WIN32(GetLastError());
            dprintf("Cert: ReadFile Failed to open cert %ls because %X, using WINHTTP_NO_CLIENT_CERT_CONTEXT\n", cfg->path, (int)cert_hr);
            //SetLastError(0);
            return S_OK;
        }

        dprintf("Cert: Read %d bytes\n", (int)cert_size);

        if (wcsstr(cfg->path, L".p12") == NULL) {
            if (!CryptStringToBinaryA(client_cert_bfr, cert_size, CRYPT_STRING_BASE64X509CRLHEADER, client_cert_decode_bfr, &pcbBinary, NULL, NULL)) {
                cert_hr = HRESULT_FROM_WIN32(GetLastError());
                dprintf("Cert: CryptStringToBinaryA Failed to decode cert %ls because %X, using WINHTTP_NO_CLIENT_CERT_CONTEXT\n", cfg->path, (int)cert_hr);
                //SetLastError(0);
                return S_OK;
            }

            client_cert_ctx = CertCreateCertificateContext(X509_ASN_ENCODING | PKCS_7_ASN_ENCODING, client_cert_decode_bfr, pcbBinary);
            if (client_cert_ctx == NULL) {
                dprintf("Cert: CertCreateCertificateContext Failed to parse cert %ls, using WINHTTP_NO_CLIENT_CERT_CONTEXT\n", cfg->path);
            }
        }

        else {
            struct _CRYPTOAPI_BLOB blob = {
                .cbData = cert_size,
                .pbData = (BYTE *)client_cert_bfr
            };

            //p12store = CertOpenStore(CERT_STORE_PROV_PKCS12, X509_ASN_ENCODING | PKCS_7_ASN_ENCODING, 0, CERT_STORE_READONLY_FLAG, &blob);
            p12store = PFXImportCertStore(&blob, NULL, 0);
            if (p12store == NULL) {
                cert_hr = HRESULT_FROM_WIN32(GetLastError());
                dprintf("Cert: CertOpenStore Failed to decode p12 data %ls because %X, using WINHTTP_NO_CLIENT_CERT_CONTEXT\n", cfg->path, (int)cert_hr);
                //SetLastError(0);
                return S_OK;
            }
        }
        

        dprintf("Cert: Using client cert from %ls\n", cfg->path);
    }
    
    return S_OK;
}

void cert_hook_insert_hooks(HMODULE target)
{
    hook_table_apply(
        target,
        "crypt32.dll",
        cert_syms,
        _countof(cert_syms));
    hook_table_apply(
        target,
        "winhttp.dll",
        winhttp_syms,
        _countof(winhttp_syms));
}

static PCCERT_CONTEXT WINAPI my_CertFindCertificateInStore(
    HCERTSTORE     hCertStore,
    DWORD          dwCertEncodingType,
    DWORD          dwFindFlags,
    DWORD          dwFindType,
    const void     *pvFindPara,
    PCCERT_CONTEXT pPrevCertContext
)
{
    char bfr[4096] = {0};
    uint8_t bfr_decode[4096] = {0};
    DWORD pcbBinary = 4096;
    DWORD num_read = 0;
    PCCERT_CONTEXT cert_ctx = NULL;

    if (client_cert_ctx != NULL) {
        if (dwFindType == CERT_FIND_ISSUER_STR || dwFindType == CERT_FIND_SUBJECT_STR) {
            wchar_t cert_issuer[129];

            CertGetNameStringW(client_cert_ctx, CERT_NAME_FRIENDLY_DISPLAY_TYPE, CERT_NAME_ISSUER_FLAG, NULL, cert_issuer, 129);
            if (!wcscmp(cert_issuer, (wchar_t *)pvFindPara)) {
                dprintf("Cert: Issuer %ls found, sending cert context %p\n", (wchar_t *)pvFindPara, cert_issuer);
                return CertDuplicateCertificateContext(client_cert_ctx);
            }

            dprintf("Cert: Issuer %ls requested, found %ls in loaded cert context %p\n", (wchar_t *)pvFindPara, cert_issuer, client_cert_ctx);
        }
    }
    
    return next_CertFindCertificateInStore(
        hCertStore,
        dwCertEncodingType,
        dwFindFlags,
        dwFindType,
        pvFindPara,
        pPrevCertContext
    );
}

static HCERTSTORE WINAPI my_CertOpenStore(
    LPCSTR            lpszStoreProvider,
    DWORD             dwEncodingType,
    HCRYPTPROV_LEGACY hCryptProv,
    DWORD             dwFlags,
    const void        *pvPara)
{
    BYTE bfr[4096] = {0};
    DWORD num_read = 0;
    if (lpszStoreProvider <= CERT_STORE_PROV_PKCS12) {
        dprintf("Cert: Open store for %p -> %S (%04X)\n", lpszStoreProvider, (wchar_t *)pvPara, (int)dwFlags);
    } else {
        dprintf("Cert: Open store for %s\n", lpszStoreProvider);
    }

    if (p12store != NULL) {
        return CertDuplicateStore(p12store);
    }

    HCERTSTORE ret = next_CertOpenStore(lpszStoreProvider, dwEncodingType, hCryptProv, dwFlags, pvPara);
    if (ret == NULL) {
        int err = GetLastError();
        if (err == 0x00000005) {
            ret = next_CertOpenStore(lpszStoreProvider, dwEncodingType, hCryptProv, 0x28000, pvPara); // This works without admin perms
            if (ret != NULL) {
                return ret;
            }
        }
        dprintf("Cert: Failed to open store %08X\n", (int)err);
    }
    return ret;
}

static BOOL WINAPI my_WinHttpSetOption(
    HINTERNET hInternet,
    DWORD     dwOption,
    LPVOID    lpBuffer,
    DWORD     dwBufferLength
)
{
    if (dwOption == WINHTTP_OPTION_CLIENT_CERT_CONTEXT) {
        dprintf("Cert: Set client cert\n");

        // Set our SSL error callback
        WINHTTP_STATUS_CALLBACK cb_check = WinHttpSetStatusCallback(hInternet, (WINHTTP_STATUS_CALLBACK)ca_error_cb, WINHTTP_CALLBACK_FLAG_SECURE_FAILURE, 0);
        if (cb_check == WINHTTP_INVALID_STATUS_CALLBACK) {
            dprintf("Cert: Failed to set SSL error callback: %08lX\n", GetLastError());
            SetLastError(0);
        }

        // Some games don't call WINHTTP_OPTION_SECURITY_FLAGS, so we do it here
        int value = SECURITY_FLAG_IGNORE_UNKNOWN_CA | SECURITY_FLAG_IGNORE_CERT_DATE_INVALID | SECURITY_FLAG_IGNORE_CERT_CN_INVALID | SECURITY_FLAG_IGNORE_CERT_WRONG_USAGE;
        next_WinHttpSetOption(hInternet, WINHTTP_OPTION_SECURITY_FLAGS, &value, 4);

        if (p12store == NULL && client_cert_ctx == NULL) {
            return next_WinHttpSetOption(hInternet, dwOption, WINHTTP_NO_CLIENT_CERT_CONTEXT, 0);
        }
    }
    else if (dwOption == WINHTTP_OPTION_SECURITY_FLAGS) {
        dprintf("Cert: Add all security ignore flags\n");
        
        // Set our SSL error callback
        WINHTTP_STATUS_CALLBACK cb_check = WinHttpSetStatusCallback(hInternet, (WINHTTP_STATUS_CALLBACK)ca_error_cb, WINHTTP_CALLBACK_FLAG_SECURE_FAILURE, 0);
        if (cb_check == WINHTTP_INVALID_STATUS_CALLBACK) {
            dprintf("Cert: Failed to set SSL error callback: %08lX\n", GetLastError());
            SetLastError(0);
        }
        
        int value = SECURITY_FLAG_IGNORE_UNKNOWN_CA | SECURITY_FLAG_IGNORE_CERT_DATE_INVALID | SECURITY_FLAG_IGNORE_CERT_CN_INVALID | SECURITY_FLAG_IGNORE_CERT_WRONG_USAGE; // the kitchen sink
        return next_WinHttpSetOption(hInternet, dwOption, &value, dwBufferLength);
    } else {
        dprintf("Cert: my_WinHttpSetOption %p %08X\n", hInternet, (int)dwOption);
    }

    return next_WinHttpSetOption(hInternet, dwOption, lpBuffer, dwBufferLength);
}

static void ca_error_cb(HINTERNET hInternet, DWORD_PTR dwContext, DWORD dwInternetStatus, LPVOID lpvStatusInformation, DWORD dwStatusInformationLength)
{
    dprintf("Cert: HTTP Secure connection failure: %04lX\n", *(DWORD *)lpvStatusInformation);
}