/* SPDX-License-Identifier: BSD-3-Clause */

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <openssl/rand.h>

#include "files.h"
#include "log.h"
#include "tpm2.h"
#include "tpm2_tool.h"
#include "tpm2_alg_util.h"
#include "tpm2_identity_util.h"
#include "tpm2_options.h"
#include "tpm2_openssl.h"

typedef struct tpm_makecred_ctx tpm_makecred_ctx;
struct tpm_makecred_ctx {
    TPM2B_NAME object_name;
    char *out_file_path;
    char *input_secret_data;
    char *public_key_path; /* path to the public portion of an object */
    TPM2B_PUBLIC public;
    TPM2B_DIGEST credential;
    struct {
        UINT8 e :1;
        UINT8 s :1;
        UINT8 n :1;
        UINT8 o :1;
    } flags;

    char *key_type; //type of key attempting to load, defaults to auto attempt
};

static tpm_makecred_ctx ctx = {
    .object_name = TPM2B_EMPTY_INIT,
    .public = TPM2B_EMPTY_INIT,
    .credential = TPM2B_EMPTY_INIT,
};

static bool write_cred_and_secret(const char *path, TPM2B_ID_OBJECT *cred,
        TPM2B_ENCRYPTED_SECRET *secret) {

    bool result = false;

    FILE *fp = fopen(path, "wb+");
    if (!fp) {
        LOG_ERR("Could not open file \"%s\" error: \"%s\"", path,
                strerror(errno));
        return false;
    }

    result = files_write_header(fp, 1);
    if (!result) {
        LOG_ERR("Could not write version header");
        goto out;
    }

    result = files_write_16(fp, cred->size);
    if (!result) {
        LOG_ERR("Could not write credential size");
        goto out;
    }

    result = files_write_bytes(fp, cred->credential, cred->size);
    if (!result) {
        LOG_ERR("Could not write credential data");
        goto out;
    }

    result = files_write_16(fp, secret->size);
    if (!result) {
        LOG_ERR("Could not write secret size");
        goto out;
    }

    result = files_write_bytes(fp, secret->secret, secret->size);
    if (!result) {
        LOG_ERR("Could not write secret data");
        goto out;
    }

    result = true;

out:
    fclose(fp);
    return result;
}

static tool_rc make_external_credential_and_save(void) {

    /*
     * Get name_alg from the public key
     */
    TPMI_ALG_HASH name_alg = ctx.public.publicArea.nameAlg;

    /*
     * Generate and encrypt seed
     */
    TPM2B_DIGEST seed = TPM2B_TYPE_INIT(TPM2B_DIGEST, buffer);
    TPM2B_ENCRYPTED_SECRET encrypted_seed = TPM2B_EMPTY_INIT;
    unsigned char label[10] = { 'I', 'D', 'E', 'N', 'T', 'I', 'T', 'Y', 0 };
    bool res = tpm2_identity_util_share_secret_with_public_key(&seed,
            &ctx.public, label, 9, &encrypted_seed);
    if (!res) {
        LOG_ERR("Failed Seed Encryption\n");
        return tool_rc_general_error;
    }

    /*
     * Perform identity structure calculations (off of the TPM)
     */
    TPM2B_MAX_BUFFER hmac_key;
    TPM2B_MAX_BUFFER enc_key;
    tpm2_identity_util_calc_outer_integrity_hmac_key_and_dupsensitive_enc_key(
            &ctx.public, &ctx.object_name, &seed, &hmac_key, &enc_key);

    /*
     * The ctx.credential needs to be marshalled into struct with
     * both size and contents together (to be encrypted as a block)
     */
    TPM2B_MAX_BUFFER marshalled_inner_integrity = TPM2B_EMPTY_INIT;
    marshalled_inner_integrity.size = ctx.credential.size
            + sizeof(ctx.credential.size);
    UINT16 cred_size = ctx.credential.size;
    if (!tpm2_util_is_big_endian()) {
        cred_size = tpm2_util_endian_swap_16(cred_size);
    }
    memcpy(marshalled_inner_integrity.buffer, &cred_size, sizeof(cred_size));
    memcpy(&marshalled_inner_integrity.buffer[2], ctx.credential.buffer,
            ctx.credential.size);

    /*
     * Perform inner encryption (encIdentity) and outer HMAC (outerHMAC)
     */
    TPM2B_DIGEST outer_hmac = TPM2B_EMPTY_INIT;
    TPM2B_MAX_BUFFER encrypted_sensitive = TPM2B_EMPTY_INIT;
    bool outer_res = tpm2_identity_util_calculate_outer_integrity(name_alg,
            &ctx.object_name, &marshalled_inner_integrity, &hmac_key, &enc_key,
            &ctx.public.publicArea.parameters.rsaDetail.symmetric,
            &encrypted_sensitive, &outer_hmac);
    if (!outer_res) {
        return tool_rc_general_error;
    }

    /*
     * Package up the info to save
     * cred_bloc = outer_hmac || encrypted_sensitive
     * secret = encrypted_seed (with pubEK)
     */
    TPM2B_ID_OBJECT cred_blob = TPM2B_TYPE_INIT(TPM2B_ID_OBJECT, credential);

    UINT16 outer_hmac_size = outer_hmac.size;
    if (!tpm2_util_is_big_endian()) {
        outer_hmac_size = tpm2_util_endian_swap_16(outer_hmac_size);
    }
    int offset = 0;
    memcpy(cred_blob.credential + offset, &outer_hmac_size,
            sizeof(outer_hmac.size));
    offset += sizeof(outer_hmac.size);
    memcpy(cred_blob.credential + offset, outer_hmac.buffer, outer_hmac.size);
    offset += outer_hmac.size;
    //NOTE: do NOT include the encrypted_sensitive size, since it is encrypted with the blob!
    memcpy(cred_blob.credential + offset, encrypted_sensitive.buffer,
            encrypted_sensitive.size);

    cred_blob.size = outer_hmac.size + encrypted_sensitive.size
            + sizeof(outer_hmac.size);

    return write_cred_and_secret(ctx.out_file_path, &cred_blob,
            &encrypted_seed) ? tool_rc_success : tool_rc_general_error;
}

