#include "src/common.h"
#include "src/pkcs7.h"
#include "src/enc_modes.h"
#include "src/base64.h"
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <time.h>
#include <openssl/evp.h>
#include <openssl/rand.h>
#include <assert.h>
#include "set2.h"

Result_t OpenSSL::Cbc(CryptoAttribs_t* i_attribs,
                    const Key_t* const i_key)
{
    if(NULL==i_attribs->output)
    {
        i_attribs->output = (uint8_t*) malloc(i_attribs->input_len);
    }

    int ret=0;
    EVP_CIPHER_CTX ctx;
    EVP_CIPHER_CTX_init(&ctx);
    OP_CHECK(
        EVP_CipherInit_ex(&ctx, EVP_aes_128_cbc(), NULL, i_key->key, i_attribs->iv,
                          i_attribs->operation==kEncrypt ? 1 : 0));
    EVP_CIPHER_CTX_set_padding(&ctx, 0);
    EVP_CIPHER_CTX_set_key_length(&ctx, i_key->len);

    OP_CHECK(
        EVP_CipherUpdate(&ctx, i_attribs->output, &ret, i_attribs->input,
            i_attribs->input_len));
    i_attribs->output_len = ret;
    OP_CHECK(
        EVP_CipherFinal_ex(&ctx, &i_attribs->output[ret], &ret));
    i_attribs->output_len += ret;
    EVP_CIPHER_CTX_cleanup(&ctx);
    return Result_OK;

end:
    return Result_Error;
}

TCASE(ecb_encrypt_decrypt_single_block)
{
    static const uint8_t expected_result[17] = "0123456789123456";
    CryptoAttribs_t encode, decode;
    Key_t key1;
    CryptoAttribs_t::Init(&encode);
    CryptoAttribs_t::Init(&decode);
    Key_t::Init(&key1);
    uint8_t input[16];
    uint8_t output[16];
    uint8_t key[16];

    encode.input = input;
    encode.input_len= sizeof(input);
    decode.input = output;
    decode.input_len= sizeof(output);
    key1.key = key;
    key1.len = sizeof(key);

    memcpy(key1.key, "YELLOW SUBMARINE", 16);
    memcpy(encode.input, expected_result, 16);
    CHECK(ecb_encrypt(&encode, &key1) == Result_OK);
    CHECK(encode.output_len==16);

    memcpy(decode.input, encode.output, 16);
    ecb_decrypt(&decode, &key1);
    CHECK(decode.output_len == 16);

    CHECK(memcmp(decode.output, expected_result, 16) == 0);
    free(encode.output);
    free(decode.output);
}
TCASE_E

TCASE(set2_challange9)
{
    const uint8_t text1[] = "Text1";
    uint8_t unpadded[32] = {0};
    uint8_t* p_unpadded = &unpadded[0];

    int ret;
    uint8_t buff1[32];
    uint8_t* p_buff1 = &buff1[0];
    char buff2[10];

    ret = pkcs7_pad(text1, 5, &p_buff1, 32);
    CHECK(ret==true, (const unsigned char* const) "Padding operation failed");
    ret = memcmp(text1, buff1, 5);
    CHECK(ret==0, (const unsigned char* const) "Content differs");
    for(int i=5; i<32; i++)
    {
        CHECK(buff1[i] == 27, (const unsigned char* const) "Wrong padding");
    }

    size_t unpad_len = 0;
    CHECK( pkcs7_unpad(&buff1[0], 32, &p_unpadded, &unpad_len) );
    CHECK(unpad_len == 5, (const uint8_t*) "Unpadded length wrong");
    CHECK( 0==memcmp(text1, unpadded, unpad_len) );
}
TCASE_E

