/*
** Copyright (C) 2020  Johannes Kepler University Linz, Institute of Networks and Security
** Copyright (C) 2020  CDL Digidow <https://www.digidow.eu/>
**
** Licensed under the EUPL, Version 1.2 or – as soon they will be approved by
** the European Commission - subsequent versions of the EUPL (the "Licence").
** You may not use this work except in compliance with the Licence.
** 
** You should have received a copy of the European Union Public License along
** with this program.  If not, you may obtain a copy of the Licence at:
** <https://joinup.ec.europa.eu/software/page/eupl>
** 
** Unless required by applicable law or agreed to in writing, software
** distributed under the Licence is distributed on an "AS IS" basis,
** WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
** See the Licence for the specific language governing permissions and
** limitations under the Licence.
**
*/

#include "member-tpm-key.h"

static TPMA_SESSION empty_session_attributes = {0};    // attributes for password either can't be set or don't make sense

struct test_context {
    TSS2_SYS_CONTEXT *sapi_ctx;
    TPM2_HANDLE primary_key_handle;
    TPM2_HANDLE signing_key_handle;
    TPM2_HANDLE persistent_key_handle;
    TPM2B_PUBLIC out_public;
    TPM2B_PRIVATE out_private;
    unsigned char tcti_buffer[256];
    unsigned char sapi_buffer[4200];

};

static void initialize(struct test_context *ctx);
static void cleanup(struct test_context *ctx);

static int clear(struct test_context *ctx);
static int create_primary(struct test_context *ctx);
static int create(struct test_context *ctx);
static int load(struct test_context *ctx);
static int save_public_key_info(const struct test_context* ctx, const char* pub_key_filename, const char* handle_filename);
static int evict_control(struct test_context *ctx);

static void initialize(struct test_context *ctx)
{
    const char *device_conf = "/dev/tpm0";
    size_t size;
    int init_ret;

    memset(ctx->tcti_buffer, 0, sizeof(ctx->tcti_buffer));
    memset(ctx->sapi_buffer, 0, sizeof(ctx->sapi_buffer));

    TSS2_TCTI_CONTEXT *tcti_ctx = (TSS2_TCTI_CONTEXT*)ctx->tcti_buffer;

    init_ret = Tss2_Tcti_Device_Init(NULL, &size, device_conf);
    if (TSS2_RC_SUCCESS != init_ret) {
        printf("Failed to get allocation size for tcti context\n");
        exit(1);
    }
    if (size > sizeof(ctx->tcti_buffer)) {
        printf("Error: device TCTI context size larger than pre-allocated buffer\n");
        exit(1);
    }
    init_ret = Tss2_Tcti_Device_Init(tcti_ctx, &size, device_conf);
    if (TSS2_RC_SUCCESS != init_ret) {
        printf("Error: Unable to initialize device TCTI context\n");
        exit(1);
    }

    ctx->sapi_ctx = (TSS2_SYS_CONTEXT*)ctx->sapi_buffer;
    size_t sapi_ctx_size = Tss2_Sys_GetContextSize(0);

    TSS2_ABI_VERSION abi_version = TSS2_ABI_VERSION_CURRENT;
    init_ret = Tss2_Sys_Initialize(ctx->sapi_ctx,
                                   sapi_ctx_size,
                                   tcti_ctx,
                                   &abi_version);
    ctx->out_public.size = 0;
    ctx->out_private.size = 0;
}

static void cleanup(struct test_context *ctx)
{
    TSS2_TCTI_CONTEXT *tcti_context = NULL;
    TSS2_RC rc;

    if (ctx->sapi_ctx != NULL) {
        rc = Tss2_Sys_GetTctiContext(ctx->sapi_ctx, &tcti_context);

        Tss2_Tcti_Finalize(tcti_context);

        Tss2_Sys_Finalize(ctx->sapi_ctx);
    }
}

int create_key(const char* pub_key_filename, const char* handle_filename)
{
    struct test_context ctx;
    initialize(&ctx);

    int ret = 0;

    ret = clear(&ctx);
    if(ret != TSS2_RC_SUCCESS) {
        printf("TPM Clear failed: %x\n",ret);
    }
  
    ret = create_primary(&ctx);
    if(ret != TSS2_RC_SUCCESS) {
        printf("TPM Create Primary failed: %x\n",ret);
    }

    ret = create(&ctx);
    if(ret != TSS2_RC_SUCCESS) {
        printf("TPM Create failed: %x\n",ret);
    }

    ret = load(&ctx);
    if(ret != TSS2_RC_SUCCESS) {
        printf("TPM Load failed: %x\n",ret);
    }

    ret = evict_control(&ctx);
    if(ret != TSS2_RC_SUCCESS) {
        printf("TPM Evict Control failed: %x\n",ret);
    }

    ret = save_public_key_info(&ctx, pub_key_filename, handle_filename);
    if(ret != TSS2_RC_SUCCESS) {
        printf("Save Public Key Info failed: %x\n",ret);
    }

    cleanup(&ctx);
    return ret;
}

