#include "AES.h"
#include <assert.h>

// write data from src to array and increment array data pointer
static inline void __Array_writeNext(Array(u8)* dst, u8* src, size_t size){
    memcpy(dst->data, src, size);
    *dst = Array_u8_sliceFrom(*dst, size);
}

// read data from array to dst and increment array data pointer
static inline void __Array_readNext(u8* dst, Array(u8)* src, size_t size){
    memcpy(dst, src->data, size);
    *src = Array_u8_sliceFrom(*src, size);
}

static void __calcKeyCheckSum(Array(u8) key, void* dst){
    br_sha256_context sha_ctx;
    br_sha256_init(&sha_ctx);
    br_sha256_update(&sha_ctx, key.data, key.len);
    br_sha256_out(&sha_ctx, dst);
}

//////////////////////////////////////////////////////////////////////////////
//                            AESBlockEncryptor                             //
//////////////////////////////////////////////////////////////////////////////

void AESBlockEncryptor_construct(AESBlockEncryptor* ptr,
    Array(u8) key, const br_block_cbcenc_class* enc_class)
{   
    ptr->enc_class = enc_class;
    AESBlockEncryptor_changeKey(ptr, key);
    
    ptr->rng_ctx.vtable = &br_hmac_drbg_vtable;
    rng_init_sha256_seedFromSystem(&ptr->rng_ctx.vtable);

}

void AESBlockEncryptor_changeKey(AESBlockEncryptor* ptr, Array(u8) key)
{
    assert(key.len == 16 || key.len == 24 || key.len == 32);
    ptr->enc_class->init((void*)ptr->enc_keys, key.data, key.len);
    __calcKeyCheckSum(key, ptr->key_checksum);
}

Result(u32) AESBlockEncryptor_encrypt(AESBlockEncryptor* ptr,
    Array(u8) src, Array(u8) dst)
{
    Deferral(4);
    u32 encrypted_size = AESBlockEncryptor_calcDstSize(src.len);
    try_assert(dst.len >= encrypted_size);
    
    // generate random initial vector
    br_hmac_drbg_generate(&ptr->rng_ctx, ptr->iv, __AES_BLOCK_IV_SIZE);
    // write IV to the beginning of dst
    __Array_writeNext(&dst, ptr->iv, __AES_BLOCK_IV_SIZE);

    EncryptedBlockHeader header;
    zeroStruct(&header);
    memcpy(header.key_checksum, ptr->key_checksum, __AES_BLOCK_KEY_CHECKSUM_SIZE);
    header.padding_size = (16 - src.len % 16) % 16;
    // write header to buffer
    memcpy(ptr->buf, &header, sizeof(header));
    // encrypt header
    ptr->enc_class->run((void*)ptr->enc_keys, ptr->iv, ptr->buf, sizeof(header));
    // write encrypted header to dst
    __Array_writeNext(&dst, ptr->buf, sizeof(header));

    // encrypt full AESBlockEncryptor buffers
    while(src.len > __AES_BUFFER_SIZE){
        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
        ptr->enc_class->run((void*)ptr->enc_keys,
            ptr->iv,
            ptr->buf, __AES_BUFFER_SIZE);
        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
    }

    // encrypt buffer with remaining data
    if(src.len > 0){
        memcpy(ptr->buf, src.data, src.len);
        memset(ptr->buf + src.len, 0, header.padding_size);
        u32 src_size_padded = src.len + header.padding_size;
        ptr->enc_class->run((void*)ptr->enc_keys,
            ptr->iv,
            ptr->buf, src_size_padded);
        memcpy(dst.data, ptr->buf, src_size_padded);
    }

    Return RESULT_VALUE(u, encrypted_size);
}


//////////////////////////////////////////////////////////////////////////////
//                            AESBlockDecryptor                             //
//////////////////////////////////////////////////////////////////////////////

void AESBlockDecryptor_construct(AESBlockDecryptor* ptr,
    Array(u8) key, const br_block_cbcdec_class* dec_class)
{
    ptr->dec_class = dec_class;
    AESBlockDecryptor_changeKey(ptr, key);
}