TCASE(set2_challange10)
{
    static const uint8_t expected_result[34] = "I'm back and I'm ringin' the bell";
    CryptoAttribs_t attribs;
    Key_t key;
    Result_t res = Result_Error;

    Key_t::Init(&key);
    key.len = 16;
    key.key = (uint8_t*) malloc(key.len);
    memcpy(key.key,"YELLOW SUBMARINE",key.len);

    CryptoAttribs_t::Init(&attribs);
    attribs.iv = (uint8_t*) malloc(key.len);
    attribs.iv_len = key.len;
    memset(attribs.iv, 0, attribs.iv_len);

    res = load_base64_to_hex(
                "sol/etc/set2_t2.txt",
                &(attribs.input),
                &(attribs.input_len));
    CHECK(res == Result_OK, (const uint8_t*) "Problem when loading input file");
    cbc_decrypt(&attribs, &key);
    CHECK( memcmp(expected_result, attribs.output, 33) == 0, (const uint8_t*)
        "Wrong plaintext decrypted");
    // cleanup
    CryptoAttribs_t::Free(&attribs);
    Key_t::Free(&key);
}
TCASE_E

TCASE(cbc_enc_dec_test)
{
    static const uint8_t test_text[49] = "The quick brown fox jumps over the lazy mad dog.";
    CryptoAttribs_t enc_attribs;
    CryptoAttribs_t dec_attribs;
    Key_t key;
    CryptoAttribs_t::Init(&enc_attribs);
    CryptoAttribs_t::Init(&dec_attribs);
    Key_t::Init(&key);

    // Setup key
    key.len = 16;
    key.key = (uint8_t*) malloc(key.len);
    memcpy(key.key, "0123456789123456", key.len);

    // Encrypt
    enc_attribs.input_len = sizeof(test_text)-1;
    enc_attribs.input = (uint8_t*) malloc(enc_attribs.input_len);
    enc_attribs.operation = kEncrypt;
    memcpy( enc_attribs.input, test_text, enc_attribs.input_len);
    enc_attribs.iv_len = key.len;
    enc_attribs.iv = (uint8_t*)malloc(enc_attribs.iv_len);
    memset(enc_attribs.iv, 0, enc_attribs.iv_len);
    cbc_encrypt(&enc_attribs, &key);

    // Setup decryption object.
    dec_attribs.iv = (uint8_t*)malloc( enc_attribs.iv_len );
    dec_attribs.iv_len = enc_attribs.iv_len;
    dec_attribs.operation = kDecrypt;
    memset(dec_attribs.iv, 0, dec_attribs.iv_len);
    dec_attribs.input = (uint8_t*)malloc(enc_attribs.output_len);
    dec_attribs.input_len = enc_attribs.output_len;
    memcpy(dec_attribs.input, enc_attribs.output, enc_attribs.output_len);

    // Decrypt
    cbc_decrypt(&dec_attribs, &key);
    CHECK( dec_attribs.output_len+1 == sizeof(test_text), (uint8_t*) "wrong size of ouptut");
    CHECK(memcmp(dec_attribs.output, test_text, dec_attribs.output_len) == 0, dec_attribs.output);

end:
    CryptoAttribs_t::Free(&enc_attribs);
    CryptoAttribs_t::Free(&dec_attribs);
    Key_t::Free(&key);
}
TCASE_E

