#include "cryptopro.h"

#include "base64.h"
#include "capi.h"
#include "library.h"
#include "logger.h"


static capi_function_list_t cp_function_list;
static library_t libcapi;

static int sign_hash_data(const cryptopro_context_t *ctx, const str_t *data, /*out*/ str_t *sign);

static int reverse_sign(const str_t *sign, /*out*/ str_t *sign_reversed);

static HCERTSTORE cert_open_store();
static void cert_close_store(HCERTSTORE hStoreHandle);

static PCCERT_CONTEXT get_signer_cert(const cryptopro_context_t *ctx, HCERTSTORE hStoreHandle);

static int set_password(PCCERT_CONTEXT cert_ctx, const str_t *password);

bool
cryptopro_init(const char* cp_file)
{
    static bool is_initialized = false;

    if (!is_initialized) {
        library_init(&libcapi, cp_file);
        if (!capi_function_list_init(&libcapi, &cp_function_list)) {
            LOG_ERROR("capi_function_list_init failed");
            return false;
        }
        is_initialized = true;
        LOG_INFO("Cryptographic Provider initialized");
    }
    return true;
}

int 
cryptopro_sign(const cryptopro_context_t *ctx, const str_t *data, /*out*/ str_t *sign)
{
    str_t signed_data = str_t_null; 
    str_t sign_reversed = str_t_null; 

    LOG_TRACE("cryptopro_sign enter");

    assert(ctx != NULL);
    assert(!str_t_is_null(*ctx->cert_thumbprint));
    assert(ctx->password != NULL);
    assert(data != NULL && !str_t_is_null(*data));
    assert(sign != NULL);

    if (sign_hash_data(ctx, data, &signed_data)) {
        goto error;
    }

    if (reverse_sign(&signed_data, &sign_reversed)) {
        goto error;
    }

    sign->len = base64_encoded_length(sign_reversed.len);
    sign->data = calloc(1, sign->len + 1);
    if (sign->data == NULL) {
        LOG_ERROR("Could not allocate memory for sign (%zu bytes)", sign->len + 1);
        goto error;
    }

    encode_base64_url(sign, &sign_reversed);

    str_t_clear(&signed_data);
    str_t_clear(&sign_reversed);

    LOG_TRACE("cryptopro_sign exit");
    return 0;

error:
    str_t_clear(&signed_data);
    str_t_clear(&sign_reversed);
    LOG_ERROR("cryptopro_sign exit with error (sign_cert_thumbprint = '%.*s')", 
                (int) ctx->cert_thumbprint->len, ctx->cert_thumbprint->data);
    return -1;
}

static void
free_cert_chain(PCCERT_CHAIN_CONTEXT chain_ctx)
{
    LOG_TRACE("free_cert_chain enter");

    if (chain_ctx) {
        cp_function_list.CertFreeCertificateChain(chain_ctx);
    }

    LOG_TRACE("free_cert_chain exit");
}

static PCCERT_CHAIN_CONTEXT
get_cert_chain(PCCERT_CONTEXT certificate)
{
    LOG_TRACE("get_cert_chain enter");

    CERT_CHAIN_PARA chain_para = {0};
    chain_para.cbSize = sizeof(CERT_CHAIN_PARA);

    PCCERT_CHAIN_CONTEXT chain_ctx = NULL;


    if (!cp_function_list.CertGetCertificateChain(NULL, 
                    certificate, 
                    NULL, 
                    NULL, 
                    &chain_para,
                    (CERT_CHAIN_CACHE_END_CERT | CERT_CHAIN_REVOCATION_CHECK_CHAIN_EXCLUDE_ROOT),
                    NULL, 
                    &chain_ctx)) {
        LOG_ERROR("CertGetCertificateChain() failed");
        goto error;
    }

    if (chain_ctx->TrustStatus.dwErrorStatus) {
        LOG_WARN("The certificate is not trusted. CERT_TRUST_STATUS: '0x%08x'", 
                    chain_ctx->TrustStatus.dwErrorStatus);
        goto error;
    }

    LOG_TRACE("get_cert_chain exit");
    return chain_ctx;

error:
    free_cert_chain(chain_ctx);
    LOG_ERROR("get_cert_chain exit with error");
    return NULL;
}