void AESBlockDecryptor_changeKey(AESBlockDecryptor* ptr, Array(u8) key)
{
    assert(key.len == 16 || key.len == 24 || key.len == 32);
    ptr->dec_class->init((void*)ptr->dec_keys, key.data, key.len);
    __calcKeyCheckSum(key, ptr->key_checksum);
}

Result(u32) AESBlockDecryptor_decrypt(AESBlockDecryptor* ptr,
    Array(u8) src, Array(u8) dst)
{
    Deferral(4);
    u32 overhead_size = AESBlockEncryptor_calcDstSize(0);
    try_assert(src.len >= overhead_size);
    try_assert(src.len % 16 == 0 && "src must be array of 16-byte blocks");
    try_assert(dst.len >= src.len - overhead_size);
    
    // read IV from the beginning of src
    __Array_readNext(ptr->iv, &src, __AES_BLOCK_IV_SIZE);

    EncryptedBlockHeader header;
    // read encrypted header from src
    __Array_readNext((void*)&header, &src, sizeof(header));
    // decrypt header
    ptr->dec_class->run((void*)ptr->dec_keys, ptr->iv, &header, sizeof(header));

    // validate decrypted data
    if(memcmp(header.key_checksum, ptr->key_checksum, __AES_BLOCK_KEY_CHECKSUM_SIZE) != 0){
        Return RESULT_ERROR("decrypted data is invalid or key is wrong", false);
    }

    // size of decrypted data without padding
    try_assert(src.len >= header.padding_size && "invalid padding size");
    u32 decrypted_size = src.len - header.padding_size;
    src.len = decrypted_size;

    // decrypt full buffers
    while(src.len > __AES_BUFFER_SIZE){
        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
        ptr->dec_class->run((void*)ptr->dec_keys, 
            ptr->iv, 
            ptr->buf, __AES_BUFFER_SIZE);
        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
    }
    
    // decrypt buffer with remaining data
    if(src.len > 0){
        memcpy(ptr->buf, src.data, src.len);
        memset(ptr->buf + src.len, 0, header.padding_size);
        u32 src_size_padded = src.len + header.padding_size;
        ptr->dec_class->run((void*)ptr->dec_keys, 
            ptr->iv,
            ptr->buf, src_size_padded);
        memcpy(dst.data, ptr->buf, src.len);
    }

    Return RESULT_VALUE(u, decrypted_size);
}


//////////////////////////////////////////////////////////////////////////////
//                            AESStreamEncryptor                            //
//////////////////////////////////////////////////////////////////////////////

void AESStreamEncryptor_construct(AESStreamEncryptor* ptr, 
    Array(u8) key, const br_block_ctr_class* ctr_class)
{
    ptr->ctr_class = ctr_class;
    AESStreamEncryptor_changeKey(ptr, key);
    ptr->block_counter = 0;
    
    br_hmac_drbg_context rng_ctx;
    rng_ctx.vtable = &br_hmac_drbg_vtable;
    rng_init_sha256_seedFromSystem(&rng_ctx.vtable);
    br_hmac_drbg_generate(&rng_ctx, ptr->iv, __AES_STREAM_IV_SIZE);
}

void AESStreamEncryptor_changeKey(AESStreamEncryptor* ptr, Array(u8) key)
{
    assert(key.len == 16 || key.len == 24 || key.len == 32);
    ptr->ctr_class->init((void*)ptr->ctr_keys, key.data, key.len);
    __calcKeyCheckSum(key, ptr->key_checksum);
}