TCASE(encode_decode_openssl)
{
    uint8_t     concatenated_blocks[16*3];
    uint8_t     iv1[16]  = {0};
    uint8_t     key[17] = "YELLOW SUBMARINE"; // 16+'\0' = 17
    memset(concatenated_blocks, 2, 16);
    memset(&concatenated_blocks[16], 5, 16);
    memset(&concatenated_blocks[32], 9, 16);

    RAND_bytes(iv1, 16);
    Key_t keyObj;
    keyObj.key = &key[0];
    keyObj.len = 16;

    // 1. decryption test
    {
        CryptoAttribs_t attribs_openssl_enc;
        CryptoAttribs_t::Init(&attribs_openssl_enc);
        attribs_openssl_enc.input = concatenated_blocks;
        attribs_openssl_enc.input_len = sizeof(concatenated_blocks);
        attribs_openssl_enc.output = (uint8_t*)malloc(attribs_openssl_enc.input_len);
        attribs_openssl_enc.output_len;
        attribs_openssl_enc.iv = &iv1[0];
        attribs_openssl_enc.iv_len = sizeof(iv1);
        attribs_openssl_enc.operation = kEncrypt;
        CHECK(OpenSSL::Cbc(&attribs_openssl_enc, &keyObj)==Result_OK);
        CHECK(attribs_openssl_enc.output_len==48, (uint8_t*)"Ciphertext has wrong size");//
        CryptoAttribs_t cbc_attribs;
        CryptoAttribs_t::Init(&cbc_attribs);
        cbc_attribs.input = attribs_openssl_enc.output;
        cbc_attribs.input_len = attribs_openssl_enc.output_len;
        cbc_attribs.iv = &iv1[0];
        cbc_attribs.iv_len = sizeof(iv1);
        CHECK( Result_OK == cbc_decrypt(&cbc_attribs, &keyObj) );
        CHECK(
            memcmp( concatenated_blocks,
                    cbc_attribs.output,
                    cbc_attribs.output_len) == 0,
                    (uint8_t*)"Input/Output differs");
        ::free(attribs_openssl_enc.output);
        ::free(cbc_attribs.output);
    }
    // 2. encryption test
    {
        CryptoAttribs_t attribs_enc;
        CryptoAttribs_t::Init(&attribs_enc);
        attribs_enc.input = concatenated_blocks;
        attribs_enc.input_len = sizeof(concatenated_blocks);
        attribs_enc.iv = &iv1[0];
        attribs_enc.iv_len = sizeof(iv1);
        attribs_enc.operation = kEncrypt;
        CHECK( Result_OK == cbc_encrypt(&attribs_enc, &keyObj) );
        CHECK( 48 == attribs_enc.output_len );
        CHECK( NULL != attribs_enc.output );

        CryptoAttribs_t attribs_openssl_dec;
        CryptoAttribs_t::Init(&attribs_openssl_dec);
        attribs_openssl_dec.input = attribs_enc.output;
        attribs_openssl_dec.input_len = attribs_enc.output_len;
        attribs_openssl_dec.iv = &iv1[0];
        attribs_openssl_dec.iv_len = sizeof(iv1);
        attribs_openssl_dec.operation = kDecrypt;
        CHECK(OpenSSL::Cbc(&attribs_openssl_dec, &keyObj)==Result_OK);
        CHECK(attribs_openssl_dec.output_len==48, (uint8_t*)"Ciphertext has wrong size");
        CHECK(
            memcmp( concatenated_blocks,
                    attribs_openssl_dec.output,
                    attribs_openssl_dec.output_len) == 0,
                    (uint8_t*)"Input/Output differs");
        ::free(attribs_enc.output);
        ::free(attribs_openssl_dec.output);
    }
}
TCASE_E

/* -------------------------------------------------------------------------
    Checks if ciphertext is encrypted with ECB or CBC.
    Requirement: In order to work, plaintext must encrypt up to 6 blocks (with
                 padding) and all the letters in the PT must be exactly the same.
    Returns:
    * 0   : if ciphertext encrypted with ECB
    * 1   : if ciphertext encrypted with CBC
    * 0xFF: in case wrong input parameters provided
   -------------------------------------------------------------------------*/
uint8_t check_ciphertext(const uint8_t* ciphertext, uint32_t len )
{
    uint32_t bs = 16;
    if( (len < 6*bs) || (len>7*bs) )
        return 0xFF;

    // check blocks 2,3,4 as first have prepaned bytes, 5th also, 6th can have padding
    // eventually 7th may be fully padded.
    for(uint32_t i=1; i<4; ++i)
    {
        if( memcmp(ciphertext+(bs*i), ciphertext+(bs*(i+1)), bs) != 0)
            return 1; // Two blocks differ -> must be CBC
    }
    return 0; // All checked blocks are the same -> ECB
}