static int save_public_key_info(const struct test_context *ctx, const char* pub_key_filename, const char* handle_filename)
{
    int write_ret = 0;

    FILE *pub_key_file_ptr = fopen(pub_key_filename, "w");

    if (NULL == pub_key_file_ptr)
        return -1;
    do {
        if (fprintf(pub_key_file_ptr, "%02X", 4) != 2)
            break;

        for (unsigned i=0; i < ctx->out_public.publicArea.unique.ecc.x.size; i++) {
            if (fprintf(pub_key_file_ptr, "%02X", ctx->out_public.publicArea.unique.ecc.x.buffer[i]) != 2) {
                write_ret = -1;
                break;
            }
        }
        if (0 != write_ret)
            break;

        for (unsigned i=0; i < ctx->out_public.publicArea.unique.ecc.y.size; i++) {
            if (fprintf(pub_key_file_ptr, "%02X", ctx->out_public.publicArea.unique.ecc.y.buffer[i]) != 2) {
                write_ret = -1;
                break;
            }
        }
        if (0 != write_ret)
            break;
    } while(0);
    (void)fclose(pub_key_file_ptr);

    (void)handle_filename;
    FILE *handle_file_ptr = fopen(handle_filename, "w");
    if (NULL == handle_file_ptr)
        return -1;
    write_ret = 0;
    do {
        for (int i=(sizeof(ctx->persistent_key_handle)-1); i >= 0; i--) {
            if (fprintf(handle_file_ptr, "%02X", (ctx->persistent_key_handle >> i*8) & 0xFF) != 2) {
                write_ret = -1;
                break;
            }
        }
        if (0 != write_ret)
            break;
    } while(0);
    (void)fclose(handle_file_ptr);

    return write_ret;
}

static int clear(struct test_context *ctx)
{
    TPMI_RH_CLEAR auth_handle = TPM2_RH_LOCKOUT;

    TSS2L_SYS_AUTH_COMMAND sessionsData = {};
    sessionsData.auths[0].sessionHandle = TPM2_RS_PW;
    sessionsData.auths[0].sessionAttributes = empty_session_attributes;
    sessionsData.count = 1;

    TSS2L_SYS_AUTH_RESPONSE sessionsDataOut = {};
    sessionsDataOut.count = 1;

    TSS2_RC ret = Tss2_Sys_Clear(ctx->sapi_ctx,
                                 auth_handle,
                                 &sessionsData,
                                 &sessionsDataOut);

    printf("Clear ret=%#X\n", ret);

    return ret;
}

static int create_primary(struct test_context *ctx)
{
    TPMI_RH_HIERARCHY hierarchy = TPM2_RH_ENDORSEMENT;

    TSS2L_SYS_AUTH_COMMAND sessionsData = {};
    sessionsData.auths[0].sessionHandle = TPM2_RS_PW;
    sessionsData.auths[0].sessionAttributes = empty_session_attributes;
    sessionsData.count = 1;

    TSS2L_SYS_AUTH_RESPONSE sessionsDataOut = {};
    sessionsDataOut.count = 1;

    TPM2B_SENSITIVE_CREATE inSensitive = {};

    TPM2B_PUBLIC in_public = {};
    in_public.publicArea.type = TPM2_ALG_ECC;
    in_public.publicArea.nameAlg = TPM2_ALG_SHA256;
    in_public.publicArea.objectAttributes = TPMA_OBJECT_FIXEDTPM |
        TPMA_OBJECT_FIXEDPARENT |
        TPMA_OBJECT_SENSITIVEDATAORIGIN |
        TPMA_OBJECT_USERWITHAUTH |
        TPMA_OBJECT_DECRYPT |
        TPMA_OBJECT_RESTRICTED;
    in_public.publicArea.parameters.eccDetail.symmetric.algorithm = TPM2_ALG_AES;
    in_public.publicArea.parameters.eccDetail.symmetric.keyBits.aes = 128;
    in_public.publicArea.parameters.eccDetail.symmetric.mode.sym = TPM2_ALG_CFB;
    in_public.publicArea.parameters.eccDetail.scheme.scheme = TPM2_ALG_NULL;
    in_public.publicArea.parameters.eccDetail.curveID = TPM2_ECC_NIST_P256;
    in_public.publicArea.parameters.eccDetail.kdf.scheme = TPM2_ALG_NULL;

    TPM2B_DATA outsideInfo = {};

    TPML_PCR_SELECTION creationPCR = {};

    TPM2B_CREATION_DATA creationData = {};
    TPM2B_DIGEST creationHash = {};
    TPMT_TK_CREATION creationTicket = {};

    TPM2B_NAME name = {};

    TPM2B_PUBLIC public_key = {};

    TSS2_RC ret = Tss2_Sys_CreatePrimary(ctx->sapi_ctx,
                                        hierarchy,
                                        &sessionsData,
                                        &inSensitive,
                                        &in_public,
                                        &outsideInfo,
                                        &creationPCR,
                                        &ctx->primary_key_handle,
                                        &public_key,
                                        &creationData,
                                        &creationHash,
                                        &creationTicket,
                                        &name,
                                        &sessionsDataOut);

    printf("CreatePrimary ret=%#X\n", ret);

    return ret;
}