static tool_rc make_credential_and_save(ESYS_CONTEXT *ectx) {
    TPM2B_ID_OBJECT *cred_blob;
    TPM2B_ENCRYPTED_SECRET *secret;
    ESYS_TR tr_handle = ESYS_TR_NONE;

    tool_rc rc = tpm2_loadexternal(ectx, 0, &ctx.public, TPM2_RH_NULL,
        &tr_handle, 0, TPM2_ALG_ERROR);
    if (rc != tool_rc_success) {
        return rc;
    }

    rc = tpm2_makecredential(ectx, tr_handle,
            &ctx.credential, &ctx.object_name, &cred_blob,
            &secret);
    if (rc != tool_rc_success) {
        return rc;
    }

    rc = tpm2_flush_context(ectx, tr_handle, NULL, TPM2_ALG_NULL);
    if (rc != tool_rc_success) {
        free(cred_blob);
        free(secret);
        return rc;
    }

    bool ret = write_cred_and_secret(ctx.out_file_path, cred_blob, secret);
    free(cred_blob);
    free(secret);
    return ret ? tool_rc_success : tool_rc_general_error;
}

static bool on_option(char key, char *value) {

    tool_rc rc;
    switch (key) {
    case 'u':
        if (ctx.flags.e) {
            LOG_ERR("Specify public key with **-u** or **-e**, not both");
            return false;
        }
        ctx.public_key_path = value;
        ctx.flags.e = 1;
        break;
    case 'e':
        if (ctx.flags.e) {
            LOG_ERR("Specify encryption key with **-u** or **-e**, not both");
            return false;
        }
        ctx.public_key_path = value;
        ctx.flags.e = 1;
        break;
    case 's':
        ctx.input_secret_data = strcmp("-", value) ? value : NULL;
        ctx.flags.s = 1;
        break;
    case 'n':
        ctx.object_name.size = BUFFER_SIZE(TPM2B_NAME, name);
        rc = tpm2_util_bin_from_hex_or_file(value,
               &ctx.object_name.size, ctx.object_name.name) ?
               tool_rc_success : tool_rc_general_error;

        if (rc != tool_rc_success) {
            LOG_ERR("Could not load name data");
            return false;
        }
 
        ctx.flags.n = 1;
        break;
    case 'o':
        ctx.out_file_path = value;
        ctx.flags.o = 1;
        break;
    case 'G':
        ctx.key_type = value;
        break;
    }

    return true;
}

static bool tpm2_tool_onstart(tpm2_options **opts) {

    const struct option topts[] = {
      {"encryption-key",  required_argument, NULL, 'e'},
      {"public",          required_argument, NULL, 'u'},
      {"secret",          required_argument, NULL, 's'},
      {"name",            required_argument, NULL, 'n'},
      {"credential-blob", required_argument, NULL, 'o'},
      { "key-algorithm",  required_argument, NULL, 'G'},
    };

    *opts = tpm2_options_new("G:u:e:s:n:o:", ARRAY_LEN(topts), topts, on_option,
        NULL, TPM2_OPTIONS_OPTIONAL_SAPI);

    return *opts != NULL;
}