Result(u32) AESStreamEncryptor_encrypt(AESStreamEncryptor* ptr,
    Array(u8) src, Array(u8) dst)
{
    Deferral(4);

    u32 encrypted_size = src.len;
    // if it is the beginning of the stream,
    if(ptr->block_counter == 0){
        // write IV generated during initialization
        __Array_writeNext(&dst, ptr->iv, __AES_STREAM_IV_SIZE);
        encrypted_size = AESStreamEncryptor_calcDstSize(encrypted_size);
        
        // encrypt checksum
        u8 key_checksum[__AES_BLOCK_KEY_CHECKSUM_SIZE];
        memcpy(key_checksum, ptr->key_checksum, __AES_BLOCK_KEY_CHECKSUM_SIZE);
        ptr->block_counter = ptr->ctr_class->run((void*)ptr->ctr_keys, 
            ptr->iv, ptr->block_counter,
            key_checksum, __AES_BLOCK_KEY_CHECKSUM_SIZE);
        // write checksum to dst
        __Array_writeNext(&dst, key_checksum, __AES_BLOCK_KEY_CHECKSUM_SIZE);
    }
    try_assert(dst.len >= encrypted_size);
    
    // encrypt full buffers
    while(src.len > __AES_BUFFER_SIZE){
        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
        ptr->block_counter = ptr->ctr_class->run((void*)ptr->ctr_keys, 
            ptr->iv, ptr->block_counter,
            ptr->buf, __AES_BUFFER_SIZE);
        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
    }

    // encrypt remaining data
    if(src.len > 0){
        memcpy(ptr->buf, src.data, src.len);
        ptr->block_counter = ptr->ctr_class->run((void*)ptr->ctr_keys,
            ptr->iv, ptr->block_counter,
            ptr->buf, src.len);
        memcpy(dst.data, ptr->buf, src.len);
    }

    Return RESULT_VALUE(u, encrypted_size);
}


//////////////////////////////////////////////////////////////////////////////
//                            AESStreamDecryptor                            //
//////////////////////////////////////////////////////////////////////////////

void AESStreamDecryptor_construct(AESStreamDecryptor* ptr,
    Array(u8) key, const br_block_ctr_class* ctr_class)
{
    ptr->ctr_class = ctr_class;
    AESStreamDecryptor_changeKey(ptr, key);
    ptr->block_counter = 0;
}

void AESStreamDecryptor_changeKey(AESStreamDecryptor* ptr, Array(u8) key)
{
    assert(key.len == 16 || key.len == 24 || key.len == 32);
    ptr->ctr_class->init((void*)ptr->ctr_keys, key.data, key.len);
    __calcKeyCheckSum(key, ptr->key_checksum);
}

Result(u32) AESStreamDecryptor_decrypt(AESStreamDecryptor* ptr,
    Array(u8) src, Array(u8) dst)
{
    Deferral(4);

    // if it is the beginning of the stream
    if(ptr->block_counter == 0){
        // read random IV
        __Array_readNext(ptr->iv, &src, __AES_STREAM_IV_SIZE);

        // read checksum
        u8 key_checksum[__AES_BLOCK_KEY_CHECKSUM_SIZE];
        __Array_readNext(key_checksum, &src, __AES_BLOCK_KEY_CHECKSUM_SIZE);
        // decrypt checksum
        ptr->block_counter = ptr->ctr_class->run((void*)ptr->ctr_keys, 
            ptr->iv, ptr->block_counter,
            key_checksum, __AES_BLOCK_KEY_CHECKSUM_SIZE);
        // validate decrypted data
        if(memcmp(key_checksum, ptr->key_checksum, __AES_BLOCK_KEY_CHECKSUM_SIZE) != 0){
            Return RESULT_ERROR("decrypted data is invalid or key is wrong", false);
        }
    }
    // size without IV
    u32 decrypted_size = src.len;
    try_assert(dst.len >= decrypted_size);

    // decrypt full buffers
    while(src.len > __AES_BUFFER_SIZE){
        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
        ptr->block_counter = ptr->ctr_class->run((void*)ptr->ctr_keys, 
            ptr->iv, ptr->block_counter,
            ptr->buf, __AES_BUFFER_SIZE);
        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
    }

    // decrypt remaining data
    if(src.len > 0){
        memcpy(ptr->buf, src.data, src.len);
        ptr->block_counter = ptr->ctr_class->run((void*)ptr->ctr_keys,
            ptr->iv, ptr->block_counter,
            ptr->buf, src.len);
        memcpy(dst.data, ptr->buf, src.len);
    }

    Return RESULT_VALUE(u, decrypted_size);
}