static bool
check_cert_chain_policy(PCCERT_CHAIN_CONTEXT chain_ctx)
{
    LOG_TRACE("check_cert_chain_policy enter");

    bool is_valid = false;

    CERT_CHAIN_POLICY_PARA policy_para = {0};
    policy_para.cbSize = sizeof(policy_para);

    CERT_CHAIN_POLICY_STATUS status = {0};
    status.cbSize = sizeof(status);

    CPSIGNATURE_EXTRA_CERT_CHAIN_POLICY_STATUS extraStatus = {0};
    extraStatus.cbSize = sizeof(extraStatus);

    status.pvExtraPolicyStatus = &extraStatus;

    if (!cp_function_list.CertVerifyCertificateChainPolicy(
            CERT_CHAIN_POLICY_BASE,
            chain_ctx, 
            &policy_para,
            &status)) {
        LOG_ERROR("CertVerifyCertificateChainPolicy() failed. Error code: 0x%08x", 
                    cp_function_list.GetLastError());
        goto exit;
    }

    if (status.dwError != 0) {
        LOG_WARN("The certificate chain cannot be validated. "
                 "CERT_CHAIN_POLICY_STATUS: '0x%08x'", status.dwError);
        goto exit;
    }

    if (extraStatus.dwError != 0) {
        LOG_WARN("The certificate chain cannot be validated. "
                 "CPSIGNATURE_EXTRA_CERT_CHAIN_POLICY_STATUS: '0x%08x'", extraStatus.dwError);
        goto exit;
    }

    is_valid = true;

exit:
    LOG_TRACE("check_cert_chain_policy exit");
    return is_valid;
}

static bool
verify_cert_chain(PCCERT_CONTEXT certificate)
{
    bool is_verified = false;

    LOG_TRACE("verify_cert_chain enter");

    PCCERT_CHAIN_CONTEXT chain_ctx = get_cert_chain(certificate);
    if (chain_ctx == NULL) {
        goto exit;
    }

    if (!check_cert_chain_policy(chain_ctx)) {
        goto exit;
    }

    is_verified = true;

exit:
    free_cert_chain(chain_ctx);

    LOG_DEBUG("cert_chain: %s", is_verified ? "verified" : "failed");

    LOG_DEBUG("verify_cert_chain exit");
    return is_verified;
}

static void
log_sign_hash_data_last_error()
{
    DWORD error = cp_function_list.GetLastError();

    switch (error) {
    case ERROR_FUNCTION_FAILED:
        LOG_ERROR("sign_hash_data exit with error. Last error code: "
                    "license is expired or not yet valid (0x%08x)", error);
        break;

    case SCARD_W_WRONG_CHV:
        LOG_ERROR("sign_hash_data exit with error. Last error code: "
                    "the wrong PIN was presented (0x%08x)", error);
        break;
        
    default:
        LOG_ERROR("sign_hash_data exit with error. Last error code: 0x%08x", error);
        break;
    }
}

int 
open_signer_cert(cryptopro_context_t *ctx)
{
    LOG_TRACE("open_signer_cert enter");

    ctx->cert_store = cert_open_store();
    if (ctx->cert_store == NULL) {
        goto error;
    }

    ctx->signer_cert = get_signer_cert(ctx, ctx->cert_store);
    if (ctx->signer_cert == NULL) {
        goto error;
    }

    LOG_TRACE("open_signer_cert exit");

    return 0;

error:
    LOG_ERROR("open_signer_cert exit with error. Last error code: 0x%08x", 
                cp_function_list.GetLastError());
    close_signer_cert(ctx);
    return -1;
}

