// Copyright 2026 Osmium Sorcerer
// SPDX-License-Identifier: MIT

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <ncrypt.h>
#include <rpc.h>

#include "tpm.h"

// Global handle for the Windows API.
static NCRYPT_PROV_HANDLE provider;

// 9 characters for prefix, 36 for random UUID, and don't forget that null terminator.
// Note: this is count of wchar_t, not size.
static const size_t keyname_len = 46;

stkb_rc tpm2_init(void)
{
    // "Microsoft Platform Crypto Provider" is the TPM.
    if (NCryptOpenStorageProvider(&provider, MS_PLATFORM_CRYPTO_PROVIDER, 0) != ERROR_SUCCESS)
        return STKB_TPM_UNAVAILABLE;
    if (NCryptIsAlgSupported(provider, BCRYPT_ECDH_P256_ALGORITHM, 0) != ERROR_SUCCESS)
        return STKB_TPM_UNAVAILABLE;
    return STKB_SUCCESS;
}


static int generate_keyname(wchar_t *out)
{
    // UUID seems like the way on Windows.
    UUID id;
    RPC_WSTR wstr = NULL;

    if (UuidCreate(&id) != RPC_S_OK)
        return 0;

    if (UuidToStringW(&id, &wstr) != RPC_S_OK)
        return 0;

    wcscpy(out, L"SoF_Auth_");
    wcscpy(out + 9, wstr);

    RpcStringFreeW(&wstr);

    return 1;
}

stkb_rc tpm2_keygen(uint8_t *out_pub, uint8_t *out_blob, size_t buffer_size, size_t *offset,
                stkb_user_input input)
{
    NCRYPT_KEY_HANDLE key_handle = 0;
    // Windows heavily abstracts the TPM and prevents you from acquiring the encrypted private structure. Instead, you
    // have to use wide null-terminated strings as canonical key identifiers and refer to them by these names later.
    // Yes, not even opaque byte arrays despite it handling "BLOBs."
    wchar_t keyname[keyname_len];
    if (sizeof(keyname) > buffer_size)
        return STKB_ERROR;
    if (!generate_keyname(keyname))
        return STKB_ERROR;

    // One might think, when NULL is passed instead of a key name, the key will not actually be persistent. But it
    // becomes impossible to export the key to be loaded later, even in an encrypted form (while NCryptExportKey
    // documents NCRYPT_OPAQUETRANSPORT_BLOB, explicitly saying "Opaque BLOBs are not transferable and must be imported
    // by using the same CSP," and NCRYPT_PROTECTED_KEY_BLOB, confusingly, also doesn't export an opaque fixed-TPM
    // restricted key blob. So we are forced to use a unique (potentially user-facing) name as an internal specifier.
    if (NCryptCreatePersistedKey(provider, &key_handle, BCRYPT_ECDH_P256_ALGORITHM, keyname, 0, 0) != ERROR_SUCCESS)
        return STKB_ERROR;

    // Leave PINs unused, let the platform handle the input dialog.
    NCRYPT_UI_POLICY ui_policy = {
        .dwVersion = 1,
        .dwFlags = NCRYPT_UI_FORCE_HIGH_PROTECTION_FLAG,
        .pszFriendlyName = input.key_name,
        .pszDescription = L"SoF authentication key",
    };
    stkb_rc ret = STKB_ERROR;
    if (NCryptSetProperty(key_handle, NCRYPT_UI_POLICY_PROPERTY, (PBYTE)&ui_policy, sizeof(ui_policy), 0) != ERROR_SUCCESS)
        goto exit;

    // If the TPM is locked out, of all functions, FinalizeKey fails by returning NTE_INVALID_HANDLE.
    // The other way it can happen (memory corruption aside) is if the user cancels the key creation
    // dialog. You can't differentiate the two.
    if (NCryptFinalizeKey(key_handle, 0) != ERROR_SUCCESS) {
        ret = STKB_LOCKOUT;
        goto exit;
    }

    BCRYPT_ECCKEY_BLOB ecc_blob = { 0 };
    BYTE pub_buffer[sizeof(ecc_blob) + 32 * 2];
    DWORD pub_buffer_off = 0;
    if (NCryptExportKey(key_handle, 0, BCRYPT_ECCPUBLIC_BLOB, NULL, pub_buffer, sizeof(pub_buffer), &pub_buffer_off,
                        0) != ERROR_SUCCESS)
        goto exit;
    memcpy(&ecc_blob, pub_buffer, sizeof(ecc_blob));
    // Check sanity of the exported public key.
    if (ecc_blob.dwMagic != BCRYPT_ECDH_PUBLIC_P256_MAGIC || ecc_blob.cbKey != 32)
        goto exit;
    // The x and y _should_ follow the blob (which is the beginning of the buffer) contiguously, each of cbKey (32)
    // bytes.
    memcpy(out_pub + 1, pub_buffer + sizeof(ecc_blob), 32);
    // Encode parity of the last byte of y (both coordinates are big-endian).
    if (pub_buffer[sizeof(ecc_blob) + 32 + 31] & 1)
        out_pub[0] = 0x03;
    else
        out_pub[0] = 0x02;

    memcpy(out_blob, keyname, sizeof(keyname));
    *offset = sizeof(keyname);

    ret = STKB_SUCCESS;

exit:
    NCryptFreeObject(key_handle);
    return ret;
}

