#ifndef DK_SCP_COMMON_H
#define DK_SCP_COMMON_H

#include "dk_common.h"
#include "dk_constants.h"

extern byte LABEL_DERIVATION_SMAC[KDF_LABEL_SIZE];
extern byte LABEL_DERIVATION_SENC[KDF_LABEL_SIZE];
extern byte LABEL_DERIVATION_SRMAC[KDF_LABEL_SIZE];
extern byte LABEL_DERIVATION_CARD_CRYPTOGRAM[KDF_LABEL_SIZE];
extern byte LABEL_DERIVATION_HOST_CRYPTOGRAM[KDF_LABEL_SIZE];

/**
 * Context of an open SCP session. Holds keys and session keys used
 * to encrypt/decrypt messages and calculate MACs. This context structure
 * contains common elements between SCP03 and SCP11. 
 * 
 * Encryption and session key lengths are not know beforehand, and are 
 * represented as a byte pointer that must be allocated with the desired
 * key length. All of these must have the same size, which must be 
 * one of 128/192/256 bits.
 * 
 * The counter is used for the MAC chaining algorithm, being incrementaded 
 * after sending a command and getting a response from the eSE. The security
 * level determines whether the message should be encrypted/decrypted and 
 * authenticated.
 * 
 * The chaining buffer contains the current chaining data, used for the
 * MAC calculation. Host and card are randomluy. The cryptograms are generated
 * for a given session and used through the transactions.
 */
typedef enum scp_context_state {
    SESSION_CLOSED = 1 << 0,
    SESSION_OPEN = 1 << 1
} scp_context_state_t;

typedef struct scp_context_t {
    byte*       key_enc;    
    size_t      key_enc_len;
    byte*       key_mac;
    size_t      key_mac_len;

    byte*       senc;
    size_t      senc_len;
    byte*       smac;
    size_t      smac_len;
    byte*       srmac;
    size_t      srmac_len;

    uint32_t    encryption_counter;
    byte        security_level;

    byte        key_information[KEY_INFO_SIZE];
    byte        byteI;

    byte        chaining[MAC_CHAINING_SIZE];

    byte        host_challenge[CHALLENGE_SIZE];
    byte        card_challenge[CHALLENGE_SIZE];

    byte        host_cryptogram[CRYPTOGRAM_SIZE];
    byte        card_cryptogram[CRYPTOGRAM_SIZE];
    
    scp_context_state_t session_state;
} __attribute__((packed)) scp_context; 

/**
 * Initializes a SCP context to an initial state.
 * 
 * @param ctx The SCP context structure
 * @return The error code for the operation
 */
DK_Result scp_init(scp_context *ctx);

/**
 * Resets a SCP context structure to it's initial state,
 * freeing all internal memory and zero'ing all buffers.
 * 
 * @param ctx The SCP context structure
 * @return The error code for the operation
 */
DK_Result scp_clear(scp_context *ctx);

/**
 * Calculates response message authentication code given a SCP
 * context and APDU response data field and status word.
 * 
 * @param ctx The SCP context structure
 * @param rdf The APDU response data field
 * @param rdf_len The length of the APDU response data field
 * @param sw The APDU status word
 * @param sw_len The length of the APDU status word
 * @param rmac Buffer that will contain the calculated R-MAC of size R_MAC_SIZE
 * @return The error code for the operation
 * @see R_MAC_SIZE
 */
DK_Result scp_calculate_rmac(
    scp_context* ctx, 
    byte* rdf, 
    size_t rdf_len, 
    byte* sw, 
    size_t sw_len, 
    byte* rmac);

/**
 * Generates session keys for a SCP context. Requires that key_enc
 * and key_mac are already set and are not null. The session keys will
 * always have the same length as the key_enc and key_mac.
 * 
 * @param ctx The SCP context structure
 * @return The error code for the operation
 */
DK_Result scp_generate_session_keys(scp_context* ctx);

/**
 * Generates a pseudorandom host challenge.
 * 
 * @param ctx The SCP context structure
 * @return The error code for the operation
 */
DK_Result scp_generate_card_challenge(scp_context *ctx);

/**
 * Generates a card cryptogram for a SCP context. Requires that session
 * keys are already generated and that the card and host challenges are
 * set and not null.
 * 
 * @param ctx The SCP context structure
 * @param cryptogram A pre-allocated buffer of size CRYPTOGRAM_SIZE that will contain the generated cryptogram
 * @return The error code for the operation
 * @see CRYPTOGRAM_SIZE
 */