void
close_signer_cert(cryptopro_context_t *ctx)
{
    LOG_TRACE("close_signer_cert enter");

    if (ctx->signer_cert) {
        cp_function_list.CertFreeCertificateContext(ctx->signer_cert);
        ctx->signer_cert = NULL;
    }

    cert_close_store(ctx->cert_store);
    ctx->cert_store = NULL;

    LOG_TRACE("close_signer_cert exit");
}

static int 
sign_hash_data(const cryptopro_context_t *ctx, const str_t *data, /*out*/ str_t *sign)
{
    int rc = -1;
    BOOL bReleaseContext = FALSE;
    HCRYPTPROV hCryptProv;
    HCRYPTHASH hash = 0;
    BYTE *pbSignedMessageBlob;
    DWORD cbSignedMessageBlob;

    if (!verify_cert_chain(ctx->signer_cert)) {
        goto exit;
    }

    if (!cp_function_list.CryptAcquireCertificatePrivateKey(
            ctx->signer_cert,
            CRYPT_ACQUIRE_SILENT_FLAG | CRYPT_ACQUIRE_CACHE_FLAG,
            NULL,
            &hCryptProv,
            NULL,
            &bReleaseContext)) {
        LOG_ERROR("Cannot acquire the certificate private key");
        goto exit;
    }

    if (!cp_function_list.CryptCreateHash(hCryptProv, CALG_GR3411_2012_256, 0, 0, &hash)) {
        LOG_ERROR("CryptCreateHash() failed");
        goto exit;
    }
            
    if (!cp_function_list.CryptHashData(hash, (BYTE *)data->data, data->len, 0)) {
        LOG_ERROR("CryptHashData() failed");
        goto exit;
    }
    
    if (!cp_function_list.CryptSignHash(hash, AT_KEYEXCHANGE, NULL, 0, NULL, &cbSignedMessageBlob)) {
        LOG_ERROR("CryptSignHash() failed");
        goto exit;
    }
    
    pbSignedMessageBlob = malloc(cbSignedMessageBlob * sizeof(BYTE));
    
    if (!cp_function_list.CryptSignHash(hash, AT_KEYEXCHANGE, NULL, 0, pbSignedMessageBlob, &cbSignedMessageBlob)) {
        LOG_ERROR("CryptSignHash() failed");
        free(pbSignedMessageBlob);
        goto exit;
    }

    sign->data = (char*)pbSignedMessageBlob;
    sign->len = (size_t)cbSignedMessageBlob;

    rc = 0;

exit:
    if (hash) {
        cp_function_list.CryptDestroyHash(hash);
    }

    if (bReleaseContext) {
        cp_function_list.CryptReleaseContext(hCryptProv, 0);
    }


    if (rc) {
        log_sign_hash_data_last_error();
    }

    return rc;
}

static int 
reverse_sign(const str_t *sign, /*out*/ str_t *sign_reversed)
{
    size_t i;

    sign_reversed->data = malloc(sign->len);
    if (sign_reversed->data == NULL) {
        LOG_ERROR("Could not allocate memory for reversed sign (%zd bytes)", sign->len);
        return -1;
    }

    for (i = 0; i < sign->len; i++) {
        sign_reversed->data[i] = sign->data[sign->len - i - 1];
    }

    sign_reversed->len = sign->len;

    return 0;
}

static HCERTSTORE 
cert_open_store()
{
    HCERTSTORE hStoreHandle;

    LOG_TRACE("cert_open_store enter");

    hStoreHandle = cp_function_list.CertOpenStore(CERT_STORE_PROV_SYSTEM,
                                    0,
                                    0,
                                    CERT_SYSTEM_STORE_CURRENT_USER | CERT_STORE_OPEN_EXISTING_FLAG,
                                    L"MY");
    
    if (hStoreHandle == NULL) {
        LOG_ERROR("The store could not be opened");
        return NULL;
    }

    LOG_DEBUG("The MY store is open");
    LOG_TRACE("cert_open_store exit");
    return hStoreHandle;
}


