/*===========================================================================*/
/**
    @file    acst.c

    @brief   Code signing tool's AHAB file, handles AHAB data and creates
             output binary with the source binary, certificates and signatures
             generated while processing csf commands.

@verbatim
=============================================================================

    Copyright 2018-2020 NXP

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this
   list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.
3. Neither the name of the copyright holder nor the names of its contributors
   may be used to endorse or promote products derived from this software without
   specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

=============================================================================
@endverbatim */

/*===========================================================================
                                INCLUDE FILES
=============================================================================*/
#include <string.h>
#include <time.h>
#include <openssl/x509v3.h>
#include "err.h"
#include "srk_helper.h"
#include "misc_helper.h"

/*===========================================================================
                               LOCAL CONSTANTS
=============================================================================*/
#define MAX_ERR_MSG_BYTES (1024)
#define FILE_EXT_BIN      ".bin"

/*===========================================================================
                                 LOCAL MACROS
=============================================================================*/
#define ALIGN(val,mask)   (((val)+((mask)-1))&~((mask)-1))

/*===========================================================================
                  LOCAL TYPEDEFS (STRUCTURES, UNIONS, ENUMS)
=============================================================================*/

/*===========================================================================
                            LOCAL VARIABLES
=============================================================================*/
char err_msg[MAX_ERR_MSG_BYTES];

/*===========================================================================
                            LOCAL FUNCTION PROTOTYPES
=============================================================================*/
/** Retrieve hash algo
 *
 * Reads an SRK table to retrieve the hash algo information
 *
 * @param[in]  srk_table Byte string of a SRK table
 *
 * @pre @a srk_table must not be NULL
 *
 * @returns the hash algo information
 */
static hash_alg_t
get_hash_alg(byte_str_t *srk_table);

/** Get signature size
 *
 * Computes the expected signature size based on the key information
 *
 * @param[in]  filename Certificate file
 *
 * @pre @a filename must not be NULL
 *
 * @returns the size of the signature generated by using this key pair
 */
static size_t
get_signature_size(const char *filename);

/** Convert certificate to NXP format
 *
 * Converts a certificate to a Certificate struct as defined by NXP
 *
 * @param[in]  filename       Certificate file
 *
 * @param[out] cert_nxp       Certificate output in NXP format
 *
 * @param[in]  permissions    Permissions specified for this certificate
 *
 * @param[in]  hash           Hash algo to be used for signature generation
 *
 * @param[in]  signature_size Size of the certificate signature
 *
 * @pre @a filename and @a cert_nxp must not be NULL
 *
 * @rpost none
 */
static void
convert_to_nxp_format(const char *filename,
                      byte_str_t *cert_nxp,
                      uint8_t    permissions,
                      hash_alg_t hash,
                      size_t     signature_size);

/** Convert data to file
 *
 * Writes a byte string to an output file
 *
 * @param[in]  data     Input byte string
 *
 * @param[in]  filename Output file
 *
 * @pre @a data and @a filename must not be NULL
 *
 * @post none
 */
static void
convert_byte_str_to_file(byte_str_t *data, const char *filename);

/** Generate a signature
 *
 * Generates a signature over the input data
 *
 * @param[in]  data Input byte string to be signed
 *
 * @param[in]  key  Certificate used to sign
 *
 * @param[in]  hash Hash algo used to sign
 *
 * @param[out] sign Generated signature
 *
 * @pre @a data, @a key and @a sign must not be NULL
 *
 * @post none
 */
static void
generate_signature(byte_str_t *data,
                   const char *data_filename_prefix,
                   const char *key,
                   hash_alg_t hash,
                   byte_str_t *sign,
                   const char *sign_filename);

/** Generate the final output
 *
 * Generates the ouput file expected by the user
 *
 * @param[in]  src     Unsigned source file
 *
 * @param[in]  data    Signed container
 *
 * @param[in]  offsets Where to find the container in the source file
 *
 * @param[in]  dst     Signed destination file
 *
 * @pre @a src, @a data and @a offsets must not be NULL
 *
 * @post Program exits with exit code 0 on success.
 */
