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

#include <string.h>

#include <tss2/tss2_esys.h>
#include <tss2/tss2_mu.h>

#include "tpm.h"

// Storage Key: asymmetric scheme is NULL, symmetric is defined and shall
// use CFB mode. The key is restricted.
static const TPM2B_PUBLIC primary_template = {
    .publicArea = {
        .type = TPM2_ALG_ECC,
        .nameAlg = TPM2_ALG_SHA256,
        .objectAttributes =
            TPMA_OBJECT_FIXEDTPM |
            TPMA_OBJECT_FIXEDPARENT |
            TPMA_OBJECT_SENSITIVEDATAORIGIN |
            TPMA_OBJECT_USERWITHAUTH |
            TPMA_OBJECT_RESTRICTED |
            TPMA_OBJECT_DECRYPT,
        .parameters.eccDetail = {
            .symmetric = {
                .algorithm = TPM2_ALG_AES,
                .keyBits.aes = 128,
                .mode.aes = TPM2_ALG_CFB,
            },
            .scheme.scheme = TPM2_ALG_NULL,
            .curveID = TPM2_ECC_NIST_P256,
            .kdf.scheme = TPM2_ALG_NULL,
        },
    },
};

// Symmetric algorithm is NULL for a nonrestricted key. This is the template
// for actual key exchange keys under the primary storage key, accessed with a
// user authentication value.
// There is TPM2_ALG_ECDH scheme, but it requires KDF.
static const TPM2B_PUBLIC ecdh_key_template = {
    .publicArea = {
        .type = TPM2_ALG_ECC,
        .nameAlg = TPM2_ALG_SHA256,
        .objectAttributes =
            TPMA_OBJECT_FIXEDTPM |
            TPMA_OBJECT_FIXEDPARENT |
            TPMA_OBJECT_SENSITIVEDATAORIGIN |
            TPMA_OBJECT_USERWITHAUTH |
            TPMA_OBJECT_DECRYPT,
        .parameters.eccDetail = {
            .symmetric.algorithm = TPM2_ALG_NULL,
            .scheme.scheme = TPM2_ALG_NULL,
            .curveID = TPM2_ECC_NIST_P256,
            .kdf.scheme = TPM2_ALG_NULL,
        },
    },
};

static const TPM2B_SENSITIVE_CREATE empty_in_sensitive = { 0 };
static const TPM2B_DATA             empty_outside_info = { 0 };
static const TPML_PCR_SELECTION     empty_creation_pcr = { 0 };

// Make sure TPM 2.0 is present and working correctly before using it.
stkb_rc tpm2_init(void)
{
    int           ret = 0;
    ESYS_CONTEXT *ctx;
    if (Esys_Initialize(&ctx, NULL, NULL) != TSS2_RC_SUCCESS)
        return STKB_TPM_UNAVAILABLE;
    // Simple self-test. Tests only what's necessary instead of all internal
    // functions and operations, doesn't disrupt the system.
    ret = Esys_SelfTest(ctx, ESYS_TR_NONE, ESYS_TR_NONE, ESYS_TR_NONE, TPM2_NO);
    Esys_Finalize(&ctx);
    if (ret != TSS2_RC_SUCCESS)
        return STKB_TPM_UNAVAILABLE;
    return STKB_SUCCESS;
}

static TSS2_RC tpm2_setup_primary(ESYS_CONTEXT **ctx, ESYS_TR *primary_handle)
{
    TSS2_RC rc = Esys_Initialize(ctx, NULL, NULL);
    if (rc != TSS2_RC_SUCCESS)
        return rc;
    rc = Esys_CreatePrimary(*ctx, ESYS_TR_RH_OWNER, ESYS_TR_PASSWORD,
                            ESYS_TR_NONE, ESYS_TR_NONE, &empty_in_sensitive,
                            &primary_template, &empty_outside_info,
                            &empty_creation_pcr, primary_handle, NULL, NULL,
                            NULL, NULL);
    if (rc != TSS2_RC_SUCCESS)
        Esys_Finalize(ctx);
    return rc;
}

static stkb_rc decode_tpm_rc(TSS2_RC tpm_rc)
{
    // Check that the code is from the TPM itself rather than ESAPI, otherwise
    // return a generic error as we're not interested in details.
    if ((tpm_rc & TSS2_RC_LAYER_MASK) != TSS2_TPM_RC_LAYER)
        return STKB_ERROR;

    // Format-One return codes are composite and in particular have an added
    // parameter value that is irrelevant to the nature of the error, clear it.
    if (tpm_rc & TPM2_RC_FMT1)
        tpm_rc &= ~TPM2_RC_N_MASK;

    switch (tpm_rc) {
    case TPM2_RC_BAD_AUTH:
    case TPM2_RC_AUTH_FAIL:
        return STKB_AUTH_FAILURE;
    case TPM2_RC_LOCKOUT:
        return STKB_LOCKOUT;
    default:
        return STKB_ERROR;
    }
}