static void
cert_close_store(HCERTSTORE hStoreHandle)
{
    LOG_TRACE("cert_close_store enter");

    if (hStoreHandle) {
        if (!cp_function_list.CertCloseStore(hStoreHandle, CERT_CLOSE_STORE_CHECK_FLAG)) {
            LOG_ERROR("CertCloseStore() failed");
        }
    }

    LOG_TRACE("cert_close_store exit");
}

static PCCERT_CONTEXT 
get_cert_by_thumbprint(HCERTSTORE hStoreHandle, const str_t* thumbprint)
{
    str_t thumbprint_bin = str_t_null;

    LOG_TRACE("get_cert_by_thumbprint enter");

    if (hex_to_bin(thumbprint, &thumbprint_bin)) {
        goto error;
    }

    CRYPT_HASH_BLOB hash = {
        .pbData = (BYTE*)thumbprint_bin.data,
        .cbData = thumbprint_bin.len
    };

    PCCERT_CONTEXT cert = cp_function_list.CertFindCertificateInStore(hStoreHandle,
                                                X509_ASN_ENCODING | PKCS_7_ASN_ENCODING,
                                                0,
                                                CERT_FIND_HASH,
                                                &hash,
                                                NULL);
    if (cert == NULL) {
        LOG_ERROR("Could not find certificate in store");
        goto error;
    }

    str_t_clear(&thumbprint_bin);

    LOG_TRACE("get_cert_by_thumbprint exit");

    return cert;

error:
    str_t_clear(&thumbprint_bin);
    LOG_ERROR("get_cert_by_thumbprint exit with error");
    return NULL;
}

static PCCERT_CONTEXT 
get_signer_cert(const cryptopro_context_t *ctx, HCERTSTORE hStoreHandle)
{
    PCCERT_CONTEXT cert_ctx = get_cert_by_thumbprint(hStoreHandle, ctx->cert_thumbprint);
    if (cert_ctx == NULL) {
        goto error;
    }

    if (!str_t_is_null(*ctx->password)) {
        if (set_password(cert_ctx, ctx->password)) {
            goto error;
        }
    }

    LOG_DEBUG("The signer's certificate was found");
    return cert_ctx;

error:
    if (cert_ctx) {
        cp_function_list.CertFreeCertificateContext(cert_ctx);
    }
    LOG_ERROR("get_signer_cert exit with error");
    return NULL;
}

static int
set_password(PCCERT_CONTEXT cert_ctx, const str_t *password)
{
    DWORD dwSize;
    if (!cp_function_list.CertGetCertificateContextProperty(
             cert_ctx, 
             CERT_KEY_PROV_INFO_PROP_ID, 
             NULL,
             &dwSize)) {
        LOG_ERROR("Error getting key property");
        LOG_ERROR("set_password exit with error. Last error code: 0x%08x", cp_function_list.GetLastError());
        return -1;
    }

    unsigned char pKeyInfoBuffer[dwSize];
    CRYPT_KEY_PROV_INFO* pKeyInfo = (CRYPT_KEY_PROV_INFO*) pKeyInfoBuffer;

    if (!cp_function_list.CertGetCertificateContextProperty(
            cert_ctx,
            CERT_KEY_PROV_INFO_PROP_ID,
            pKeyInfo,
            &dwSize)){
        LOG_ERROR("The second call to the function failed");
        goto error;
    }

    DWORD keyType = (pKeyInfo->dwKeySpec == AT_SIGNATURE) ? PP_SIGNATURE_PIN
                                                          : PP_KEYEXCHANGE_PIN;

    CRYPT_KEY_PROV_PARAM keyProvParam;
    memset(&keyProvParam, 0, sizeof(CRYPT_KEY_PROV_PARAM));
    keyProvParam.dwFlags = 0;
    keyProvParam.dwParam = keyType;
    keyProvParam.cbData  = (DWORD)(password->len + 1);
    keyProvParam.pbData  = (BYTE*) password->data;

    pKeyInfo->cProvParam = 1;
    pKeyInfo->rgProvParam = &keyProvParam;

    if (!cp_function_list.CertSetCertificateContextProperty(
                cert_ctx,
                CERT_KEY_PROV_INFO_PROP_ID,
                0,
                pKeyInfo)) {
        LOG_ERROR("Error setting key property");
        goto error;
    }

    return 0;

error:
    LOG_ERROR("set_password exit with error. Last error code: 0x%08x", cp_function_list.GetLastError());
    return -1;
}