static void
generate_output(FILE       *dst_tmp,
                byte_str_t *data,
                offsets_t  *offsets,
                const char *dst);

/*===========================================================================
                            LOCAL FUNCTIONS
=============================================================================*/
/*--------------------------
  ahab_hash_2_cst_hash
---------------------------*/
hash_alg_t ahab_hash_2_cst_hash(uint8_t h)
{
    switch (h)
    {
        case HASH_SHA256:
            return SHA_256;
        case HASH_SHA384:
            return SHA_384;
        case HASH_SHA512:
            return SHA_512;
        default:
            return INVALID_DIGEST;
    }
}

/*--------------------------
  get_hash_alg
---------------------------*/
hash_alg_t get_hash_alg(byte_str_t *srk_table)
{
    uint32_t srk_tab_hdr_bytes = sizeof(struct ahab_container_srk_table_s);
    uint32_t srk_rec_hdr_bytes = sizeof(struct ahab_container_srk_s);

    /* Check overflow */
    if ((srk_tab_hdr_bytes + srk_rec_hdr_bytes) < srk_tab_hdr_bytes)
    {
        error("Wrong srk struct definitions");
    }

    /* Check if enough data */
    if (srk_table->entry_bytes < (srk_tab_hdr_bytes + srk_rec_hdr_bytes))
    {
        error("SRK table too small");
    }

    /* Get hash algo info from SRK record 0                   */
    /* For AHAB, all SRKs in a table must be of the same type */
    return ahab_hash_2_cst_hash(((struct ahab_container_srk_s *)(srk_table->entry + srk_tab_hdr_bytes))->hash);
}

/*--------------------------
  get_signature_size
---------------------------*/
size_t get_signature_size(const char *filename)
{
    X509     *cert = NULL;
    EVP_PKEY *pkey = NULL;
    size_t   size  = offsetof(struct ahab_container_signature_s, data);

    /* Open the certificate file */
    cert = read_certificate(filename);
    if (NULL == cert)
    {
        snprintf(err_msg, MAX_ERR_MSG_BYTES, "Unable to read %s", filename);
        error(err_msg);
    }

    /* Extract public key information */
    pkey = X509_get_pubkey(cert);
    if (NULL == pkey)
    {
        X509_free(cert);
        snprintf(err_msg,
                 MAX_ERR_MSG_BYTES,
                 "Unable to retrieve the public key from %s",
                 filename);
        error(err_msg);
    }

    /* Compute the size of the signature generated by using this key pair */
    switch (EVP_PKEY_id(pkey))
    {
        case EVP_PKEY_RSA:
            size += EVP_PKEY_size(pkey);
            break;

        case EVP_PKEY_EC:
            /* The EVP_PKEY_size() output includes the DER encoding */
            size += (EVP_PKEY_bits(pkey) >> 3) * 2;
            if (EVP_PKEY_bits(pkey) & 0x7) size += 2; /* Valid for P-521 */
            break;
    }

    EVP_PKEY_free(pkey);
    X509_free(cert);

    return size;
}

