#include "SCP03_crypto.h"
#include "tz_debug.h"
#include "tz_utils.h"

#ifdef USE_SCRYPTO    
    #include "openssl/cmac.h"
#endif

#if defined(USE_BLOWFISH) || defined(USE_TRUSTY_UNISOC)
#include <tee_internal_api.h>

uint16_t  scp03_gen_random( uint8_t *out, uint32_t size ) {
    uint32_t desired;

    while (size != 0) {
        desired = size;
        TEE_GenerateRandom( out, desired );

        out += desired;
        size -= desired;
    }

    return 0;
}
#endif

#ifdef USE_MOBICORE
#include <TlApi/TlApi.h>

uint16_t  scp03_gen_random( uint8_t *out, uint32_t size ) {
    int32_t res;
    uint32_t desired;

    while (size != 0) {
        desired = size;
        res = tlApiRandomGenerateData( TLAPI_ALG_SECURE_RANDOM, out, &desired );

        if (TLAPI_OK != res ) {
            LOGE("tlApiRandomGenerateData failed with code %d", res);
            return 1;
        }

        out += desired;
        size -= desired;
    }

    return 0;
}
#endif

#ifdef USE_QSEE
#include <qsee_prng.h>

uint16_t  scp03_gen_random( uint8_t *out, uint32_t size ) {
    int32_t  res;
    uint8_t  *out_ptr = out;

    while ( size != 0 ) {
        res = qsee_prng_getdata( out_ptr, (size < QSEE_MAX_PRNG) ? size : QSEE_MAX_PRNG );
        if ( res <= 0 ) {
            LOGE( "qsee_prng_getdata failed with code %d", res );
            return 1;
        }

        out_ptr += res;
        size -= res;
    }

    return 0;
}
#endif

#ifdef USE_SCRYPTO    
int SCP03_AES_CMAC ( unsigned char *key, unsigned char *input, int length, unsigned char *mac) {
    
    return AES_CMAC(mac, key, KEY_LEN_BYTE, input, length);
}
#else /* if USE_SCRYPTO */

/* License information: 
Name    - INTERNET DRAFT "AES-CMAC-96 Algorithm"
Source  - https://tools.ietf.org/html/draft-songlee-aes-cmac-96-00
Authors - JunHyuk Song, Jicheol Lee
Company - Samsung Electronics
License - 2008 IETF Trust https://trustee.ietf.org/documents/trust-legal-provisions/
*/

#define AES_128 0

unsigned char const_Rb[16] = {
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x87
};
unsigned char const_Zero[16] = {
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
};

void xor_128(unsigned char *a, unsigned char *b, unsigned char *out) {
    int i;

    for (i = 0 ; i < 16 ; i++) {
        out[i] = a[i] ^ b[i];
    }
}

/* AES-CMAC Generation Function */
void leftshift_onebit(unsigned char *input,unsigned char *output) {
    int i;
    unsigned char overflow = 0;

    for (i = 15 ; i >= 0 ; i--) {
        output[i] = input[i] << 1;
        output[i] |= overflow;
        overflow = (input[i] & 0x80) ? 1:0;
    }
}

void generate_subkey(unsigned char *key, unsigned char *K1, unsigned char *K2) {
    unsigned char L[16];
    unsigned char Z[16];
    unsigned char tmp[16];
    int i;
    AES_KEY aes = {0};

    for (i=0 ; i < 16 ; i++) {
        Z[i] = 0;
    }

    AES_set_encrypt_key(key, KEY_LEN_BIT, &aes);
    AES_encrypt(Z, L, &aes);

    if ((L[0] & 0x80) == 0) { /* If MSB(L) = 0, then K1 = L << 1 */
        leftshift_onebit(L, K1);
    } else {    /* Else K1 = ( L << 1 ) (+) Rb */
        leftshift_onebit(L, tmp);
        xor_128(tmp, const_Rb, K1);
    }

    if ( (K1[0] & 0x80) == 0 ) {
        leftshift_onebit(K1, K2);
    } else {
        leftshift_onebit(K1, tmp);
        xor_128(tmp, const_Rb, K2);
    }
}

void padding ( unsigned char *lastb, unsigned char *pad, int length) {
    int j;

    /* original last block */
    for (j = 0 ; j < 16 ; j++) {
        if ( j < length ) {
            pad[j] = lastb[j];
        } else if (j == length) {
            pad[j] = 0x80;
        } else {
            pad[j] = 0x00;
        }
    }
}