static int create(struct test_context *ctx)
{
    TSS2L_SYS_AUTH_COMMAND sessionsData = {};
    sessionsData.auths[0].sessionHandle = TPM2_RS_PW;
    sessionsData.auths[0].sessionAttributes = empty_session_attributes;
    sessionsData.count = 1;

    TSS2L_SYS_AUTH_RESPONSE sessionsDataOut = {};
    sessionsDataOut.count = 1;

    TPM2B_SENSITIVE_CREATE inSensitive = {};

    TPM2B_PUBLIC in_public = {};
    in_public.publicArea.type = TPM2_ALG_ECC;
    in_public.publicArea.nameAlg = TPM2_ALG_SHA256;
    in_public.publicArea.objectAttributes = TPMA_OBJECT_FIXEDTPM |
        TPMA_OBJECT_FIXEDPARENT |
        TPMA_OBJECT_SENSITIVEDATAORIGIN |
        TPMA_OBJECT_USERWITHAUTH |
        TPMA_OBJECT_SIGN_ENCRYPT;
    in_public.publicArea.parameters.eccDetail.symmetric.algorithm = TPM2_ALG_NULL;
    in_public.publicArea.parameters.eccDetail.scheme.scheme = TPM2_ALG_ECDAA;
    in_public.publicArea.parameters.eccDetail.scheme.details.ecdaa.hashAlg = TPM2_ALG_SHA256;
    in_public.publicArea.parameters.eccDetail.scheme.details.ecdaa.count = 1;
    in_public.publicArea.parameters.eccDetail.curveID = TPM2_ECC_BN_P256;
    in_public.publicArea.parameters.eccDetail.kdf.scheme = TPM2_ALG_NULL;

    TPM2B_DATA outsideInfo = {};

    TPML_PCR_SELECTION creationPCR = {};

    TPM2B_CREATION_DATA creationData = {};
    TPM2B_DIGEST creationHash = {};
    TPMT_TK_CREATION creationTicket = {};

    TSS2_RC ret = Tss2_Sys_Create(ctx->sapi_ctx,
                                  ctx->primary_key_handle,
                                  &sessionsData,
                                  &inSensitive,
                                  &in_public,
                                  &outsideInfo,
                                  &creationPCR,
                                  &ctx->out_private,
                                  &ctx->out_public,
                                  &creationData,
                                  &creationHash,
                                  &creationTicket,
                                  &sessionsDataOut);

    printf("Create ret=%#X\n", ret);

    return ret;
}

static int load(struct test_context *ctx)
{
    TSS2L_SYS_AUTH_COMMAND sessionsData = {};
    sessionsData.auths[0].sessionHandle = TPM2_RS_PW;
    sessionsData.auths[0].sessionAttributes = empty_session_attributes;
    sessionsData.count = 1;

    TSS2L_SYS_AUTH_RESPONSE sessionsDataOut = {};
    sessionsDataOut.count = 1;

    TPM2B_NAME name = {};

    int ret = Tss2_Sys_Load(ctx->sapi_ctx,
                            ctx->primary_key_handle,
                            &sessionsData,
                            &ctx->out_private,
                            &ctx->out_public,
                            &ctx->signing_key_handle,
                            &name,
                            &sessionsDataOut);

    printf("Load ret=%#X\n", ret);

    return ret;
}

static int evict_control(struct test_context *ctx)
{
    TSS2L_SYS_AUTH_COMMAND sessionsData = {};
    sessionsData.auths[0].sessionHandle = TPM2_RS_PW;
    sessionsData.auths[0].sessionAttributes = empty_session_attributes;
    sessionsData.count = 1;

    TSS2L_SYS_AUTH_RESPONSE sessionsDataOut = {};
    sessionsDataOut.count = 1;

    ctx->persistent_key_handle = 0x81010000;

    TSS2_RC ret = Tss2_Sys_EvictControl(ctx->sapi_ctx,
                                        TPM2_RH_OWNER,
                                        ctx->signing_key_handle,
                                        &sessionsData,
                                        ctx->persistent_key_handle,
                                        &sessionsDataOut);

    printf("EvictControl ret=%#X\n", ret);

    return ret;
}