// SPDX-License-Identifier: BSD-3-Clause
/*
 * Copyright 2018-2020, 2022-2023 NXP
 */

/*===========================================================================*/
/**
    @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.
 */

/*===========================================================================
                                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
=============================================================================*/
static 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:
        case EVP_PKEY_RSA_PSS:
            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:
        case EVP_PKEY_RSA_PSS:
            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_RSA_PSS:
            sig_fmt = SIG_FMT_RSA_PSS;
            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;
}