/*--------------------------
  convert_to_nxp_format
---------------------------*/
void convert_to_nxp_format(const char *filename,
                           byte_str_t *cert_nxp,
                           uint8_t    permissions,
                           hash_alg_t hash,
                           size_t     signature_size)
{
    X509         *cert   = NULL;
    EVP_PKEY     *pkey   = NULL;
    bool         ca_flag = false;
    const tgt_t  target  = TGT_AHAB;
    srk_entry_t  srk_entry;
    struct ahab_container_certificate_s *ahab_cert;
    uint32_t     cert_hdr_bytes = offsetof(struct ahab_container_certificate_s,
                                           public_key);

    /* Open the certificate file */
    cert = read_certificate(filename);
    if (NULL == cert)
    {
        snprintf(err_msg,
                    MAX_ERR_MSG_BYTES,
                    "Unable to read %s",
                    filename);
        error(err_msg);
    }

    /* Get the public key  */
    pkey = X509_get_pubkey(cert);
    if (NULL == pkey)
    {
        X509_free(cert);
        error("Unable to retrieve the \'Certificate\' public key");
    }

    /* Determine if the certificate has the CA flag set */
    if (X509_CA_CERT == X509_check_ca(cert))
    {
        ca_flag = TRUE;
    }

    /* Build the SRK record */
    switch (EVP_PKEY_id(pkey))
    {
        case EVP_PKEY_RSA:
            srk_entry_pkcs1(target, pkey, &srk_entry, ca_flag, get_digest_name(hash));
            break;

        case EVP_PKEY_EC:
            srk_entry_ec(target, pkey, &srk_entry, ca_flag, get_digest_name(hash));
            break;

        default:
            EVP_PKEY_free(pkey);
            X509_free(cert);
            error("Unsupported type of public key, must be either RSA or EC");
    }

    EVP_PKEY_free(pkey);

    /* Build the complete certificate */
    /* The Certificate signature must start on 64-bit boundaries */
    cert_nxp->entry_bytes = ALIGN(cert_hdr_bytes + srk_entry.entry_bytes, 8);
    cert_nxp->entry       = malloc(cert_nxp->entry_bytes);
    if (NULL == cert_nxp->entry)
    {
        X509_free(cert);
        error("Cannot allocate buffer fro handling the \'Certificate\'");
    }
    memset(cert_nxp->entry, 0, cert_nxp->entry_bytes);

    /* Check overflow */
    if ((signature_size + cert_nxp->entry_bytes) < signature_size)
    {
        X509_free(cert);
        error("Wrong \'Certificate\' or \'Signature\' definitions");
    }

    ahab_cert = (struct ahab_container_certificate_s *)(cert_nxp->entry);
    ahab_cert->tag              = CERTIFICATE_TAG;
    ahab_cert->version          = 0;
    ahab_cert->length           = cert_nxp->entry_bytes + signature_size;
    ahab_cert->signature_offset = cert_nxp->entry_bytes;

    cert_nxp->entry[6] = ~permissions;
    cert_nxp->entry[7] = permissions;

    memcpy(cert_nxp->entry + cert_hdr_bytes,
           srk_entry.entry,
           srk_entry.entry_bytes);

    free(srk_entry.entry);
}

/*--------------------------
  convert_byte_str_to_file
---------------------------*/
void convert_byte_str_to_file(byte_str_t *data, const char *filename)
{
    FILE *file = NULL;

    /* Open the destination file */
    if ((file = fopen(filename, "wb")) == NULL)
    {
        snprintf(err_msg,
                 MAX_ERR_MSG_BYTES,
                 "Unable to create binary file %s",
                 filename);
        error(err_msg);
    }

    /* Write the data */
    if (data->entry_bytes != fwrite(data->entry, 1, data->entry_bytes, file))
    {
        snprintf(err_msg,
                 MAX_ERR_MSG_BYTES,
                 "Unable to write to binary file %s",
                 filename);
        error(err_msg);
    }

    fclose (file);
}