stkb_rc tpm2_compute_ss(uint8_t *ss, const uint8_t *blob, size_t blob_len, const uint8_t *pk, stkb_user_input input)
{
    (void)input;
    // Because wchar_t is 16-byte-aligned, we can't cast a byte array, we have to do this redundant copy.
    wchar_t keyname[keyname_len];
    if (blob_len != sizeof(keyname))
        return STKB_ERROR;
    memcpy(keyname, blob, sizeof(keyname));
    stkb_rc ret = STKB_ERROR;
    NCRYPT_KEY_HANDLE key_handle = 0;
    NCRYPT_KEY_HANDLE pk_handle = 0;
    NCRYPT_SECRET_HANDLE shared_point = 0;
    if (NCryptOpenKey(provider, &key_handle, keyname, 0, 0) != ERROR_SUCCESS)
        goto exit;
    // The API expects us to do the same ceremony to import the key in its format instead of using it directly.
    BCRYPT_ECCKEY_BLOB ecc_blob = {
        .dwMagic = BCRYPT_ECDH_PUBLIC_P256_MAGIC,
        .cbKey = 32,
    };
    BYTE pub_buffer[sizeof(ecc_blob) + 32 * 2];
    memcpy(pub_buffer, &ecc_blob, sizeof(ecc_blob));
    memcpy(pub_buffer + sizeof(ecc_blob), pk, 64);
    if (NCryptImportKey(provider, 0, BCRYPT_ECCPUBLIC_BLOB, NULL, &pk_handle, pub_buffer, sizeof(pub_buffer), 0) != ERROR_SUCCESS)
        goto exit;
    // Here it can also fail due to cancelation of dialog, too many failed attempts, or a TPM lockout.
    if (NCryptSecretAgreement(key_handle, pk_handle, &shared_point, 0) != ERROR_SUCCESS) {
        ret = STKB_LOCKOUT;
        goto exit;
    }
    // You cannot extract the x coordinate of the shared point either, you have to derive the key. Even if you want to
    // directly use the secret. But fine, apparently RAW_SECRET is a cryptographic key derivation function, which is
    // some unrecognized cryptographic genius. We'll use it to get our affine big-endian x coordinate, the canonical
    // ECDH shared secret. If you specify the size to be 32, it will only output the x coordinate, after all, that's
    // what a point is: 32 bytes of x, then 32 bytes of y, no padding, no leading bytes, no headers. Right?
    BYTE secret[32];
    DWORD written = 0;
    if (NCryptDeriveKey(shared_point, BCRYPT_KDF_RAW_SECRET, NULL, secret, sizeof(secret), &written, 0) != ERROR_SUCCESS)
        goto exit;
    // Of course it wouldn't be that easy. Even if we guessed the behavior of this function with respect to the buffer
    // size you pass to it, CNG API decides that it's a good idea to output point coordinates in little-endian order,
    // opposite of how the standard defines it (and what every other implementation correctly does, including the
    // Platform Crypto backend that CNG relies on, and CNG's own NCryptExportKey).
    for (size_t i = 0; i < sizeof(secret); ++i)
        ss[i] = secret[sizeof(secret) - 1 - i];

    ret = STKB_SUCCESS;

exit:
    NCryptFreeObject(key_handle);
    NCryptFreeObject(pk_handle);
    NCryptFreeObject(shared_point);
    return ret;
}