#include <windows.h>
#include <wincrypt.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
);

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

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

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

void cert_hook_insert_hooks(HMODULE target)
{
    hook_table_apply(
        target,
        "crypt32.dll",
        cert_syms,
        _countof(cert_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)
{
    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;
}