/*--------------------------
  generate_signature
---------------------------*/
void generate_signature(byte_str_t *data,
                        const char *data_filename_prefix,
                        const char *key,
                        hash_alg_t hash,
                        byte_str_t *sig,
                        const char *sig_filename)
{
    uint32_t   sig_hdr_bytes  = offsetof(struct ahab_container_signature_s, data);
    X509       *skey          = NULL;
    EVP_PKEY   *pkey          = NULL;
    sig_fmt_t  sig_fmt        = SIG_FMT_UNDEF;
    size_t     sig_bytes      = 0;
    byte_str_t sig_data       = {NULL, 0};
    uint8_t    *sha1          = NULL;
    size_t     digest_size    = 0;
    char       *data_filename = NULL;
    time_t     t;
    struct tm  *t_info;
    char       t_str[4+2+2+2+2+2]; /* YYYY + MM + DD + HH + MM + "-\0" */
    struct ahab_container_signature_s *ahab_sig;

    /* Get signing key information */
    skey = read_certificate(key);
    if (NULL == skey)
    {
        snprintf(err_msg, MAX_ERR_MSG_BYTES, "Unable to read %s", key);
        error(err_msg);
    }

    /* Get the public key  */
    pkey = X509_get_pubkey(skey);
    if (NULL == pkey)
    {
        X509_free(skey);
        snprintf(err_msg,
                 MAX_ERR_MSG_BYTES,
                 "Unable to retrieve the public key from %s",
                 key);
        error(err_msg);
    }

    /* Build the signature struct */
    sig->entry_bytes = get_signature_size(key);
    sig_bytes        = sig->entry_bytes - sig_hdr_bytes;
    sig->entry       = malloc(sig->entry_bytes);
    if (NULL == sig->entry)
    {
        error("Cannot allocate buffer for creating a signature");
    }
    memset(sig->entry, 0, sig->entry_bytes);

    ahab_sig = (struct ahab_container_signature_s *)(sig->entry);
    ahab_sig->tag     = SIGNATURE_TAG;
    ahab_sig->version = 0;
    ahab_sig->length  = sig->entry_bytes;

    switch (EVP_PKEY_id(pkey))
    {
        case EVP_PKEY_RSA:
            sig_fmt = SIG_FMT_PKCS1;
            break;

        case EVP_PKEY_EC:
            sig_fmt = SIG_FMT_ECDSA;
            break;
    }

    EVP_PKEY_free(pkey);
    X509_free(skey);

    /* Generate data filename */
    data_filename = malloc(strlen(data_filename_prefix)
                           + sizeof(t_str) - 1
                           + HASH_BYTES_SHA1 * 2
                           + sizeof(FILE_EXT_BIN));
    if (NULL == data_filename)
    {
        error("Unable to allocate memory for data filename");
    }
    strcpy(data_filename, data_filename_prefix);

    /* Add time information to filename */
    time(&t);
    t_info = localtime(&t);
    strftime(t_str, sizeof(t_str), "%Y%m%d%H%M-", t_info);
    strcat(data_filename, t_str);

    /* Add sha1 information to filename */
    sha1 = generate_hash(data->entry,
                         data->entry_bytes,
                         get_digest_name(SHA_1),
                         &digest_size);
    if (NULL == sha1)
    {
        error("Fail to generate the sha1 digest to be used as filename");
    }
    while (digest_size)
    {
        char tmp[4];
        snprintf(tmp, 4, "%02x", sha1[HASH_BYTES_SHA1 - digest_size--]);
        strcat(data_filename, tmp);
    }
    free(sha1);

    /* Add extension to filename */
    memcpy(data_filename
           + strlen(data_filename_prefix)
           + sizeof(t_str) - 1
           + HASH_BYTES_SHA1 * 2,
           FILE_EXT_BIN,
           sizeof(FILE_EXT_BIN));

    /* Create a tmp file for the signing process */
    convert_byte_str_to_file(data, data_filename);

    if (NULL == sig_filename)
    {
        /* Generate the signature */
        if (SUCCESS != gen_sig_data(data_filename,
                                    key,
                                    hash,
                                    sig_fmt,
                                    sig->entry + sig_hdr_bytes,
                                    &sig_bytes,
                                    g_mode))
        {
            error("Unable to generate the signature");
        }
    }
    else
    {
        read_file(sig_filename, &sig_data, NULL);

        if (sig_bytes != sig_data.entry_bytes)
        {
            snprintf(err_msg, MAX_ERR_MSG_BYTES,
                     "The length of the signature file %s is not valid", sig_filename);
            error(err_msg);
        }

        if (CAL_SUCCESS != ver_sig_data(data_filename,
                                        key,
                                        hash,
                                        sig_fmt,
                                        sig_data.entry,
                                        sig_data.entry_bytes))
        {
            snprintf(err_msg, MAX_ERR_MSG_BYTES,
                     "The signature file %s is not valid", sig_filename);
            error(err_msg);
        }

        memcpy(sig->entry + sig_hdr_bytes, sig_data.entry, sig_data.entry_bytes);

        free(sig_data.entry);
    }

    if (MODE_HSM != g_mode)
    {
        if (0 != remove(data_filename))
        {
            snprintf(err_msg, MAX_ERR_MSG_BYTES, "Unable to delete %s", data_filename);
            error(err_msg);
        }
    }

    if (sig->entry_bytes != (sig_hdr_bytes + sig_bytes))
    {
        error("Unexpected signature length");
    }
}