Result_t encryption_oracle( const uint8_t* pt,
                            uint32_t pt_len,
                            uint8_t* ct,
                            uint32_t* ct_len,
                            uint8_t* mode)
{
    uint8_t         iv1[16]  = {0};
    uint8_t         key[16]  = {0};
    uint8_t         prepend_byte_size;
    uint8_t         append_byte_size;
    uint8_t         random_bytes[10];
    uint8_t         encryption_mode = 0xFF;
    struct timespec tv;
    CryptoAttribs_t attribs;
    Key_t           keyObj;

    // get seed for good randomnes
    clock_gettime(CLOCK_MONOTONIC, &tv);
    srand((unsigned int)tv.tv_nsec);

    // choose random sizes and enc mode (0 ECB, 1 CBC)
    prepend_byte_size = 5+(rand() % 6);
    append_byte_size  = 5+(rand() % 6);
    encryption_mode   = rand() % 2;

    // generate rando data
    RAND_bytes(iv1, 16);
    RAND_bytes(key, 16);

    // initialize encryption attribs
    CryptoAttribs_t::Init(&attribs);
    attribs.operation = kEncrypt;
    attribs.input_len = 0;
    attribs.input = (uint8_t*)malloc(7*16);
    attribs.output_len = 0;
    attribs.output = (uint8_t*)malloc(7*16);
    attribs.padding = kPadding_PKCS7;

    // copy input to the input buffer
    RAND_bytes(random_bytes, prepend_byte_size);
    int tmplen=0;
    memcpy(attribs.input+tmplen, random_bytes, prepend_byte_size);
    tmplen+=prepend_byte_size;
    memcpy(attribs.input+tmplen, pt, pt_len);
    tmplen+=pt_len;
    RAND_bytes(random_bytes, append_byte_size);
    memcpy(attribs.input+tmplen, random_bytes, append_byte_size);
    tmplen+=append_byte_size;
    attribs.input_len = tmplen;

    // key
    Key_t::Init(&keyObj);
    keyObj.key = key;
    keyObj.len = 16;

    Result_t ret = Result_OK;
    do
    {
        if(encryption_mode)
        {
            attribs.iv = (uint8_t*)malloc(16);
            attribs.iv_len = 16;
            RAND_bytes(attribs.iv, 16);
            if( Result_OK!=cbc_encrypt(&attribs, &keyObj) )
            {
                ret = Result_Error;
                break;
            }
        }
        else
        {
            if( Result_OK!=ecb_encrypt(&attribs, &keyObj) )
            {
                ret = Result_Error;
                break;
            }
        }

        // Copy results
        memcpy(ct, attribs.output, attribs.output_len);
        *ct_len = attribs.output_len;
        *mode = encryption_mode;

        ret = Result_OK;
    } while(false);

    CryptoAttribs_t::Free(&attribs);
    return ret;
}

TCASE(set2_challange11)
{
    // let's run it 1000 times
    for(size_t i=0; i<1000; ++i)
    {
        uint8_t plaintext[5*16] = {0};
        uint8_t ciphertext[7*16] = {0};
        uint8_t mode = 0xFF;
        uint8_t guessed_mode = 0xFF;
        uint32_t ciphertext_len = 7*16;

        // set data in blocks so that 5 blocks is 11111...
        memset(plaintext, 1, 5*16);
        Result_t ret = encryption_oracle(plaintext, 5*16, ciphertext, &ciphertext_len, &mode);
        CHECK(ret == Result_OK, (uint8_t*)"Error occured on encryption");
        CHECK(mode != 0xFF);

        guessed_mode = check_ciphertext(ciphertext, ciphertext_len);
        CHECK(mode == guessed_mode);
    }
}
TCASE_E

