use aes::cipher::generic_array::GenericArray;
use aes::cipher::{Block, BlockEncryptMut, KeyInit};
use aes::Aes128;

fn main() {
    let key = [0_u8; 16];
    let iv = [0_u8; 16];

    let plaintext = b"hello world";
    let ciphertext = encrypt_ctr(&key, &iv, plaintext);
    assert_eq!("0e8c27b880aa5b54fa209e", hex::encode(&ciphertext));

    let plaintext = b"hello world.hello world.hello world.hello world.hello world.hello world.";
    let ciphertext = encrypt_ctr(&key, &iv, plaintext);
    assert_eq!("0e8c27b880aa5b54fa209e77a251474237c28ba18812544f5e1a713bcbc7323571e4bee008d3cf\
fe9c08b5d603de9a569ff0c6c7266b2e4c8591edd1fceead8c4f22664e3c1ff0f4", hex::encode(&ciphertext));

    let key = [1_u8; 16];
    let mut iv = [0_u8; 16];
    for i in 0..8 {
        iv[i] = i as u8 + 1;
    }
    let plaintext = b"hello world.hello world.hello world.hello world.hello world.hello world.";
    let ciphertext = encrypt_ctr(&key, &iv, plaintext);
    assert_eq!("84ad8d80732490c061177a58bd26d032d6fcff2e66f9afe3cf95717485d3a4485d4a2a7bd835df\
3d0756b8192e3bf5a287ad8dd81942c43bc812c82d666ebbb34df4e2a5069467d9", hex::encode(&ciphertext));

    let key = [1_u8; 16];
    let iv = [1, 2, 3, 4, 5, 6, 7, 0, 1, 1, 2, 2, 3, 3, 4, 4];
    let plaintext = b"hello world.hello world.hello world.hello world.hello world.hello world.";
    let ciphertext = encrypt_ctr(&key, &iv, plaintext);
    assert_eq!("d373d6b3491c7b8940996f9efefc98d0f564d066a557cf0aa29935a638f99d077a85338e574ee6\
7f69da41b8a87569c324b4d421f1350e11d6be30721605b9dd6401bc3780da3452", hex::encode(&ciphertext));

    let key = [1_u8; 16];
    let iv = [1, 2, 3, 4, 5, 6, 7, 0, 1, 1, 2, 2, 3, 3, 4, 4];
    let plaintext = b"hello world.hello world.hello world.hello world.hello world.h\
ello world.hello world.hello world.hello world.hello world.hello world.hello world.hello w\
orld.hello world.hello world.hello world.hello world.hello world.hello world.hello world.h\
ello world.hello world.hello world.hello world.hello world.hello world.hello world.hello w\
orld.hello world.hello world.hello world.hello world.hello world.hello world.hello world.h\
ello world.hello world.hello world.hello world.hello world.hello world.hello world.hello w\
orld.hello world.hello world.hello world.hello world.hello world.hello world.hello world.h\
ello world.hello world.hello world.hello world.hello world.hello world.hello world.hello w\
orld.hello world.hello world.hello world.hello world.hello world.hello world.hello world.h\
ello world.hello world.hello world.hello world.hello world.hello world.hello world.hello w\
orld.hello world.hello world.hello world.hello world.hello world.hello world.hello world.h\
ello world.hello world.hello world.hello world.";
    let ciphertext = encrypt_ctr(&key, &iv, plaintext);
    assert_eq!("d373d6b3491c7b8940996f9efefc98d0f564d066a557cf0aa29935a638f99d077a85338e574ee6\
7f69da41b8a87569c324b4d421f1350e11d6be30721605b9dd6401bc3780da3452033164403472b0ea84a7894b\
500100959ccc2cfbe2b6792691f1a1b047755c1961a730c7bcd941ab50b4ece7501d624d5af457875370dbc7e6\
0e73fd6ab10b5033fb3b3e2caed7489cc02bf174a3e7f04bd1e5c81c811314689b2bb7f6f0f0333d9b3bf637ff\
7ac34b32bc4905638204dbd97d4cefd1a823b5d49364c72d987807991eb703e47d56631f9f1a62c9f691bd3641\
509884144fc6e1b63e8177c8e1247ccc0769cb59cd973869f2ac3dd6cb90534ee42f69995e4af48de870936d21\
cf1b4ccd52e4bd0eafaff46dd6f34c91072ffd61f73050580b790deba9dc30d9557c4142a0dfa52b70e050a563\
abd69e87150da9325751a58bb3e162ea3a69347684e50710ea2bd25f6c285903d7b200a56a7eb042639410da17\
471c93e2b26b13091cc22ec81c3912a7121a1337a560d9733d904b6c6752993b2943861d3b770693281d3519fb\
69158df2353913aebe1e6c16787b8d2ff363c1c31bb8b491cec0214bd3ef61357569c8ecf9487087e22eca48b2\
d1c2c718ad7d3984eb7d0ca27b96653d90847ed138b581f8a01c0bea33c435e7f124f4009de059c46c006431fa\
5cf41ad9f16e6457389f928216ebd4efd048840a4dbc80da170c666a5293209866e9371f47005a12d6dc76f7fd\
57fc5b265cc38f4aef6e5496836d1c4c6d96b2c807997b3f61fd45f0acf505f3512ea04a0ad89cc35b344bdc3f\
2126b24bee40406ff00a733555316b8ccb57d661326d6b2d48a7a7fe33fbec23d69890b9368f6321d60b231cf2\
5680bcf7e8f43e443c19150b72394b5a9617fb8cc9666d43968fecdb56c22db6421ffbe3351fe479d9e7841288\
15fbeb7f1b5f9dc1f8b3d772741f7ef9884b5ca81d972d8537ea4c12cbf29c7e362ccd14ed83781db312da5d83\
8e11b20b050b6be472d8133f8f6525649dc9b967d40f430feb4cb3408a623a12658959f8cd8a115d8f63dd6e8b\
b8449937620b1cf783a0761fcbd4f48dcb572c1463a7c756bf9dd77d45f92cfe9a04dfb4d007a448d9b1a8c936\
9bcd35db598e7cdec566d075ff4c8d445e94ef49c0d0a78bb45921bf4c6e6b750a4d0f3310fc9a767a368ced10\
293987bf88026c69dcd805aea561d9b207c2710e6e90658ec5e0f8560449889ba8af852ea36b272aebb73ef59e\
74d1c0fe963ecbf1049743263681f5733cdf3b0b89b52667f63f65e4e7327ca45a2c1bd385f7c5fbf74b82e4f8\
0425c237813f2d39c4d7a06b49eec49b26dd234323747cf41eb4c2a7647b8548dcd72d59a4f0da2b1120759cfc\
dae5df80c55a9015fc48f985bda3d416f7301d72a4348953", hex::encode(&ciphertext));
}