DK_Result scp_generate_card_cryptogram(
    scp_context* ctx, 
    byte* cryptogram);

DK_Result scp_verify_card_cryptogram(
    scp_context *ctx,
    byte *cryptogram);

/**
 * Generates a pseudorandom host challenge.
 * 
 * @param ctx The SCP context structure
 * @return The error code for the operation
 */
DK_Result scp_generate_host_challenge(scp_context *ctx);

/**
 * Generates a host cryptogram for a SCP context. Requires that session
 * keys are already generated and that the card and host challenges are
 * set and not null.
 * 
 * @param ctx The SCP context structure
 * @param cryptogram A pre-allocated buffer of size CRYPTOGRAM_SIZE that will contain the generated cryptogram
 * @return The error code for the operation
 * @see CRYPTOGRAM_SIZE
 */
DK_Result scp_generate_host_cryptogram(
    scp_context* ctx, 
    byte* cryptogram);

/**
 * Adds zeros to the left of the data until it's size is a multiple of
 * AES_BLOCK_SIZE. Allocates a new buffer for response, it is the
 * responsability of the caller to free this memory.
 * 
 * @param data The data buffer to be paddded
 * @param data_len The length of the data buffer
 * @param out The buffer containing the padded data. This is allocated in the heap, and should be freed when not necessary anymore.
 * @param out_len The length of the output buffer
 * @return The error code for the operation
 */
DK_Result scp_left_pad(
    byte* data, 
    size_t data_len, 
    byte** out, 
    size_t* out_len);

/**
 * Pads a data buffer according to NIST 800-38a, Appendix A. If the data length
 * is not a multiple of AES_BLOCK_SIZE, adds a delimiter (0x80) to the end of
 * the data and completes the size with 0x00 until it is a multiple of 
 * AES_BLOCK_SIZE.
 * 
 * @param data The data buffer to be paddded
 * @param data_len The length of the data buffer
 * @param out The buffer containing the stripped data. This must be preallocated with at least data_len + AES_BLOCK_SIZE bytes.
 * @param out_len The length of the output buffer
 * @return The error code for the operation
 */
DK_Result scp_aes_pad(
    byte* data, 
    size_t data_len, 
    byte* out, 
    size_t* out_len);

/**
 * Strips the padding of a data buffer according to NIST 800-38a, Appendix A. 
 * 
 * @param data The data buffer to be stripped
 * @param data_len The length of the data buffer
 * @param out The buffer containing the stripped data. This must be preallocated with at least data_len - AES_BLOCK_SIZE bytes.
 * @param out_len The length of the output buffer
 * @return The error code for the operation
 */
DK_Result scp_strip_aes_pad(
    byte *data, 
    size_t data_len, 
    byte* out, 
    size_t* out_len);

/**
 * Generates an initialization vector to be used in the AES encryption/decryption
 * of an OCE command.
 * 
 * @param ctx The SCP context structure
 * @param iv The generated IV buffer, which must have a length of AES_BLOCK_SIZE bytes
 * @return The error code for the operation
 */
DK_Result scp_generate_command_icv(
    scp_context* ctx, 
    byte* iv);

/**
 * Generates an initialization vector to be used in the AES encryption/decryption
 * of an OCE reponse.
 * 
 * @param ctx The SCP context structure
 * @param iv The generated IV buffer, which must have a length of AES_BLOCK_SIZE bytes
 * @return The error code for the operation
 */
DK_Result scp_generate_response_icv(
    scp_context* ctx, 
    byte* iv);

/**
 * Generates an initialization vector to be used in the AES encryption/decryption
 * of an OCE command or response.
 * 
 * @param ctx The SCP context structure
 * @param iv The generated IV buffer, which must have a length of AES_BLOCK_SIZE bytes
 * @param is_response If 0, the IV is for a command, else it's for a response.
 * @return The error code for the operation
 */
DK_Result scp_generate_icv(
    scp_context* ctx, 
    byte* iv,
    byte is_response);

/**
 * Wraps an APDU command by encrypting and calculating a MAC for the 
 * encrypted data.
 * 
 * @param ctx The SCP context structure
 * @param apdu A buffer contaning the APDU
 * @param apdu_len The length of the APDU buffer. Must be at least 6 bytes long
 * @param wrapped_apdu A buffer that will hold the wrapped APDU. This must be preallocated with at least CPDU_MAX_SIZE bytes
 * @param wrapped_apdu_len A pointer to an size_t variable that will hold the wrapped APDU length
 * @return The error code for the operation
 */
