/* =========================================================================
 * Copyright (c) 2020, 2023 Qualcomm Technologies, Inc.
 * All Rights Reserved.
 * Confidential and Proprietary - Qualcomm Technologies, Inc.
 * =========================================================================
 */

#include <stdbool.h>
#include <string.h>
#include <stringl.h>
#include "CCipher.h"
#include "CCrypto.h"
#include "CHWKeyFactory.h"
#include "CHmac.h"
#include "ICipher.h"
#include "ICrypto.h"
#include "IHWKeyFactory.h"
#include "IHmac.h"
#include "IOpener.h"
#include "crypto_hwkm_test_vectors.h"
#include "object.h"
#include "qsee_env.h"
#include "qsee_heap.h"
#include "qsee_log.h"
#include "qsee_utils.h"
#include "sampleapp.h"

#define ASSERT_OK(func)                                                \
    if (!Object_isOK(func)) {                                          \
        qsee_log(QSEE_LOG_MSG_ERROR,                                   \
                 "HW Key factory test %s failed on line %d", __func__, \
                 __LINE__);                                            \
        goto __Exit;                                                   \
    }

// To test IDL interface which return objects, we check if return code
// shows Object is ok. If it is not, we set the output object to
// Object_NULL to avoid releasing it
#define ASSERT_OK_CLEAN(func, output_obj)                              \
    if (!Object_isOK(func)) {                                          \
        output_obj = Object_NULL;                                      \
        qsee_log(QSEE_LOG_MSG_ERROR,                                   \
                 "HW Key factory test %s failed on line %d", __func__, \
                 __LINE__);                                            \
        goto __Exit;                                                   \
    }

#define ASSERT_TRUE(cond)                                              \
    if (!(cond)) {                                                     \
        qsee_log(QSEE_LOG_MSG_ERROR,                                   \
                 "HW Key factory test %s failed on line %d", __func__, \
                 __LINE__);                                            \
        goto __Exit;                                                   \
    }

#define ASSERT_ERROR(func)                                             \
    if (!Object_isERROR(func)) {                                       \
        qsee_log(QSEE_LOG_MSG_ERROR,                                   \
                 "HW Key factory test %s failed on line %d", __func__, \
                 __LINE__);                                            \
        goto __Exit;                                                   \
    }

#define ASSERT_ERROR_CLEAN(func, output_obj)                           \
    if (!Object_isERROR(func)) {                                       \
        qsee_log(QSEE_LOG_MSG_ERROR,                                   \
                 "HW Key factory test %s failed on line %d", __func__, \
                 __LINE__);                                            \
        goto __Exit;                                                   \
    } else {                                                           \
        output_obj = Object_NULL;                                      \
    }

