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

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

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

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
)
{
    uint8_t bfr[4096] = {0};
    wchar_t cert_path[MAX_PATH] = {0};
    DWORD num_read = 0;

    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

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

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

        if (f != INVALID_HANDLE_VALUE) {
            ReadFile(f, bfr, sizeof(bfr), &num_read, NULL);
            CloseHandle(f);

            if (bfr[0]) {
                return CertCreateCertificateContext(X509_ASN_ENCODING | PKCS_7_ASN_ENCODING, bfr, num_read);
            }
        }
    }
    
    return next_CertFindCertificateInStore(
        hCertStore,
        dwCertEncodingType,
        dwFindFlags,
        dwFindType,
        pvFindPara,
        pPrevCertContext
    );
}