#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"

static CRITICAL_SECTION cert_lock;
static wchar_t path[MAX_PATH];

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

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

WINHTTPAPI BOOL hook_WinHttpSetOption(
    HINTERNET hInternet,
    DWORD     dwOption,
    LPVOID    lpBuffer,
    DWORD     dwBufferLength
);

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

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

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

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  = hook_CertFindCertificateInStore,
        .link   = (void **) &next_CertFindCertificateInStore,
    },
    {
        .name   = "CertOpenStore",
        .patch  = hook_CertOpenStore,
        .link   = (void **) &next_CertOpenStore,
    },
};

static const struct hook_symbol winhttp_syms[] = {
    {
        .name   = "WinHttpSetOption",
        .patch  = hook_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",
        (struct hook_symbol *) cert_syms,
        _countof(cert_syms));

    proc_addr_table_push(
        NULL,
        "Winhttp.dll",
        (struct hook_symbol *) winhttp_syms,
        _countof(winhttp_syms));
    
    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",
        (struct hook_symbol *) winhttp_syms,
        _countof(winhttp_syms));
}

PCCERT_CONTEXT WINAPI hook_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;
    wchar_t cert_path[MAX_PATH] = {0};
    DWORD num_read = 0;
    PCCERT_CONTEXT cert_ctx = NULL;

    if (dwFindType == CERT_FIND_ISSUER_STR || dwFindType == CERT_FIND_SUBJECT_STR) {
        wcscat_s(cert_path, _countof(cert_path), path);
        wcscat_s(cert_path, _countof(cert_path), L"/");
        wcscat_s(cert_path, _countof(cert_path), (wchar_t *)pvFindPara); // use the search string as a name
        wcscat_s(cert_path, _countof(cert_path), L".cer");

        // dprintf("Cert: Look for override cert at %S\n", cert_path);

        HANDLE f = CreateFileW((LPCWSTR)cert_path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

        if (f != INVALID_HANDLE_VALUE) {
            //dprintf("Cert: Read file %S\n", cert_path);
            ReadFile(f, bfr, sizeof(bfr), &num_read, NULL);
            CloseHandle(f);

            if (bfr[0]) {
                dprintf("Cert: Override %S\n", cert_path);
                if (CryptStringToBinary(bfr, 0, CRYPT_STRING_BASE64X509CRLHEADER, bfr_decode, &pcbBinary, NULL, NULL)) {
                    cert_ctx = CertCreateCertificateContext(X509_ASN_ENCODING, bfr_decode, num_read);
                    if (cert_ctx != NULL) {
                        return cert_ctx;
                    }                    
                    dprintf("Cert: Override FAIL %08X\n", (int)GetLastError());
                }
                dprintf("Cert: CryptStringToBinary FAIL %08X\n", (int)GetLastError());
            }
        }
    }
    
    return next_CertFindCertificateInStore(
        hCertStore,
        dwCertEncodingType,
        dwFindFlags,
        dwFindType,
        pvFindPara,
        pPrevCertContext
    );
}

HCERTSTORE WINAPI hook_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);
    }

    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;
    /*HANDLE f = CreateFileW(path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

        if (f != INVALID_HANDLE_VALUE) {
            dprintf("Cert: Read file %S\n", path);
            ReadFile(f, bfr, sizeof(bfr), &num_read, NULL);
            CloseHandle(f);

            if (bfr[0]) {
                CRYPT_INTEGER_BLOB blob = {
                    .pbData = bfr,
                    .cbData = num_read
                };
                dprintf("Cert: detour open of %S to %S\n", (wchar_t *)pvPara, path);
                HCERTSTORE ret = next_CertOpenStore(CERT_STORE_PROV_PKCS12, X509_ASN_ENCODING | PKCS_7_ASN_ENCODING, hCryptProv, dwFlags, &blob);
                dprintf("Cert: Open %p\n", ret);
                return ret;
            }
        }

    return next_CertOpenStore(lpszStoreProvider, dwEncodingType, hCryptProv, dwFlags, pvPara);*/

}
WINHTTPAPI BOOL hook_WinHttpSetOption(
    HINTERNET hInternet,
    DWORD     dwOption,
    LPVOID    lpBuffer,
    DWORD     dwBufferLength
)
{
    if (dwOption == WINHTTP_OPTION_CLIENT_CERT_CONTEXT) {
        // This is U G L Y and will fail on servers that actually check the client cert.
        dprintf("Cert: Block WINHTTP_OPTION_CLIENT_CERT_CONTEXT\n");
        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);
        }

        // Sneak in security disable while we're here
        int value = SECURITY_FLAG_IGNORE_UNKNOWN_CA | SECURITY_FLAG_IGNORE_CERT_DATE_INVALID | SECURITY_FLAG_IGNORE_CERT_CN_INVALID; // the kitchen sink
        if (!next_WinHttpSetOption(hInternet, WINHTTP_OPTION_SECURITY_FLAGS, &value, 4)) {
            dprintf("Cert: Failed to set ignore security flags: %08lX\n", GetLastError());
            SetLastError(0);
        }
        return true;
    }
    else if (dwOption == WINHTTP_OPTION_SECURITY_FLAGS) {
        dprintf("Cert: Add all security ignore flags\n");
        
        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; // the kitchen sink
        return next_WinHttpSetOption(hInternet, dwOption, &value, dwBufferLength);
    } else {
        dprintf("Cert: hook_WinHttpSetOption %p %08X\n", hInternet, (int)dwOption);
    }

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

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);
}