/*--------------------------
  generate_output
---------------------------*/
void generate_output(FILE        *dst_tmp,
                     byte_str_t  *data,
                     offsets_t   *offsets,
                     const char  *dst)
{
    FILE     *file_dst = NULL;
    char     chr;
    uint32_t count, read_size;

    /* Create destination file */
    if ((file_dst = fopen(dst, "wb")) == NULL)
    {
        snprintf(err_msg,
                 MAX_ERR_MSG_BYTES,
                 "Unable to create binary file %s",
                 dst);
        error(err_msg);
    }

    /* Fill destination file with source data until first offset */
    fseek(dst_tmp, 0, SEEK_SET);
    count = 0;
    while (offsets->first != count)
    {
        if (1 != fread(&chr, 1, 1, dst_tmp))
        {
            snprintf(err_msg, MAX_ERR_MSG_BYTES, "Unexpected read termination");
            error(err_msg);
        }

        if (1 != fwrite(&chr, 1, 1, file_dst))
        {
            snprintf(err_msg,
                    MAX_ERR_MSG_BYTES,
                    "Unable to write to binary file %s",
                    dst);
            error(err_msg);
        }

        count++;
    }

    /* Fill destination file with input data */
    if (data->entry_bytes != fwrite(data->entry, 1, data->entry_bytes, file_dst))
    {
        snprintf(err_msg,
                 MAX_ERR_MSG_BYTES,
                 "Unable to write to binary file %s",
                 dst);
        error(err_msg);
    }

    /* Fill destination file with remaining source data */
    fseek(dst_tmp, data->entry_bytes, SEEK_CUR);
    do
    {
        read_size = fread(&chr, 1, 1, dst_tmp);

        if (1 < read_size)
        {
            snprintf(err_msg, MAX_ERR_MSG_BYTES, "Unexpected read");
            error(err_msg);
        }

        if (read_size != fwrite(&chr, 1, read_size, file_dst))
        {
            snprintf(err_msg,
                    MAX_ERR_MSG_BYTES,
                    "Unable to write to binary file %s",
                    dst);
            error(err_msg);
        }
    } while (0 != read_size);

    fclose(dst_tmp);
    fclose(file_dst);

    printf("CSF Processed successfully and signed image available in %s\n", dst);
    exit(0);
}