#define TEST(func, passed, failed)                          \
    if ((func) == SAMPLEAPP_SUCCESS) {                      \
        qsee_log(QSEE_LOG_MSG_DEBUG, "Passed Test " #func); \
        (passed)++;                                         \
    } else {                                                \
        qsee_log(QSEE_LOG_MSG_ERROR, "Failed Test " #func); \
        (failed)++;                                         \
    }

/**
 * Simple test case showing basic uses of deriveFromHWSource()
 *
 * Derives a key from HW source and ensures that there are no errors
 */
int test_IHWKeyFactory_deriveFromHWSource(Object factory)
{
    int status = SAMPLEAPP_FAILED;
    Object wrap_key = Object_NULL;
    Object op_key = Object_NULL;
    Object no_throttle_key = Object_NULL;
    Object mixing_key = Object_NULL;
    IHWKey_policy mixing_policy = {.keyType = IHWKey_KEY_TYPE_DERIVATION};
    IHWKey_policy wrap_policy = {.keyType = IHWKey_KEY_TYPE_WRAPPING,
                                 .kdfDepth = 10};
    IHWKey_policy op_policy = {
        .keyType = IHWKey_KEY_TYPE_OPERATIONAL,
        .kdfDepth = 10,
        .opPolicy = {.algorithm = IHWKey_ALGORITHM_AES128_CBC,
                     .operations = IHWKey_OPERATION_ENCRYPT_ALLOWED}};
    IHWKey_policy no_throttle_policy = {.keyType = IHWKey_KEY_TYPE_DERIVATION};
    IHWKey_bindings bindings = {.bindings = IHWKey_BINDING_KEY_POLICY_VERSION};

    ASSERT_OK_CLEAN(
        IHWKeyFactory_generate(factory, &mixing_policy, &mixing_key),
        mixing_key);

    ASSERT_OK_CLEAN(
        IHWKeyFactory_deriveFromHWSource(
            factory, &op_policy, &bindings,
            IHWKeyFactory_HW_SOURCE_CHIP_RANDOM_KEY, mixing_key, &op_key),
        op_key);

    ASSERT_OK_CLEAN(
        IHWKeyFactory_deriveFromHWSource(
            factory, &wrap_policy, &bindings,
            IHWKeyFactory_HW_SOURCE_CHIP_RANDOM_KEY, mixing_key, &wrap_key),
        wrap_key);

    ASSERT_OK_CLEAN(IHWKeyFactory_deriveFromHWSource(
                        factory, &no_throttle_policy, &bindings,
                        IHWKeyFactory_HW_SOURCE_CHIP_RANDOM_KEY, Object_NULL,
                        &no_throttle_key),
                    no_throttle_key);

    status = SAMPLEAPP_SUCCESS;

__Exit:
    Object_ASSIGN_NULL(wrap_key);
    Object_ASSIGN_NULL(op_key);
    Object_ASSIGN_NULL(no_throttle_key);
    Object_ASSIGN_NULL(mixing_key);
    return status;
}

/**
 * Simple test case showing a basic use of generate()
 *
 * Generates a key and ensures that there are no errors
 */
int test_IHWKeyFactory_generate(Object factory)
{
    int status = SAMPLEAPP_FAILED;
    Object key = Object_NULL;
    IHWKey_policy policy = {
        .keyType = IHWKey_KEY_TYPE_OPERATIONAL,
        .opPolicy = {.operations = IHWKey_OPERATION_ENCRYPT_ALLOWED |
                                   IHWKey_OPERATION_DECRYPT_ALLOWED,
                     .algorithm = IHWKey_ALGORITHM_AES256_ECB}};

    ASSERT_OK_CLEAN(IHWKeyFactory_generate(factory, &policy, &key), key);
    status = SAMPLEAPP_SUCCESS;

__Exit:
    Object_ASSIGN_NULL(key);
    return status;
}

/**
 * Simple test case showing a basic use of importKey()
 *
 * Imports a key and ensures that there are no errors
 * Uses the key for crypto and ensures that the output is correct
 */
int test_IHWKeyFactory_importKey(Object factory)
{
    int status = SAMPLEAPP_FAILED;
    Object key = Object_NULL;
    Object ctx = Object_NULL;
    const uint8_t sw_key[] = {0x33, 0x88, 0x84, 0xF4, 0xF4, 0x47, 0xC9, 0x13,
                              0xC8, 0x6E, 0x96, 0x8C, 0xF2, 0xB0, 0x32, 0xBC,
                              0xE8, 0xF6, 0x8B, 0x74, 0x07, 0x57, 0x74, 0x5C,
                              0x6B, 0xF8, 0xDF, 0x29, 0x17, 0x5E, 0xE5, 0xAD};
    const uint8_t iv[] = {0x9E, 0x72, 0xE7, 0x8D, 0x5E, 0x26, 0xBD, 0x04,
                          0xCD, 0xF4, 0xA9, 0xD4, 0x21, 0xCC, 0xD3, 0x86};
    const uint8_t pt[] = {0x5B, 0xFA, 0x4F, 0xF1, 0x58, 0x70, 0xB0, 0x30,
                          0xC9, 0x4A, 0xBF, 0xD3, 0x37, 0x25, 0x77, 0x9E};
    const uint8_t ct[] = {0x58, 0xF6, 0x4B, 0x2C, 0x78, 0xFD, 0x30, 0x29,
                          0xE7, 0xEF, 0xE5, 0xD3, 0x83, 0x29, 0xFC, 0xAD};
    uint8_t output[sizeof(pt)] = {0};
    size_t outlen = 0;
    IHWKey_policy policy = {
        .keyType = IHWKey_KEY_TYPE_OPERATIONAL,
        .destination = IHWKey_DESTINATION_DEFAULT,
        .opPolicy = {.algorithm = IHWKey_ALGORITHM_AES256_CBC,
                     .operations = IHWKey_OPERATION_ENCRYPT_ALLOWED |
                                   IHWKey_OPERATION_DECRYPT_ALLOWED}};

    ASSERT_OK(qsee_open(CCipherAES256_UID, &ctx));

    ASSERT_OK_CLEAN(
        IHWKeyFactory_importKey(factory, &policy, sw_key, sizeof(sw_key), &key),
        key);

    ASSERT_OK(ICipher_setParamAsU32(ctx, ICipher_PARAM_MODE, ICipher_MODE_CBC));
    ASSERT_OK(ICipher_setParamAsData(ctx, ICipher_PARAM_IV, iv, sizeof(iv)));
    ASSERT_OK(ICipher_setParamAsObject(ctx, ICipher_PARAM_HW_KEY_OBJECT, key));

    ASSERT_OK(
        ICipher_encrypt(ctx, pt, sizeof(pt), output, sizeof(output), &outlen));

    ASSERT_TRUE(timesafe_memcmp(ct, output, sizeof(ct)) == 0);

    status = SAMPLEAPP_SUCCESS;

__Exit:
    Object_ASSIGN_NULL(key);
    Object_ASSIGN_NULL(ctx);
    return status;
}

/**
 * Test case demonstrating a sample call flow to wrap a key and unwrap it
 * A client who wanted to persist a key would follow this call flow
 *
 * Workflow to persist key:
 * key = generate(...) (or derive(), deriveFromHWSource(), importKey())
 * wrapping_key = deriveFromHWSource(..., bindings, source, ...)
 * blob = key.wrap(wrapping_key)
 *
 * persist(blob)
 *
 * Workflow to unwrap key from from persistent storage:
 * wrapping_key = deriveFromHWSource(..., bindings, source, ...)
 * key = unwrap(wrapping_key, blob)
 */
int test_IHWKeyFactory_persist_usecase(Object factory)
{
    int status = SAMPLEAPP_FAILED;
    Object key = Object_NULL;
    Object wrapping_key = Object_NULL;
    Object unwrapped_key = Object_NULL;
    uint8_t* wrapped_key = NULL;
    uint8_t* wrapped_key2 = NULL;
    uint32_t wrapped_key_len = 0;
    size_t len_out = 0;
    IHWKey_policy policy = {
        .permissions = IHWKey_PERMISSION_WRAP_KEY_EXPORT,
        .keyType = IHWKey_KEY_TYPE_OPERATIONAL,
        .opPolicy = {.algorithm = IHWKey_ALGORITHM_AES256_ECB,
                     .operations = IHWKey_OPERATION_ENCRYPT_ALLOWED}};
    IHWKey_policy wrapping_policy = {.keyType = IHWKey_KEY_TYPE_WRAPPING};
    IHWKey_bindings wrapping_bindings = {
        .bindings =
            IHWKey_BINDING_KEY_POLICY_VERSION | IHWKey_BINDING_CHILD_KEY_POLICY,
        .context = {0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99,
                    0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF},
        .contextLen = 16};

    ASSERT_OK_CLEAN(IHWKeyFactory_generate(factory, &policy, &key), key);
    ASSERT_OK_CLEAN(IHWKeyFactory_deriveFromHWSource(
                        factory, &wrapping_policy, &wrapping_bindings,
                        IHWKeyFactory_HW_SOURCE_CHIP_RANDOM_KEY, Object_NULL,
                        &wrapping_key),
                    wrapping_key);

    ASSERT_OK(IHWKey_getBlobSize(wrapping_key, &wrapped_key_len));

    wrapped_key = QSEE_ZALLOC_ARRAY(uint8_t, wrapped_key_len);
    ASSERT_TRUE(wrapped_key != NULL);

    wrapped_key2 = QSEE_ZALLOC_ARRAY(uint8_t, wrapped_key_len);
    ASSERT_TRUE(wrapped_key2 != NULL);

    ASSERT_OK(
        IHWKey_wrap(key, wrapping_key, wrapped_key, wrapped_key_len, &len_out));

    ASSERT_OK_CLEAN(IHWKeyFactory_unwrap(factory, wrapping_key, wrapped_key,
                                         wrapped_key_len, &unwrapped_key),
                    unwrapped_key);

    ASSERT_OK(IHWKey_wrap(unwrapped_key, wrapping_key, wrapped_key2,
                          wrapped_key_len, &len_out));

    ASSERT_TRUE(timesafe_memcmp(wrapped_key, wrapped_key2, wrapped_key_len) ==
                0);

    status = SAMPLEAPP_SUCCESS;

__Exit:
    QSEE_FREE_PTR(wrapped_key);
    QSEE_FREE_PTR(wrapped_key2);

    Object_ASSIGN_NULL(key);
    Object_ASSIGN_NULL(wrapping_key);
    Object_ASSIGN_NULL(unwrapped_key);
    return status;
}

/**
 * Simple test case showing a basic use of derive()
 *
 * Derives a key from a parent key and ensures that there are no errors
 */
int test_IHWKeyFactory_derive(Object factory)
{
    int status = SAMPLEAPP_FAILED;
    Object derivation_key = Object_NULL;
    Object key = Object_NULL;
    IHWKey_policy policy = {.keyType = IHWKey_KEY_TYPE_DERIVATION};
    IHWKey_bindings bindings = {.bindings = IHWKey_BINDING_KEY_POLICY_VERSION};

    ASSERT_OK_CLEAN(IHWKeyFactory_generate(factory, &policy, &derivation_key),
                    derivation_key);
    ASSERT_OK_CLEAN(IHWKeyFactory_derive(factory, &policy, &bindings,
                                         derivation_key, Object_NULL, &key),
                    key);

    status = SAMPLEAPP_SUCCESS;

__Exit:
    Object_ASSIGN_NULL(derivation_key);
    Object_ASSIGN_NULL(key);
    return status;
}

/**
 * Run a cipher test case using a HW key
 *
 * Uses importKey() to generate key object
 * Passes key object to ICipher using setParamAsObject()
 * Sets up other required ICipher parameters for specified cipher algorithm
 * Encrypts plaintext and compares result to ciphertext
 * Decrypts ciphertext and compares result to plaintext
 */
int test_IHWKeyFactory_ICipher(Object factory, bool encrypt,
                               const hwkm_cipher_test_vector_t* test)
{
    int status = SAMPLEAPP_FAILED;
    Object key = Object_NULL;
    Object ctx = Object_NULL;
    uint8_t* output = NULL;
    size_t outlen = test->datalen;
    size_t datalen = test->datalen;
    const uint8_t* data = encrypt ? test->pt : test->ct;
    const uint8_t* expected = encrypt ? test->ct : test->pt;
    IHWKey_policy policy = {
        .keyType = IHWKey_KEY_TYPE_OPERATIONAL,
        .destination = IHWKey_DESTINATION_DEFAULT,
        .opPolicy = {.algorithm = test->alg,
                     .operations = IHWKey_OPERATION_ENCRYPT_ALLOWED |
                                   IHWKey_OPERATION_DECRYPT_ALLOWED}};

    ASSERT_OK_CLEAN(qsee_open(test->uid, &ctx), ctx);

    ASSERT_OK_CLEAN(IHWKeyFactory_importKey(factory, &policy, test->key,
                                            test->keylen, &key),
                    key);

    if (test->mode == ICipher_MODE_CCM) {
        if (encrypt) {
            outlen += test->maclen;
        } else {
            datalen += test->maclen;
        }
    }

    output = QSEE_ZALLOC_ARRAY(uint8_t, outlen);
    ASSERT_TRUE(output != NULL);

    ASSERT_OK(ICipher_setParamAsObject(ctx, ICipher_PARAM_HW_KEY_OBJECT, key));
    ASSERT_OK(ICipher_setParamAsU32(ctx, ICipher_PARAM_MODE, test->mode));

    if (test->iv != NULL) {
        ASSERT_OK(ICipher_setParamAsData(ctx, ICipher_PARAM_IV, test->iv,
                                         test->ivlen));
    }

    if (test->mode == ICipher_MODE_CCM) {
        ASSERT_OK(ICipher_setParamAsData(ctx, ICipher_PARAM_NONCE, test->nonce,
                                         test->noncelen));
        ASSERT_OK(ICipher_setParamAsU32(ctx, ICipher_PARAM_CCM_HDR_LEN,
                                        test->headerlen));
        ASSERT_OK(ICipher_setParamAsU32(ctx, ICipher_PARAM_CCM_MAC_LEN,
                                        test->maclen));
        ASSERT_OK(ICipher_setParamAsU32(ctx, ICipher_PARAM_CCM_PAYLOAD_LEN,
                                        test->payloadlen));
    }

    if (encrypt) {
        ASSERT_OK(ICipher_encrypt(ctx, data, datalen, output, outlen, &outlen));
    } else {
        ASSERT_OK(ICipher_decrypt(ctx, data, datalen, output, outlen, &outlen));
    }
    ASSERT_TRUE(timesafe_memcmp(output, expected, outlen) == 0);

    status = SAMPLEAPP_SUCCESS;

__Exit:
    QSEE_FREE_PTR(output);
    Object_ASSIGN_NULL(key);
    Object_ASSIGN_NULL(ctx);
    return status;
}

int test_IHWKeyFactory_IHmac(Object factory,
                             const hwkm_hmac_test_vector_t* test)
{
    int status = SAMPLEAPP_FAILED;
    Object key = Object_NULL;
    Object ctx = Object_NULL;
    uint8_t* output = NULL;
    size_t outlen = 0;
    IHWKey_policy policy = {
        .keyType = IHWKey_KEY_TYPE_OPERATIONAL,
        .destination = IHWKey_DESTINATION_DEFAULT,
        .opPolicy = {.algorithm = test->alg,
                     .operations = IHWKey_OPERATION_ENCRYPT_ALLOWED |
                                   IHWKey_OPERATION_DECRYPT_ALLOWED}};

    output = QSEE_ZALLOC_ARRAY(uint8_t, test->digestlen);
    ASSERT_TRUE(output != NULL);

    ASSERT_OK_CLEAN(qsee_open(test->uid, &ctx), ctx);

    ASSERT_OK_CLEAN(IHWKeyFactory_importKey(factory, &policy, test->key,
                                            test->keylen, &key),
                    key);

    ASSERT_OK(IHmac_setParamAsObject(ctx, IHmac_PARAM_KEY, key));
    ASSERT_OK(IHmac_update(ctx, test->msg, test->msglen));
    ASSERT_OK(IHmac_final(ctx, output, test->digestlen, &outlen));

    ASSERT_TRUE(timesafe_memcmp(output, test->digest, test->digestlen) == 0);

    status = SAMPLEAPP_SUCCESS;

__Exit:
    QSEE_FREE_PTR(output);
    Object_ASSIGN_NULL(key);
    Object_ASSIGN_NULL(ctx);
    return status;
}

int test_IHWKeyFactory_cmac(Object factory, const hwkm_cmac_test_vector_t* test)
{
    int status = SAMPLEAPP_FAILED;
    Object key = Object_NULL;
    Object ctx = Object_NULL;
    uint8_t* output = NULL;
    size_t outlen = 0;
    IHWKey_policy policy = {
        .keyType = IHWKey_KEY_TYPE_OPERATIONAL,
        .destination = IHWKey_DESTINATION_DEFAULT,
        .opPolicy = {.algorithm = test->alg,
                     .operations = IHWKey_OPERATION_ENCRYPT_ALLOWED |
                                   IHWKey_OPERATION_DECRYPT_ALLOWED}};

    output = QSEE_ZALLOC_ARRAY(uint8_t, test->digestlen);
    ASSERT_TRUE(output != NULL);

    ASSERT_OK_CLEAN(qsee_open(CCrypto_UID, &ctx), ctx);

    ASSERT_OK_CLEAN(IHWKeyFactory_importKey(factory, &policy, test->key,
                                            test->keylen, &key),
                    key);

    ASSERT_OK(ICrypto_hwkey_cmac(ctx, test->cmac_alg, key, test->msg,
                                 test->msglen, output, test->digestlen,
                                 &outlen));

    ASSERT_TRUE(timesafe_memcmp(output, test->digest, test->digestlen) == 0);

    status = SAMPLEAPP_SUCCESS;

__Exit:
    QSEE_FREE_PTR(output);
    Object_ASSIGN_NULL(key);
    Object_ASSIGN_NULL(ctx);
    return status;
}

/**
 * Test IHWKeyFactory with all algorithms supported by QTEE interfaces
 */
int test_IHWKeyFactory_all_algorithms(Object factory)
{
    uint32_t passed = 0;
    uint32_t failed = 0;
    uint32_t skipped = 0;

    for (size_t i = 0; i < ARR_SIZE(hwkm_cipher_test_vectors); i++) {
        hwkm_cipher_test_vector_t test = hwkm_cipher_test_vectors[i];

        if (test.enabled) {
            if (test_IHWKeyFactory_ICipher(factory, true, &test) ==
                SAMPLEAPP_SUCCESS) {
                qsee_log(QSEE_LOG_MSG_DEBUG,
                         "Passed hwkm cipher encrypt test case %d", i);
                passed++;
            } else {
                qsee_log(QSEE_LOG_MSG_ERROR,
                         "Failed hwkm cipher encrypt test case %d", i);
                failed++;
            }
        } else {
            qsee_log(QSEE_LOG_MSG_ERROR,
                     "Skipped hwkm cipher encrypt test case %d", i);
            skipped++;
        }
    }

    for (size_t i = 0; i < ARR_SIZE(hwkm_cipher_test_vectors); i++) {
        hwkm_cipher_test_vector_t test = hwkm_cipher_test_vectors[i];

        if (test.enabled) {
            if (test_IHWKeyFactory_ICipher(factory, false, &test) ==
                SAMPLEAPP_SUCCESS) {
                qsee_log(QSEE_LOG_MSG_DEBUG,
                         "Passed hwkm cipher decrypt test case %d", i);
                passed++;
            } else {
                qsee_log(QSEE_LOG_MSG_ERROR,
                         "Failed hwkm cipher decrypt test case %d", i);
                failed++;
            }
        } else {
            qsee_log(QSEE_LOG_MSG_ERROR,
                     "Skipped hwkm cipher encrypt test case %d", i);
            skipped++;
        }
    }

    // TODO: work with HW team to enable on Waipio
    // Tracked by CR: https://orbit/CR/2851448

    // for (size_t i = 0; i < ARR_SIZE(hwkm_hmac_test_vectors); i++) {
    //     hwkm_hmac_test_vector_t test = hwkm_hmac_test_vectors[i];

    //     if (test.enabled) {
    //         if (test_IHWKeyFactory_IHmac(factory, &test) == SAMPLEAPP_SUCCESS) {
    //             qsee_log(QSEE_LOG_MSG_DEBUG, "Passed hwkm hmac test case %d",
    //                      i);
    //             passed++;
    //         } else {
    //             qsee_log(QSEE_LOG_MSG_ERROR, "Failed hwkm hmac test case %d",
    //                      i);
    //             failed++;
    //         }
    //     } else {
    //         qsee_log(QSEE_LOG_MSG_ERROR, "Skipped hwkm hmac test case %d", i);
    //         skipped++;
    //     }
    // }

    for (size_t i = 0; i < ARR_SIZE(hwkm_cmac_test_vectors); i++) {
        hwkm_cmac_test_vector_t test = hwkm_cmac_test_vectors[i];

        if (test.enabled) {
            if (test_IHWKeyFactory_cmac(factory, &test) == SAMPLEAPP_SUCCESS) {
                qsee_log(QSEE_LOG_MSG_DEBUG, "Passed hwkm cmac test case %d",
                         i);
                passed++;
            } else {
                qsee_log(QSEE_LOG_MSG_ERROR, "Failed hwkm cmac test case %d",
                         i);
                failed++;
            }
        } else {
            qsee_log(QSEE_LOG_MSG_ERROR, "Skipped hwkm cmac test case %d", i);
            skipped++;
        }
    }

    qsee_log(QSEE_LOG_MSG_DEBUG,
             "hwkm crypto test results: passed %d, failed %d, skipped %d",
             passed, failed, skipped);

    return (failed == 0) ? SAMPLEAPP_SUCCESS : SAMPLEAPP_FAILED;
}

/**
 * Test that encrypt() and decrypt() methods of ICipher fail when invoked
 * with parameters which do not match the algorithm or allowed operations
 * specified by the key's policy
 */
int test_IHWKeyFactory_incorrect_policy(Object factory)
{
    int status = SAMPLEAPP_FAILED;
    Object key = Object_NULL;
    Object ctx = Object_NULL;
    uint8_t sw_key[32] = {0};
    uint8_t data[32] = {0};
    uint8_t out[sizeof(data)] = {0};
    size_t outlen = 0;
    IHWKey_policy policy = {
        .keyType = IHWKey_KEY_TYPE_OPERATIONAL,
        .destination = IHWKey_DESTINATION_DEFAULT,
        .opPolicy = {.algorithm = IHWKey_ALGORITHM_AES128_CBC,
                     .operations = IHWKey_OPERATION_ENCRYPT_ALLOWED}};

    ASSERT_OK_CLEAN(qsee_open(CCipherAES128_UID, &ctx), ctx);

    ASSERT_OK_CLEAN(
        IHWKeyFactory_importKey(factory, &policy, sw_key, sizeof(sw_key), &key),
        key);

    ASSERT_OK(ICipher_setParamAsObject(ctx, ICipher_PARAM_HW_KEY_OBJECT, key));

    // Selecting incorrect mode here, mismatched with policy
    ASSERT_OK(ICipher_setParamAsU32(ctx, ICipher_PARAM_MODE, ICipher_MODE_ECB));

    // When encrypt operation occurs, CE should encounter an error (algorithm
    // and mode combination requested does not match policy)
    ASSERT_ERROR(
        ICipher_encrypt(ctx, data, sizeof(data), out, sizeof(out), &outlen));

    ASSERT_OK(ICipher_reset(ctx));

    ASSERT_OK(ICipher_setParamAsObject(ctx, ICipher_PARAM_HW_KEY_OBJECT, key));
    ASSERT_OK(ICipher_setParamAsU32(ctx, ICipher_PARAM_MODE, ICipher_MODE_CBC));
    ASSERT_OK(ICipher_setParamAsData(ctx, ICipher_PARAM_IV, data, 16));

    // When decrypt operation occurs, CE should encounter an error (decrypt not
    // allowed in policy)
    ASSERT_ERROR(
        ICipher_decrypt(ctx, data, sizeof(data), out, sizeof(out), &outlen));

    status = SAMPLEAPP_SUCCESS;

__Exit:
    Object_ASSIGN_NULL(key);
    Object_ASSIGN_NULL(ctx);
    return status;
}

/**
 * This test ensures that when keys are derived from the same parent key
 * with different contexts, the derived keys have different values
 */
int test_IHWKeyFactory_deriveFromHWSource_different_ctx(Object factory)
{
    int status = SAMPLEAPP_FAILED;
    Object key1 = Object_NULL;
    Object key2 = Object_NULL;
    Object ctx = Object_NULL;
    uint8_t data[32] = {0};
    uint8_t out1[sizeof(data)] = {0};
    uint8_t out2[sizeof(data)] = {0};
    size_t outlen = 0;
    IHWKey_bindings bindings1 = {.context = {0xAA, 0xBB, 0xCC, 0xDD},
                                 .contextLen = 4};
    IHWKey_bindings bindings2 = {.context = {0xFF, 0xFF, 0xFF, 0xFF},
                                 .contextLen = 4};
    IHWKey_policy policy = {
        .keyType = IHWKey_KEY_TYPE_OPERATIONAL,
        .destination = IHWKey_DESTINATION_DEFAULT,
        .opPolicy = {.algorithm = IHWKey_ALGORITHM_AES128_ECB,
                     .operations = IHWKey_OPERATION_ENCRYPT_ALLOWED}};

    ASSERT_OK_CLEAN(qsee_open(CCipherAES128_UID, &ctx), ctx);

    ASSERT_OK_CLEAN(
        IHWKeyFactory_deriveFromHWSource(
            factory, &policy, &bindings1,
            IHWKeyFactory_HW_SOURCE_CHIP_RANDOM_KEY, Object_NULL, &key1),
        key1);

    ASSERT_OK_CLEAN(
        IHWKeyFactory_deriveFromHWSource(
            factory, &policy, &bindings2,
            IHWKeyFactory_HW_SOURCE_CHIP_RANDOM_KEY, Object_NULL, &key2),
        key2);

    ASSERT_OK(ICipher_setParamAsObject(ctx, ICipher_PARAM_HW_KEY_OBJECT, key1));
    ASSERT_OK(ICipher_setParamAsU32(ctx, ICipher_PARAM_MODE, ICipher_MODE_ECB));
    ASSERT_OK(
        ICipher_encrypt(ctx, data, sizeof(data), out1, sizeof(out1), &outlen));

    ASSERT_OK(ICipher_reset(ctx));

    ASSERT_OK(ICipher_setParamAsObject(ctx, ICipher_PARAM_HW_KEY_OBJECT, key2));
    ASSERT_OK(ICipher_setParamAsU32(ctx, ICipher_PARAM_MODE, ICipher_MODE_ECB));
    ASSERT_OK(
        ICipher_encrypt(ctx, data, sizeof(data), out2, sizeof(out2), &outlen));

    // Encrypted data should be different
    ASSERT_TRUE(timesafe_memcmp(out1, out2, sizeof(out1)) != 0);

    status = SAMPLEAPP_SUCCESS;

__Exit:
    Object_ASSIGN_NULL(key1);
    Object_ASSIGN_NULL(key2);
    Object_ASSIGN_NULL(ctx);
    return status;
}

/**
 * This test ensures that a key cannot be unwrapped by a wrapping key derived
 * with a different context than the wrapping key used to wrapp it.
 */
int test_IHWKeyFactory_cannot_unwrap_with_different_wrap_key(Object factory)
{
    int status = SAMPLEAPP_FAILED;
    Object key = Object_NULL;
    Object wrap_key1 = Object_NULL;
    Object wrap_key2 = Object_NULL;
    Object unwrapped_key = Object_NULL;
    uint8_t data[32] = {0};
    uint8_t* blob = NULL;
    uint32_t bloblen = 0;
    size_t outlen = 0;
    IHWKey_bindings bindings1 = {.context = {0xAA, 0xBB, 0xCC, 0xDD},
                                 .contextLen = 4};
    IHWKey_bindings bindings2 = {.context = {0xFF, 0xFF, 0xFF, 0xFF},
                                 .contextLen = 4};
    IHWKey_policy policy = {
        .keyType = IHWKey_KEY_TYPE_OPERATIONAL,
        .destination = IHWKey_DESTINATION_DEFAULT,
        .permissions = IHWKey_PERMISSION_WRAP_KEY_EXPORT,
        .opPolicy = {.algorithm = IHWKey_ALGORITHM_AES128_ECB,
                     .operations = IHWKey_OPERATION_ENCRYPT_ALLOWED}};
    IHWKey_policy wrapping_policy = {.keyType = IHWKey_KEY_TYPE_WRAPPING,
                                     .destination = IHWKey_DESTINATION_DEFAULT};

    ASSERT_OK_CLEAN(IHWKeyFactory_generate(factory, &policy, &key), key);

    ASSERT_OK_CLEAN(
        IHWKeyFactory_deriveFromHWSource(
            factory, &wrapping_policy, &bindings1,
            IHWKeyFactory_HW_SOURCE_CHIP_RANDOM_KEY, Object_NULL, &wrap_key1),
        wrap_key1);

    ASSERT_OK_CLEAN(
        IHWKeyFactory_deriveFromHWSource(
            factory, &wrapping_policy, &bindings2,
            IHWKeyFactory_HW_SOURCE_CHIP_RANDOM_KEY, Object_NULL, &wrap_key2),
        wrap_key2);

    ASSERT_OK(IHWKey_getBlobSize(key, &bloblen));

    blob = QSEE_ZALLOC_ARRAY(uint8_t, bloblen);
    ASSERT_TRUE(blob != NULL);

    ASSERT_OK(IHWKey_wrap(key, wrap_key1, blob, bloblen, &outlen));

    // Trying to unwrap with a wrapping key derived with different bindings
    // should fail because the wrapping keys have different values
    ASSERT_ERROR_CLEAN(
        IHWKeyFactory_unwrap(factory, wrap_key2, blob, bloblen, &unwrapped_key),
        unwrapped_key);

    status = SAMPLEAPP_SUCCESS;

__Exit:
    Object_ASSIGN_NULL(key);
    Object_ASSIGN_NULL(wrap_key1);
    Object_ASSIGN_NULL(wrap_key2);
    Object_ASSIGN_NULL(unwrapped_key);
    QSEE_FREE_PTR(blob);
    return status;
}

int run_hw_key_factory_tests(void)
{
    uint32_t passed = 0;
    uint32_t failed = 0;
    Object factory = Object_NULL;

    int status = qsee_open(CHWKeyFactory_UID, &factory);
    if (status == IOpener_ERROR_NOT_SUPPORTED) {
        // Do not run tests if no HWKM support on chipset
        return SAMPLEAPP_SUCCESS;
    }
    ASSERT_OK_CLEAN(status, factory);

    TEST(test_IHWKeyFactory_deriveFromHWSource(factory), passed, failed);
    TEST(test_IHWKeyFactory_generate(factory), passed, failed);
    TEST(test_IHWKeyFactory_importKey(factory), passed, failed);
    TEST(test_IHWKeyFactory_persist_usecase(factory), passed, failed);
    TEST(test_IHWKeyFactory_derive(factory), passed, failed);
    TEST(test_IHWKeyFactory_all_algorithms(factory), passed, failed);
    TEST(test_IHWKeyFactory_deriveFromHWSource_different_ctx(factory), passed,
         failed);
    TEST(test_IHWKeyFactory_cannot_unwrap_with_different_wrap_key(factory),
         passed, failed);
    TEST(test_IHWKeyFactory_incorrect_policy(factory), passed, failed);

    qsee_log(QSEE_LOG_MSG_DEBUG,
             "HW Key Factory test results: %d Passed, %d Failed", passed,
             failed);

__Exit:
    Object_ASSIGN_NULL(factory);
    return (Object_isOK(status) && (failed == 0)) ? SAMPLEAPP_SUCCESS
                                                  : SAMPLEAPP_FAILED;
}