DK_Result scp_wrap_command(
    scp_context *ctx, 
    byte *apdu, 
    size_t apdu_len, 
    byte *wrapped_apdu, 
    size_t *wrapped_apdu_len);

/**
 * Encrypts the data field of the command APDU. The size will always
 * change due to added padding, and the relevant fields of the APDU
 * are updated to reflect that.
 * 
 * @param ctx The SCP context structure
 * @param apdu A buffer contaning the APDU
 * @param apdu_len The length of the APDU buffer. Must be at least 6 bytes long
 * @param wrapped_apdu A buffer that will hold the wrapped APDU. This must be preallocated with at least CPDU_MAX_SIZE bytes
 * @param wrapped_apdu_len A pointer to an size_t variable that will hold the wrapped APDU length
 * @return The error code for the operation
 */
DK_Result scp_encrypt_command(
    scp_context *ctx, 
    byte *apdu, 
    size_t apdu_len, 
    byte *wrapped_apdu, 
    size_t *wrapped_apdu_len);

/**
 * Generates a MAC for the command APDU (C-MAC) and appends it to the 
 * data field of the APDU, updating relevant parameters accordingly.
 * 
 * @param ctx The SCP context structure
 * @param apdu A buffer contaning the APDU
 * @param apdu_len The length of the APDU buffer. Must be at least 6 bytes long
 * @param wrapped_apdu A buffer that will hold the wrapped APDU. This must be preallocated with at least CPDU_MAX_SIZE bytes
 * @param wrapped_apdu_len A pointer to an size_t variable that will hold the wrapped APDU length
 * @return The error code for the operation
 */
DK_Result scp_generate_apdu_cmac(
    scp_context *ctx, 
    byte *apdu, 
    size_t apdu_len, 
    byte *wrapped_apdu, 
    size_t *wrapped_apdu_len);

/**
 * Unwraps a response APDU by checking its R-MAC and decrypting and
 * stripping the padding from the data. 
 * 
 * @param ctx The SCP context structure
 * @param wrapped_apdu A buffer contaning the wrapped APDU
 * @param wrapped_apdu_len The length of the wrapped APDU buffer.
 * @param apdu A buffer that will hold the unwrapped APDU. This must be preallocated with at least RPDU_MAX_SIZE bytes
 * @param apdu_len A pointer to an size_t variable that will hold the unwrapped APDU length
 * @return The error code for the operation
 */
DK_Result scp_unwrap_response(
    scp_context *ctx, 
    byte *wrapped_apdu, 
    size_t wrapped_apdu_len, 
    byte *apdu, 
    size_t* apdu_len);

/**
 * Verifies if the response APDU MAC (R-MAC), returning DK_SUCCESS
 * if it matches the received value or a TEE_GENERIC_ERROR otherwise.
 * 
 * @param ctx The SCP context structure
 * @param apdu A buffer contaning the APDU
 * @param apdu_len The length of the APDU buffer. Must be at least 6 bytes long
 * @return The error code for the operation
 */
DK_Result scp_verify_apdu_rmac(
    scp_context *ctx, 
    byte* apdu, 
    size_t apdu_len);

/**
 * Decrypts the response APDU data field and strips the R-MAC and 
 * padding added for the encryption process.
 * 
 * @param ctx The SCP context structure
 * @param wrapped_apdu A buffer contaning the wrapped APDU
 * @param wrapped_apdu_len The length of the wrapped APDU buffer.
 * @param apdu A buffer that will hold the unwrapped APDU. This must be preallocated with at least RPDU_MAX_SIZE bytes
 * @param apdu_len A pointer to an size_t variable that will hold the unwrapped APDU length
 * @return The error code for the operation
 */
DK_Result scp_decrypt_response(
    scp_context *ctx, 
    byte* wrapped_apdu, 
    size_t wrapped_apdu_len, 
    byte* apdu, 
    size_t *apdu_len);

#ifdef DK_DEBUG
extern DK_Result (*test_scp_generate_host_challenge_ptr)(scp_context *ctx);
DK_Result mock_scp_generate_host_challenge_test_open_secure_channel(scp_context *ctx);
DK_Result mock_scp_generate_host_challenge_test_generate_init_update(scp_context *ctx);
DK_Result mock_scp_generate_host_challenge_test_full(scp_context *ctx);
void scp_dump_context(scp_context *ctx);
#endif

#endif