static void
set_unknown_alg_error(/*out*/ char** verify_error, const str_t* alg)
{
    LOG_TRACE("set_unknown_alg_error enter");

    const str_t err_format = str_t_const("Unknown alg: '%.*s'");

    size_t size = err_format.len + alg->len;

    *verify_error = malloc(size);
    if (*verify_error == NULL) {
        LOG_ERROR("set_unknown_alg_error failed. Could not allocate memory for error string "
                "(%zd bytes)", size);
        goto exit;
    }

    int n = snprintf(*verify_error, size, err_format.data, (int) alg->len, alg->data);

    if (n < 0 || (size_t)n >= size) {
        LOG_ERROR("set_unknown_alg_error failed. Error occurred in concatenation err_string");
        free(*verify_error);
        *verify_error = NULL;
        goto exit;
    }

    LOG_ERROR("%s", *verify_error);

exit:
    LOG_TRACE("set_unknown_alg_error exit");
}

static int
alg_id_from_str(const str_t* alg, /*out*/ ALG_ID* alg_id, /*out*/ char** verify_error)
{
    LOG_TRACE("alg_id_from_str enter");

    if (strncmp(alg->data, "GOST3410_2012_256", alg->len) == 0) {
        *alg_id = CALG_GR3411_2012_256;
    } else if (strncmp(alg->data, "GOST3410_2012_512", alg->len) == 0) {
        *alg_id = CALG_GR3411_2012_512;
    } else {
        set_unknown_alg_error(verify_error, alg);
        LOG_ERROR("alg_id_from_str exit with error");
        return -1;
    }

    LOG_TRACE("alg_id_from_str exit");
    return 0;
}

static void
get_verify_error(char** verify_error)
{
    LOG_TRACE("get_verify_error enter");

    DWORD err = cp_function_list.GetLastError();

    const char* err_string;

    switch (err) {
    case NTE_BAD_SIGNATURE:
        err_string = "sign is invalid: bad signature (0x%08x)";
        break;

    case NTE_BAD_ALGID:
        err_string = "sign is invalid: bad alg id (0x%08x)";
        break;
    
    default:
        err_string = "sign is invalid. Last error code: 0x%08x";
        break;
    }

    size_t size = strlen(err_string) + 4 /*error code*/ + 1 /*terminating null*/;

    if (*verify_error != NULL) {
        free(*verify_error); // в ответе отдаем последнюю ошибку
    }

    *verify_error = malloc(size);
    if (*verify_error == NULL) {
        LOG_ERROR("get_verify_error failed. Could not allocate memory for error string "
                    "(%zd bytes)", size);
        return;
    }

    int n = snprintf(*verify_error, size, err_string, err);
    
    if (n < 0 || (size_t)n >= size) {
        LOG_ERROR("get_verify_error failed. Error occurred in concatenation err_string");
        free(*verify_error);
        *verify_error = NULL;
        return;
    }

    LOG_TRACE("get_verify_error exit");
}