int SCP03_AES_CMAC ( unsigned char *key, unsigned char *input, int length, unsigned char *mac) {
    unsigned char X[16],Y[16], M_last[16], padded[16];
    unsigned char K1[16], K2[16];
    int n, i, flag;
    AES_KEY aes = {0};

    generate_subkey(key, K1, K2);
    n = (length+15) / 16;       /* n is number of rounds */
    if (n == 0) {
        n = 1;
        flag = 0;
    } else {
        if ( (length%16) == 0 ) { /* last block is a complete block */
            flag = 1;
        } else { /* last block is not complete block */
            flag = 0;
        }
    }

    if (flag) { /* last block is complete block */
        xor_128(&input[16*(n-1)], K1, M_last);
    } else {
        padding(&input[16*(n-1)], padded, length % 16);
        xor_128(padded,K2,M_last);
    }

    AES_set_encrypt_key(key, KEY_LEN_BIT, &aes);
    for (i = 0 ; i < 16 ; i++) {
        X[i] = 0;
    }

    for (i = 0 ; i < n-1 ; i++) {
        xor_128(X, &input[16*i], Y); /* Y := Mi (+) X  */
        AES_encrypt(Y, X, &aes); /* X := AES-128(KEY, Y); */
    }

    xor_128(X, M_last, Y);
    AES_encrypt(Y, X, &aes);
    for (i = 0 ; i < 16 ; i++) {
        mac[i] = X[i];
    }

    return 16;
}
#endif /* if not USE_SCRYPTO */

void scp03_aes_encrypt (uint8_t* enc_key, uint8_t* data, int data_length_byte, uint8_t* encrypted_data, uint8_t* IV, int padding_type) {
    int i;
    int block_size;
    int encrypted_data_size_byte;
    uint8_t* padded_data;
    AES_KEY aes_key = {0};
    uint8_t tempIV[AES_BLOCK_SIZE] = {0x00, };

    if (padding_type == NO_PADDING) {
        block_size = data_length_byte / 16;
    } else {
        block_size = data_length_byte / 16 + 1;
    }
    encrypted_data_size_byte = block_size*16;

    padded_data = (uint8_t*)tz_malloc(encrypted_data_size_byte);
    if (padded_data == NULL) {
        LOGE("memory is not allocated");
        return;
    }
    memcpy(padded_data,data,data_length_byte);

    if (padding_type == PADDING) {
        for (i = data_length_byte ; i < encrypted_data_size_byte ; i++) {
            if (i == data_length_byte) {
                padded_data[i] = 0x80;
            } else {
                padded_data[i] = 0x00;
            }
        }
    }

    AES_set_encrypt_key(enc_key, KEY_LEN_BIT, &aes_key);
    if (IV == NULL) {
        AES_cbc_encrypt(padded_data, encrypted_data, encrypted_data_size_byte, &aes_key, tempIV, 1);
    } else {
        AES_cbc_encrypt(padded_data, encrypted_data, encrypted_data_size_byte, &aes_key, IV, 1);
    }
    tz_free(padded_data);
}

void scp03_aes_decrypt (uint8_t* dec_key, uint8_t* encrypted_data, uint16_t encrypted_data_length_byte, uint8_t* original_data, int* original_data_len, uint8_t* IV_dec, int padding_type) {
    AES_KEY aes_dec_key = {0};
    int i=1;

    uint8_t* decrypted_data = NULL;
    decrypted_data = (uint8_t*)tz_malloc(encrypted_data_length_byte);
    if ( NULL == decrypted_data )
    {
        LOGE("memory is not allocated");
        return;
    }

    AES_set_decrypt_key(dec_key, KEY_LEN_BIT, &aes_dec_key);
    AES_cbc_encrypt(encrypted_data, decrypted_data, encrypted_data_length_byte, &aes_dec_key, IV_dec, 0);
    if (padding_type == PADDING) {
        while (decrypted_data[encrypted_data_length_byte-i] != 0x80) {
            i++;
        }
        *original_data_len = encrypted_data_length_byte-i;
    } else {
        *original_data_len = encrypted_data_length_byte;
    }
    //*original_data = (uint8_t*)malloc(*original_data_len);
    memcpy(original_data,decrypted_data,*original_data_len);
    tz_free(decrypted_data);
}