TCASE(encrypt_padding_pkcs7)
{
    uint8_t text[3] = {'D', 'E', 'F'};
    uint8_t expected_no_padding_dec[16] = {0};
    uint8_t iv[16] = {0};
    uint8_t key[16] = {0};

    // "ABC" must be padded with 13,13,13....
    memcpy(expected_no_padding_dec, text, 3);
    memset(expected_no_padding_dec+3, 13, 13);

    // key
    Key_t keyObj = {0};
    Key_t::Init(&keyObj);
    keyObj.key = key;
    keyObj.len = 16;

    CryptoAttribs_t attribs_enc;
    CryptoAttribs_t::Init(&attribs_enc);
    attribs_enc.input = &text[0];
    attribs_enc.input_len = 3;
    attribs_enc.iv = &iv[0];
    attribs_enc.iv_len = sizeof(iv);
    attribs_enc.operation = kEncrypt;
    attribs_enc.padding = kPadding_PKCS7;
    CHECK( Result_OK == cbc_encrypt(&attribs_enc, &keyObj) );
    CHECK( attribs_enc.output_len == 16, (uint8_t*) "Wrong out size");

    // Decrypt with openssl and no padding. Check padding value
    {
        CryptoAttribs_t attribs_openssl_dec;
        CryptoAttribs_t::Init(&attribs_openssl_dec);
        attribs_openssl_dec.input = attribs_enc.output;
        attribs_openssl_dec.input_len = attribs_enc.output_len;
        attribs_openssl_dec.iv = &iv[0];
        attribs_openssl_dec.iv_len = sizeof(iv);
        attribs_openssl_dec.operation = kDecrypt;
        attribs_openssl_dec.padding = kPadding_None;
        CHECK(OpenSSL::Cbc(&attribs_openssl_dec, &keyObj)==Result_OK);
        CHECK(attribs_openssl_dec.output_len==16, (uint8_t*)"Ciphertext has wrong size");
        CHECK( 0==memcmp(   attribs_openssl_dec.output,
                            expected_no_padding_dec,
                            16), (uint8_t*)"Wrong padding decrypted" );

        ::free(attribs_openssl_dec.output);
    }

    // Decrypt padding
    {
        CryptoAttribs_t attribs_dec;
        CryptoAttribs_t::Init(&attribs_dec);
        attribs_dec.input = attribs_enc.output;
        attribs_dec.input_len = attribs_enc.output_len;
        attribs_dec.iv = &iv[0];
        attribs_dec.iv_len = sizeof(iv);
        attribs_dec.operation = kDecrypt;
        attribs_dec.padding = kPadding_PKCS7;
        CHECK( Result_OK == cbc_decrypt(&attribs_dec, &keyObj) );
        CHECK(attribs_dec.output_len==3, (uint8_t*)"Ciphertext has wrong size");
        CHECK( 0==memcmp(   attribs_dec.output,
                            expected_no_padding_dec,
                            3), (uint8_t*)"Wrong padding decrypted" );

        ::free(attribs_dec.output);
    }
    ::free(attribs_enc.output);
}
TCASE_E

TCASE(set2_challange12_not_finished)
{
    uint8_t CIPHERTEXT[] = "Um9sbGluJyBpbiBteSA1LjAKV2l0aCBteSByYWctdG9wIGRvd24gc28gbXkg"
                           "aGFpciBjYW4gYmxvdwpUaGUgZ2lybGllcyBvbiBzdGFuZGJ5IHdhdmluZyBq"
                           "dXN0IHRvIHNheSBoaQpEaWQgeW91IHN0b3A/IE5vLCBJIGp1c3QgZHJvdmUg"
                           "YnkK";

    uint8_t HEX_STRING[sizeof CIPHERTEXT];
    uint8_t CONCAT_STRING[sizeof(CIPHERTEXT)*2];
    CryptoAttribs_t EncryptForBS;

    Key_t key;
    uint32_t hex_len = 0;
    hex_len = base64_to_hex(CIPHERTEXT, sizeof(CIPHERTEXT), HEX_STRING);


}
TCASE_E