int 
cryptopro_verify(const str_t* cert_thumbprint, const str_t* alg, const str_t* data, 
                    const str_t* sign, bool* is_verified, char** verify_error)
{
    int rc = -1;
    HCERTSTORE hStoreHandle = NULL;
    PCCERT_CONTEXT certificate = NULL;
    HCRYPTPROV hCryptProv = 0;
    HCRYPTHASH hash = 0;
    HCRYPTKEY hPubKey = 0;
    str_t sign_reversed = str_t_null;
    ALG_ID alg_id;

    LOG_TRACE("cryptopro_verify enter");

    *is_verified = false;

    if (alg_id_from_str(alg, &alg_id, verify_error)) {
        goto exit;
    }

    if (reverse_sign(sign, &sign_reversed)) {
        goto exit;
    }

    hStoreHandle = cert_open_store();
    if (hStoreHandle == NULL) {
        goto exit;
    }

    certificate = get_cert_by_thumbprint(hStoreHandle, cert_thumbprint);
    if (certificate == NULL) {
        goto exit;
    }

    if (!verify_cert_chain(certificate)) {
        goto exit;
    }

    if (!cp_function_list.CryptAcquireContext(&hCryptProv, NULL, NULL, PROV_GOST_2012_256, 
                                                CRYPT_VERIFYCONTEXT)) {
        LOG_ERROR("CryptAcquireContext() failed");
        goto exit;
    }

    if (!cp_function_list.CryptImportPublicKeyInfo(hCryptProv, X509_ASN_ENCODING,
                                                    &certificate->pCertInfo->SubjectPublicKeyInfo, 
                                                    &hPubKey)) {
        LOG_ERROR("CryptImportPublicKeyInfo() failed");
        goto exit;
    }

    if (!cp_function_list.CryptCreateHash(hCryptProv, alg_id, 0, 0, &hash)) {
        LOG_ERROR("CryptCreateHash() failed");
        goto exit;
    }
            
    if (!cp_function_list.CryptHashData(hash, (CONST BYTE *) data->data, data->len, 0)) {
        LOG_ERROR("CryptHashData() failed");
        goto exit;
    }

    if (cp_function_list.CryptVerifySignature(hash, 
                                                (CONST BYTE *) sign_reversed.data, 
                                                sign_reversed.len, 
                                                hPubKey,
                                                NULL, 
                                                0)) {
        LOG_DEBUG("sign is valid");
        *is_verified = true;
    } else {
        get_verify_error(verify_error);
        if (*verify_error == NULL) {
            goto exit;
        }
        LOG_WARN("%s, cert_thumbprint: %.*s", *verify_error, 
                    (int) cert_thumbprint->len, cert_thumbprint->data);
    }

exit:

    if (*is_verified || *verify_error) {
        rc = 0;
    }

    str_t_clear(&sign_reversed);

    if (hash) {
        cp_function_list.CryptDestroyHash(hash);
    }

    if (hPubKey) {
        cp_function_list.CryptDestroyKey(hPubKey);
    }

    if (hCryptProv) {
        cp_function_list.CryptReleaseContext(hCryptProv, 0);
    }

    if (certificate) {
        cp_function_list.CertFreeCertificateContext(certificate);
    }

    cert_close_store(hStoreHandle);

    if (rc == 0) {
        LOG_TRACE("cryptopro_verify exit");
    } else {
        LOG_ERROR("cryptopro_verify exit with error. Last error code: 0x%08x", 
                    cp_function_list.GetLastError());
    }

    return rc;  
}

int 
cryptopro_gen_random(unsigned char* data, size_t len)
{
    HCRYPTPROV   hCryptProv = 0;

    LOG_TRACE("cryptopro_gen_random enter");

    if (!cp_function_list.CryptAcquireContext(&hCryptProv, NULL, NULL, PROV_GOST_2012_256, 
                                                CRYPT_VERIFYCONTEXT)) {
        LOG_ERROR("CryptAcquireContext() failed");
        goto error;
    }

    if (!cp_function_list.CryptGenRandom(hCryptProv, len, data)) {
        LOG_ERROR("CryptGenRandom() failed");
        goto error;
    }

    cp_function_list.CryptReleaseContext(hCryptProv, 0);

    LOG_TRACE("cryptopro_gen_random exit");
    return 0;

error:
    if (hCryptProv) {
        cp_function_list.CryptReleaseContext(hCryptProv, 0);
    }

    LOG_ERROR("cryptopro_gen_random exit with error. Last error code: 0x%08x", 
                cp_function_list.GetLastError());
    return -1;
}