static void set_default_TCG_EK_template(TPMI_ALG_PUBLIC alg) {
    /* Values for RSA-2048 and ECC-256 keys. */
    TPMI_AES_KEY_BITS bits = 128;
    TPMI_ALG_HASH name_alg = TPM2_ALG_SHA256;

    ctx.public.publicArea.objectAttributes = 0;

    switch (alg) {
        case TPM2_ALG_RSA:
            if (ctx.public.publicArea.parameters.rsaDetail.keyBits > 2048) {
                /* Values for RSA-3072 and RSA-4096 keys. */
                bits = 256;
                name_alg = TPM2_ALG_SHA384;
                ctx.public.publicArea.objectAttributes = TPMA_OBJECT_USERWITHAUTH;
            }
            ctx.public.publicArea.parameters.rsaDetail.symmetric.algorithm =
                    TPM2_ALG_AES;
            ctx.public.publicArea.parameters.rsaDetail.symmetric.keyBits.aes = bits;
            ctx.public.publicArea.parameters.rsaDetail.symmetric.mode.aes =
                    TPM2_ALG_CFB;
            break;
        case TPM2_ALG_ECC:
            if (ctx.public.publicArea.unique.ecc.x.size > 32) {
                /* Values for ECC-384 keys. */
                bits = 256;
                name_alg = TPM2_ALG_SHA384;
                ctx.public.publicArea.objectAttributes = TPMA_OBJECT_USERWITHAUTH;
            }
            if (ctx.public.publicArea.unique.ecc.x.size > 48) {
                /* Values for ECC-512 keys. */
                name_alg = TPM2_ALG_SHA512;
            }

            ctx.public.publicArea.parameters.eccDetail.symmetric.algorithm =
                    TPM2_ALG_AES;
            ctx.public.publicArea.parameters.eccDetail.symmetric.keyBits.aes = bits;
            ctx.public.publicArea.parameters.eccDetail.symmetric.mode.sym =
                    TPM2_ALG_CFB;
            break;
    }

    ctx.public.publicArea.objectAttributes |=
          TPMA_OBJECT_RESTRICTED  | TPMA_OBJECT_ADMINWITHPOLICY
        | TPMA_OBJECT_DECRYPT     | TPMA_OBJECT_FIXEDTPM
        | TPMA_OBJECT_FIXEDPARENT | TPMA_OBJECT_SENSITIVEDATAORIGIN;

    ctx.public.publicArea.nameAlg = name_alg;
}

static tool_rc process_input(tpm2_option_flags flags) {

    TPMI_ALG_PUBLIC alg = TPM2_ALG_NULL;
    if (ctx.key_type) {
        if (!flags.quiet) {
            LOG_WARN("Because **-G** is specified, assuming input encryption "
                     "public key is in PEM format.");
        }
        alg = tpm2_alg_util_from_optarg(ctx.key_type,
            tpm2_alg_util_flags_asymmetric);
        if (alg == TPM2_ALG_ERROR ||
           (alg != TPM2_ALG_RSA && alg != TPM2_ALG_ECC)) {
            LOG_ERR("Unsupported key type, got: \"%s\"", ctx.key_type);
            return tool_rc_general_error;
        }
    }

    if (ctx.public_key_path) {
        bool result = alg != TPM2_ALG_NULL ?
            tpm2_openssl_load_public(ctx.public_key_path, alg,
            &ctx.public) : files_load_public(ctx.public_key_path, &ctx.public);
        if (!result) {
            return tool_rc_general_error;
        }
    }

    /*
     * Since it is a PEM we will fixate the key properties from TCG EK
     * template since we had to choose "a template".
     */
    if (ctx.key_type) {
        set_default_TCG_EK_template(alg);
    }

    if (!ctx.flags.s) {
        LOG_ERR("Specify the secret either as a file or a '-' for stdin");
        return tool_rc_option_error;
    }

    if (!ctx.flags.e || !ctx.flags.n || !ctx.flags.o) {
        LOG_ERR("Expected mandatory options e, n, o.");
        return tool_rc_option_error;
    }

    /*
     * Maximum size of the allowed secret-data size  to fit in TPM2B_DIGEST
     */
    ctx.credential.size = TPM2_SHA512_DIGEST_SIZE;

    bool result = files_load_bytes_from_buffer_or_file_or_stdin(NULL,
        ctx.input_secret_data, &ctx.credential.size, ctx.credential.buffer);
    if (!result) {
        return tool_rc_general_error;
    }

    /*
     * If input was read from stdin, check if a larger data set was specified
     * and error out.
     */
    if (ctx.credential.size > TPM2_SHA512_DIGEST_SIZE) {
        LOG_ERR("Size is larger than buffer, got %d expected less than or equal"
        "to %d", ctx.credential.size, TPM2_SHA512_DIGEST_SIZE);
        return tool_rc_general_error;
    }

    return tool_rc_success;
}

static tool_rc tpm2_tool_onrun(ESYS_CONTEXT *ectx, tpm2_option_flags flags) {

    UNUSED(flags);

    tool_rc rc = process_input(flags);
    if (rc != tool_rc_success) {
        return rc;
    }

    // Run it outside of a TPM
    return ectx ?
            make_credential_and_save(ectx) :
                make_external_credential_and_save();
}

// Register this tool with tpm2_tool.c
TPM2_TOOL_REGISTER("makecredential", tpm2_tool_onstart, tpm2_tool_onrun, NULL, NULL)