stkb_rc tpm2_keygen(uint8_t *out_pub, uint8_t *out_buffer, size_t buffer_size,
                    size_t *written, stkb_user_input input)
{
    if (input.pin_len > TPM2_SHA256_DIGEST_SIZE)
        return STKB_BAD_PIN;
    TPM2B_SENSITIVE_CREATE in_sensitive = { 0 };
    memcpy(in_sensitive.sensitive.userAuth.buffer, input.pin, input.pin_len);
    in_sensitive.sensitive.userAuth.size = (UINT16)input.pin_len;
    ESYS_CONTEXT *ctx;
    ESYS_TR       primary_handle;
    stkb_rc       ret = STKB_ERROR;
    if (tpm2_setup_primary(&ctx, &primary_handle) != TSS2_RC_SUCCESS)
        return ret;
    TPM2B_PRIVATE *ecdh_priv = NULL;
    TPM2B_PUBLIC  *ecdh_pub  = NULL;
    TSS2_RC rc = Esys_Create(ctx, primary_handle, ESYS_TR_PASSWORD,
                             ESYS_TR_NONE, ESYS_TR_NONE, &in_sensitive,
                             &ecdh_key_template, &empty_outside_info,
                             &empty_creation_pcr, &ecdh_priv, &ecdh_pub, NULL,
                             NULL, NULL);
    if (rc != TSS2_RC_SUCCESS) {
        ret = decode_tpm_rc(rc);
        goto exit;
    }

    // TPM 2.0 Library, Part 1, 44.5.3 Padding:
    //
    // > In ECC points returned by the TPM, the x and y values, if non-empty,
    // > are required to be the size of their associated curve (e.g., 32 bytes
    // > for NIST P-256).
    if (ecdh_pub->publicArea.unique.ecc.x.size != 32 ||
        ecdh_pub->publicArea.unique.ecc.y.size != 32)
        goto exit;

    // Compress the public key: knowing x, it's enough to only encode parity of
    // y to fully reconstruct the point.
    if (ecdh_pub->publicArea.unique.ecc.y.buffer[31] & 1)
        out_pub[0] = 0x03;
    else
        out_pub[0] = 0x02;
    memcpy(out_pub + 1, ecdh_pub->publicArea.unique.ecc.x.buffer, 32);

    // Serialize both structures, we'll need them for Load.
    size_t off = 0;
    if (Tss2_MU_TPM2B_PRIVATE_Marshal(ecdh_priv, out_buffer, buffer_size, &off) !=
        TSS2_RC_SUCCESS)
        goto exit;
    if (Tss2_MU_TPM2B_PUBLIC_Marshal(ecdh_pub, out_buffer, buffer_size - off,
                                     &off) != TSS2_RC_SUCCESS)
        goto exit;

    *written = off;
    ret     = STKB_SUCCESS;

exit:
    Esys_Free(ecdh_priv);
    Esys_Free(ecdh_pub);
    Esys_FlushContext(ctx, primary_handle);
    Esys_Finalize(&ctx);
    return ret;
}

// pk is an uncompressed 64-byte point (x || y) of the peer.
stkb_rc tpm2_compute_ss(uint8_t *ss, const uint8_t *blob, size_t blob_len,
                        const uint8_t *pk, stkb_user_input input)
{
    if (input.pin_len > TPM2_SHA256_DIGEST_SIZE)
        return STKB_BAD_PIN;
    TPM2B_AUTH auth_value = { 0 };
    memcpy(auth_value.buffer, input.pin, input.pin_len);
    auth_value.size = (UINT16)input.pin_len;

    ESYS_CONTEXT *ctx;
    ESYS_TR       primary_handle;
    stkb_rc       ret = STKB_ERROR;
    if (tpm2_setup_primary(&ctx, &primary_handle) != TSS2_RC_SUCCESS)
        return ret;

    ESYS_TR          ecdh_key_handle = ESYS_TR_NONE;
    TPM2B_PRIVATE    ecdh_priv;
    TPM2B_PUBLIC     ecdh_pub;
    TPM2B_ECC_POINT *shared_point = NULL;
    size_t           off          = 0;
    if (Tss2_MU_TPM2B_PRIVATE_Unmarshal(blob, blob_len, &off, &ecdh_priv) !=
        TSS2_RC_SUCCESS)
        goto exit;
    if (Tss2_MU_TPM2B_PUBLIC_Unmarshal(blob, blob_len, &off, &ecdh_pub) !=
        TSS2_RC_SUCCESS)
        goto exit;

    TSS2_RC rc = Esys_Load(ctx, primary_handle, ESYS_TR_PASSWORD, ESYS_TR_NONE,
                           ESYS_TR_NONE, &ecdh_priv, &ecdh_pub,
                           &ecdh_key_handle);
    if (rc != TSS2_RC_SUCCESS) {
        ret = decode_tpm_rc(rc);
        goto exit;
    }

    if (Esys_TR_SetAuth(ctx, ecdh_key_handle, &auth_value) != TSS2_RC_SUCCESS)
        goto exit;

    TPM2B_ECC_POINT in_point;
    in_point.point.x.size = 32;
    in_point.point.y.size = 32;
    memcpy(in_point.point.x.buffer, pk, 32);
    memcpy(in_point.point.y.buffer, pk + 32, 32);
    rc = Esys_ECDH_ZGen(ctx, ecdh_key_handle, ESYS_TR_PASSWORD, ESYS_TR_NONE,
                        ESYS_TR_NONE, &in_point, &shared_point);
    if (rc != TSS2_RC_SUCCESS) {
        ret = decode_tpm_rc(rc);
        goto exit;
    }

    // Deliver the x coordinate of the shared point.
    memcpy(ss, shared_point->point.x.buffer, 32);

    ret = STKB_SUCCESS;

exit:
    Esys_Free(shared_point);
    if (primary_handle != ESYS_TR_NONE)
        Esys_FlushContext(ctx, primary_handle);
    if (ecdh_key_handle != ESYS_TR_NONE)
        Esys_FlushContext(ctx, ecdh_key_handle);
    Esys_Finalize(&ctx);
    return ret;
}