fn encrypt_ctr(key: &[u8; 16], iv: &[u8; 16], plaintext: &[u8]) -> Vec<u8> {
    let mut aes128 = Aes128::new_from_slice(key).unwrap();
    let plaintext_len = plaintext.len();
    let blocks_count = (plaintext_len / 16) + if plaintext_len % 16 == 0 { 0 } else { 1 };
    let iv_u128 = u128::from_be_bytes(iv.clone());

    let mut cipher_text = vec![0_u8; plaintext_len];
    for i in 0..blocks_count {
        let enc_iv_u128 = iv_u128 + i as u128;
        let mut block: Block<Aes128> = GenericArray::from(enc_iv_u128.to_be_bytes()).clone();
        aes128.encrypt_block_mut(&mut block);

        let block_slice = block.as_slice();
        for j in 0..block_slice.len() {
            let oi = i * 16 + j;
            if oi < plaintext_len {
                cipher_text[oi] = block_slice[j] ^ plaintext[oi];
            }
        }
    }
    cipher_text
}

#[test]
fn test_aes() {
    let mut aes128 = Aes128::new_from_slice(
        &[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    ).unwrap();

    let mut a: Block<Aes128> = GenericArray::from_slice(&[0_u8; 16]).clone();
    aes128.encrypt_block_mut(&mut a);
    println!("{}", hex::encode(a.as_slice()));

    assert_eq!("66e94bd4ef8a2c3b884cfa59ca342b2e", hex::encode(a.as_slice()));
}