#include "ssp.h"
#include "ssp_public.h"
#include "ssp_apdu.h"
#include "ssp_util.h"
#include "crypto_module.h"
#include "tz_platform.h"
#include "ssp_service.h"

#include "tz_debug.h"

/*
 * SSP Public APIs
 */
SSPSTATUS ssp_transaction_public( uint8_t *wrappedSessionInfo, uint32_t wrappedSessionInfoSize, p_secEse_7816_cpdu_t cpdu, p_secEse_7816_rpdu_t rpdu ) {
    SSPSTATUS           status;
    ssp_session_t       hsession;
    apdu_header_t       hcmd;
    apdu_header_t       cmac_hcmd;
    secEse_7816_cpdu_t  cur_cmd;
    secEse_7816_rpdu_t  cur_rsp;

    uint8_t             iv[AES_BLOCK_SIZE];
    uint8_t             tmp_mcv[AES_BLOCK_SIZE];
    uint8_t             tmp_cmac[AES_BLOCK_SIZE];
    uint8_t             response[MAX_SSP_RAPDU_DATA_SIZE];
    uint32_t            enc_data_len;
    uint8_t             tmp_buf[MAX_SSP_CAPDU_DATA_SIZE];
    uint32_t            tmp_buf_len = sizeof(tmp_buf);
    uint8_t             tmp_buf2[MAX_SSP_CAPDU_DATA_SIZE];
    uint32_t            tmp_buf_len2 = sizeof(tmp_buf2);


    LOGD("ssp_transaction is called");
    if ( cpdu == NULL || rpdu == NULL ) {
        status = SSPSTATUS_INVALID_ARGUMENT;
        goto ret_point;
    }

    cur_cmd = *cpdu;

    LOGD("Incoming CPDU : CLA=0x%02X, INS=0x%02X, P1=0x%02X, P2=0x%02X, LC=%u, LE=%u", cur_cmd.cla, cur_cmd.ins, cur_cmd.p1, cur_cmd.p2, cur_cmd.lc, cur_cmd.le);
    hex_print_tag_debug("Incoming CPDU", cur_cmd.pdata, cur_cmd.lc);

#if defined(USE_BLOWFISH) || defined(USE_QSEE) || defined(USE_TRUSTY_UNISOC)
    CHECK_OP_STATUS( ssp_unwrap_secure_object( SEM_TID, (const uint8_t *)wrappedSessionInfo, wrappedSessionInfoSize, (uint8_t *)&hsession, &tmp_buf_len) );
#else
    CHECK_OP_STATUS( ssp_unwrap_secure_object( (const uint8_t *)wrappedSessionInfo, wrappedSessionInfoSize, (uint8_t *)&hsession, &tmp_buf_len) );
#endif

    hex_print_tag_debug("hsession->cmac", hsession.cmac, AES_BLOCK_SIZE);
    hex_print_tag_debug("hsession->enc_key", hsession.enc_key, AES_256_KEY_SIZE);
    hex_print_tag_debug("hsession->mac_key", hsession.mac_key, AES_256_KEY_SIZE);

    /* Generate IV */
    CHECK_OP_STATUS( crypto_aes_256_block_encrypt( hsession.cmac, iv, hsession.enc_key ) );

    if ( cur_cmd.lc != 0 ) {
        /* Encrypt plain data */
        CHECK_OP_STATUS( crypto_aes_cbc_256_encrypt( cpdu->pdata, cur_cmd.lc, tmp_buf, &tmp_buf_len, iv, hsession.enc_key ) );
    } else {
        tmp_buf_len = 0;
    }

    cur_cmd.cla |= 0x04;
    cur_cmd.lc = tmp_buf_len + AES_BLOCK_SIZE;
    cur_cmd.le_type = 1;
    cur_cmd.le = MAX_SSP_RAPDU_DATA_SIZE;

    /* Make C-MAC */
    *(uint32_t*)&hcmd = *(uint32_t*)&cur_cmd;
    hcmd.lc = cur_cmd.lc;
    memcpy(&cmac_hcmd, &cur_cmd, sizeof(cmac_hcmd));
    hcmd.cla |= hsession.channel;
#ifdef OT
    if (hsession.applet_state.version_major == 0x02 && hsession.applet_state.version_minor == 0x00 && hsession.applet_state.version_release == 0x08) {
        cmac_hcmd.cla |= hsession.channel;
    }
#endif

    /* Save mcv for C-MAC calculation */
    memcpy(tmp_mcv, hsession.cmac, AES_BLOCK_SIZE);

    CHECK_OP_STATUS( crypto_aes_256_cmac( hsession.cmac, hsession.mac_key,
                                  tmp_mcv, AES_BLOCK_SIZE,
                                  &cmac_hcmd, sizeof(hcmd),
                                  tmp_buf, tmp_buf_len,
                                  NULL ) );

    /* Compose and run command */
    BUFFER_APPEND( tmp_buf, tmp_buf_len, hsession.cmac, AES_BLOCK_SIZE );
    cur_cmd.pdata = tmp_buf;
    cur_rsp.pdata = response;

    LOGD("Transaction cmd : CLA=0x%02X, INS=0x%02X, P1=0x%02X, P2=0x%02X, LC=%u, LE=%u", cur_cmd.cla, cur_cmd.ins, cur_cmd.p1, cur_cmd.p2, cur_cmd.lc, cur_cmd.le);
    hex_print_tag_debug("Encrypted Cmd", tmp_buf, tmp_buf_len);

    CHECK_OP_STATUS( ese_status_2_ssp_status( secEseTransmit( hsession.channel, &cur_cmd, &cur_rsp ) ) );

    LOGD("Transaction rsp : SW1=0x%02X, SW2=0x%02X, Len=%u", cur_rsp.sw1, cur_rsp.sw2, cur_rsp.len);
    hex_print_tag_debug("Encrypted Rsp", cur_rsp.pdata, cur_rsp.len);

    if ( cur_rsp.sw1 == 0x69 && cur_rsp.sw2 == 0x82 ) {
        LOGE( "eSE returned error 69 82 during secure transaction" );
        status = SSPSTATUS_INVALID_CMAC;
        goto ret_point;
    }

    if ( cur_rsp.len == 0 || cur_rsp.len % AES_BLOCK_SIZE != 0 ) {
        LOGE( "Unexpected response : SW1=0x%02X, SW2=0x%02X, Len=%u", cur_rsp.sw1, cur_rsp.sw2, cur_rsp.len );
        status = SSPSTATUS_UNEXPECTED_RESPONSE;
        goto ret_point;
    }

    /* Generate IV_2 */
    CHECK_OP_STATUS( crypto_aes_256_block_encrypt( hsession.cmac, iv, hsession.enc_key ) );

    /* Save mcv for C-MAC calculation */
    memcpy(tmp_mcv, hsession.cmac, AES_BLOCK_SIZE);

    /* Save data from APDU response */
    enc_data_len = cur_rsp.len - AES_BLOCK_SIZE;
    memcpy( tmp_buf, response, enc_data_len );
    memcpy( hsession.cmac, response + enc_data_len, AES_BLOCK_SIZE );

    /* C-MAC verification */
    CHECK_OP_STATUS( crypto_aes_256_cmac( tmp_cmac, hsession.mac_key,
                                  tmp_mcv, AES_BLOCK_SIZE,
                                  tmp_buf, enc_data_len,
                                  &cur_rsp.sw1, 2,
                                  NULL ) );

    if ( crypto_secure_mac_cmp( hsession.cmac, tmp_cmac ) != 0 ) {
        LOGE( "Invalid MAC in response to secure command" );
        status = SSPSTATUS_INVALID_CMAC;
        goto ret_point;
    }

    CHECK_OP_STATUS( crypto_aes_cbc_256_decrypt( tmp_buf, enc_data_len, tmp_buf2, &tmp_buf_len2, iv, hsession.enc_key ) );

    /* Output response data */
    memcpy( rpdu->pdata, tmp_buf2, tmp_buf_len2 );
    rpdu->len = tmp_buf_len2;
    rpdu->sw1 = cur_rsp.sw1;
    rpdu->sw2 = cur_rsp.sw2;

    LOGD("Outgoing RPDU : SW1=0x%02X, SW2=0x%02X, Len=%u", rpdu->sw1, rpdu->sw2, rpdu->len);
    hex_print_tag_debug("Outgoing RPDU", rpdu->pdata, rpdu->len);

    CHECK_OP_STATUS( ssp_wrap_secure_object( hsession.tid, (const uint8_t *)&hsession, sizeof(ssp_session_t), wrappedSessionInfo, &wrappedSessionInfoSize) );

ret_point:
    return status;
}