/*--------------------------
  encrypt_images
---------------------------*/
void encrypt_images(ahab_data_t *ahab_data,
                    byte_str_t *cont_hdr,
                    const char *key,
                    uint8_t key_length,
                    FILE *destination)
{
    struct ahab_container_header_s *container_header
        = (struct ahab_container_header_s *)cont_hdr->entry;

    if (container_header->nrImages > AHAB_MAX_NR_IMAGES) {
        error("Invalid number of images");
    }

    if (!(((1U << container_header->nrImages) - 1) & ahab_data->image_indexes)) {
        error("At least one image must be encrypted");
    }

    /** Encrypt the images **/

    size_t digest_size;

    struct ahab_container_image_s *image = get_ahab_image_array(container_header);

    for (uint8_t i = 0; i < container_header->nrImages; i++) {

        if (0 == image->image_size || 0 == (ahab_data->image_indexes & (1U << i))) {
            image++;
            continue;
        }

        uint8_t hash_type = ahab_container_image_get_hash(image);
        int32_t hash_size = ahab_get_hash_size_by_sha_type(hash_type);

        if (hash_size < 0) {
            error("Unsupported hash algorithm for image integrity");
        }

        uint8_t *hash = malloc(hash_size);

        if (NULL ==  hash) {
            error("Cannot allocate memory for the hash value of the plaintext");
        }

        /* Retrieve image data */

        offsets_t offsets = {
            .first  = ahab_data->offsets.first + image->image_offset,
            .second = offsets.first + image->image_size
        };

        byte_str_t image_data;

        read_file(ahab_data->source, &image_data, &offsets);

        if (NULL == image_data.entry) {
            error("Cannot allocate memory for the image index %d", i);
        }

        /* Create a temporary file with the data to be enrypted */
        FILE *plaintext = fopen(FILE_PLAIN_DATA, "wb");
        if (NULL == plaintext) {
            error("Cannot create temporary file for plaintext");
        }
        if (image_data.entry_bytes != fwrite(image_data.entry, 1, image_data.entry_bytes, plaintext)) {
            error("Fail to update the plaintext temporary file");
        }
        fclose(plaintext);

        /* Encrypt image data */

        size_t iv_length = 16; /* The IV size for AES-CBC is 128 bits */
        uint8_t *iv = malloc(SHA256_DIGEST_LENGTH);

        if (NULL == iv) {
            error("Cannot allocate memory for the IV");
        }

        iv = generate_hash(image_data.entry,
                        image_data.entry_bytes,
                        get_digest_name(SHA_256),
                        &digest_size);

        if (SHA256_DIGEST_LENGTH != digest_size) {
            error("Fail to generate IV of image index %d", i);
        }

        int32_t err_value = gen_auth_encrypted_data(
            FILE_PLAIN_DATA,
            FILE_ENCRYPTED_DATA,
            AES_CBC,
            NULL,
            0,
            iv + SHA256_DIGEST_LENGTH - iv_length,
            iv_length,
            NULL,
            0,
            key_length,
            g_cert_dek,
            key,
            g_reuse_dek);

        if (0 != remove(FILE_PLAIN_DATA)) {
            error("Cannot delete %s", FILE_PLAIN_DATA);
        }

        if (CAL_SUCCESS != err_value) {
            error("Fail to generate encrypted data for image index %d", i);
        }

        /* Replace image data by encrypted data */

        byte_str_t enc_data;

        read_file(FILE_ENCRYPTED_DATA, &enc_data, NULL);

        if (NULL == enc_data.entry) {
            error("Cannot allocate memory for the encrypted data of the image index %d", i);
        }

        if (0 != remove(FILE_ENCRYPTED_DATA)) {
            error("Cannot delete %s", FILE_ENCRYPTED_DATA);
        }

        if (image->image_size != enc_data.entry_bytes) {
            error("Unexpected size of encryted data of image index %d (%d)", i, enc_data.entry_bytes);
        }

        hash = generate_hash(enc_data.entry,
                            enc_data.entry_bytes,
                            get_digest_name(ahab_hash_2_cst_hash(hash_type)),
                            &digest_size);

        if (digest_size != (size_t)hash_size) {
            error("Fail to generate hash of encrypted data of image index %d", i);
        }

        memcpy(image->iv, iv, SHA256_DIGEST_LENGTH);
        memcpy(image->hash, hash, hash_size);
        memset(image->hash + hash_size, 0, SHA512_DIGEST_LENGTH - hash_size);
        image->flags |= IMAGE_FLAGS_ENCRYPTED;

        fseek(destination, offsets.first, SEEK_SET);

        if (enc_data.entry_bytes != fwrite(enc_data.entry, 1, enc_data.entry_bytes, destination)) {
            error("Fail to replace image index %d with encrypted data", i);
        }

        free(enc_data.entry);
        free(image_data.entry);
        free(hash);
        free(iv);

        image++;
    }
}

/*===========================================================================
                               GLOBAL FUNCTIONS
=============================================================================*/