SSPSTATUS ssp_closeSession_public( uint8_t *wrappedSessionInfo, uint32_t wrappedSessionInfoSize ) {
    SSPSTATUS       status;
    ssp_session_t   hsession;
    apdu_header_t   hcmd;

    uint8_t buff;
    uint32_t buflen = 0;

    uint32_t        tmp_buf_len = sizeof(hsession);;

    /* TODO: temporary workaround */

    LOGD("ssp_closeSession is called");

#if defined(USE_BLOWFISH) || defined(USE_QSEE) || defined(USE_TRUSTY_UNISOC)
    CHECK_OP_STATUS( ssp_unwrap_secure_object( SEM_TID, (const uint8_t *)wrappedSessionInfo, wrappedSessionInfoSize, (uint8_t *)&hsession, &tmp_buf_len) );
#else
    CHECK_OP_STATUS( ssp_unwrap_secure_object( (const uint8_t *)wrappedSessionInfo, wrappedSessionInfoSize, (uint8_t *)&hsession, &tmp_buf_len) );
#endif

    SET_APDU_SSP_COMMAND( hcmd, 0x80, 0xD5, 0 );
    CHECK_OP_STATUS( ssp_send_service_command( hsession.channel, hcmd, NULL, &buff, &buflen ) );

ret_point:
    ssp_close_service_connection( hsession.channel );

    return status;
}