/*--------------------------
  convert_certificate
---------------------------*/
int32_t handle_ahab_signature(void)
{
    ahab_data_t *ahab_data = &g_ahab_data;
    byte_str_t  srk_table  = {NULL, 0};
    byte_str_t  cert       = {NULL, 0};
    byte_str_t  cert_sign  = {NULL, 0};
    byte_str_t  cont_hdr   = {NULL, 0};
    byte_str_t  sign_blk   = {NULL, 0};
    struct ahab_container_signature_block_s *ahab_sig_blk;
    byte_str_t  container  = {NULL, 0};
    byte_str_t  cont_unsig = {NULL, 0};
    byte_str_t  signature  = {NULL, 0};
    byte_str_t  blob       = {NULL, 0};
    hash_alg_t  hash       = INVALID_DIGEST;
    const char  *sign_key  = NULL;

    /* Get SRK table */
    if (NULL == ahab_data->srk_table)
    {
        error("A SRK table must be installed");
    }
    read_file(ahab_data->srk_table, &srk_table, NULL);

    /* Parse SRK table to extract hash algo information */
    hash = get_hash_alg(&srk_table);
    if (INVALID_DIGEST == hash)
    {
        error("Invalid hash algo defined in SRK table");
    }

    /* Handle a Certificate if present */
    if (NULL != ahab_data->certificate)
    {
        convert_to_nxp_format(ahab_data->certificate,
                              &cert,
                              ahab_data->permissions,
                              hash,
                              get_signature_size(ahab_data->srk_entry));

        generate_signature(&cert,
                           "certificate-",
                           ahab_data->srk_entry,
                           hash,
                           &cert_sign,
                           ahab_data->cert_sign);

        sign_key = ahab_data->certificate;
    }
    else
    {
        sign_key = ahab_data->srk_entry;
    }

    /* Get Container header */
    read_file(ahab_data->source, &cont_hdr, &ahab_data->offsets);

    /* Check if signature block offset matches header information */
    if ((ahab_data->offsets.second - ahab_data->offsets.first) !=
        ((struct ahab_container_header_s *)(cont_hdr.entry))->signature_block_offset)
    {
        error("Offsets are not consistent with the input binary to be signed");
    }

    /* Copy source to temporary file */
    FILE *source = fopen(ahab_data->source, "rb");
    if (NULL == source ) {
        error("Cannot open %s", ahab_data->source);
    }
    FILE *destination = tmpfile();
    if (NULL == destination) {
        error("Cannot create temporary file");
    }

    char byte;
    while (0 != fread(&byte, 1, 1, source)) {
        if (1 != fwrite(&byte, 1, 1, destination)) {
            error("Fail to fill the temporary file");
        }
    }

    /* Handle a DEK if requested */
    if (NULL != ahab_data->dek)
    {
        /* Encrypt the images */
        encrypt_images(ahab_data, &cont_hdr, ahab_data->dek, ahab_data->dek_length, destination);

        blob.entry_bytes = AHAB_BLOB_HEADER + CAAM_BLOB_OVERHEAD_NORMAL + ahab_data->dek_length;
        blob.entry = malloc(blob.entry_bytes);
        if (NULL == blob.entry)
        {
            error("Cannot allocate memory for the Blob placeholder");
        }
        memset(blob.entry, 0xCA, blob.entry_bytes);
    }

    /* Create the Signature Block */
    sign_blk.entry_bytes = sizeof(struct ahab_container_signature_block_s)
                           + ALIGN(srk_table.entry_bytes, 8)
                           + ALIGN(get_signature_size(sign_key), 8)
                           + ALIGN(cert.entry_bytes + cert_sign.entry_bytes, 8)
                           + blob.entry_bytes;
    sign_blk.entry = malloc(sign_blk.entry_bytes);
    if (NULL == sign_blk.entry)
    {
        error("Cannot allocate memory for the Signature block");
    }
    memset(sign_blk.entry, 0, sign_blk.entry_bytes);

    /* Update the Signature block */
    ahab_sig_blk = (struct ahab_container_signature_block_s *)(sign_blk.entry);
    ahab_sig_blk->tag                = SIGNATURE_BLOCK_TAG;
    ahab_sig_blk->version            = 0;
    ahab_sig_blk->length             = sign_blk.entry_bytes;
    ahab_sig_blk->srk_table_offset   = sizeof(struct ahab_container_signature_block_s);
    ahab_sig_blk->signature_offset   = ahab_sig_blk->srk_table_offset
                                     + ALIGN(srk_table.entry_bytes, 8);
    ahab_sig_blk->certificate_offset = ahab_sig_blk->signature_offset
                                     + ALIGN(get_signature_size(sign_key), 8);
    ahab_sig_blk->blob_offset        = ahab_sig_blk->certificate_offset
                                     + ALIGN(cert.entry_bytes + cert_sign.entry_bytes, 8);
    ahab_sig_blk->key_identifier     = ahab_data->key_identifier;

    /* Copy SRK table data */
    memcpy(sign_blk.entry + ahab_sig_blk->srk_table_offset,
           srk_table.entry,
           srk_table.entry_bytes);

    free(srk_table.entry);

    /* Copy Certificate data */
    if (NULL == cert.entry)
    {
        ahab_sig_blk->certificate_offset = 0;
    }
    else
    {
        memcpy(sign_blk.entry + ahab_sig_blk->certificate_offset,
               cert.entry,
               cert.entry_bytes);
        memcpy(sign_blk.entry
               + ahab_sig_blk->certificate_offset
               + cert.entry_bytes,
               cert_sign.entry,
               cert_sign.entry_bytes);

        free(cert.entry);
        free(cert_sign.entry);
    }

    /* Copy Blob data */
    if (NULL == blob.entry)
    {
        ahab_sig_blk->blob_offset = 0;
        ahab_sig_blk->key_identifier = 0;
    }
    else
    {
        memcpy(sign_blk.entry + ahab_sig_blk->blob_offset,
               blob.entry,
               blob.entry_bytes);

        free(blob.entry);

        printf("The DEK BLOB must be inserted at offset 0x%zx (its expected size is %zd bytes)\n",
            ahab_data->offsets.first
            + cont_hdr.entry_bytes
            + ahab_sig_blk->blob_offset,
            blob.entry_bytes);
    }

    /* Update Container header length */
    container.entry_bytes = cont_hdr.entry_bytes + sign_blk.entry_bytes;
    ((struct ahab_container_header_s *)(cont_hdr.entry))->length =
        container.entry_bytes;

    /* Update Container header flags */
    ((struct ahab_container_header_s *)(cont_hdr.entry))->flags =
        ahab_data->srk_set     << HEADER_FLAGS_SRK_SET_SHIFT |
        ahab_data->srk_index   << HEADER_FLAGS_SRK_SHIFT     |
        ahab_data->revocations << HEADER_FLAGS_REVOKING_MASK_SHIFT;

    /* Build the signed Container */
    container.entry = malloc(container.entry_bytes);
    if (NULL == container.entry)
    {
        error("Cannot allocate memory for the Container");
    }
    memset(container.entry, 0, container.entry_bytes);

    memcpy(container.entry,
           cont_hdr.entry,
           cont_hdr.entry_bytes);

    memcpy(container.entry + cont_hdr.entry_bytes,
           sign_blk.entry,
           sign_blk.entry_bytes);

    /* Generate Container signature */
    cont_unsig.entry       = container.entry;
    cont_unsig.entry_bytes = cont_hdr.entry_bytes
                           + ahab_sig_blk->signature_offset;

    free(sign_blk.entry);

    generate_signature(&cont_unsig,
                       "container-",
                       sign_key,
                       hash,
                       &signature,
                       ahab_data->signature);

    memcpy(container.entry + cont_unsig.entry_bytes,
           signature.entry,
           signature.entry_bytes);

    free(signature.entry);

    /* Generate the output file */
    generate_output(destination,
                    &container,
                    &ahab_data->offsets,
                    ahab_data->destination);

    free(container.entry);

    return SUCCESS;
}
