main.rs

//! In Module 1, we discussed Block ciphers like AES. Block ciphers have a fixed length input.
//! Real wold data that we wish to encrypt _may_ be exactly the right length, but is probably not.
//! When your data is too short, you can simply pad it up to the correct length.
//! When your data is too long, you have some options.
//!
//! In this exercise, we will explore a few of the common ways that large pieces of data can be broken
//! up and combined in order to encrypt it with a fixed-length block cipher.
//!
//! WARNING: ECB MODE IS NOT SECURE.
//! Seriously, ECB is NOT secure. Don't use it irl. We are implementing it here to understand _why_ it
//! is not secure and make the point that the most straight-forward approach isn't always the best, and
//! can sometimes be trivially broken.

#![feature(slice_flatten)]
use aes::cipher::{generic_array::GenericArray, BlockDecrypt, BlockEncrypt, KeyInit};
use aes::Aes128;
use rand::RngCore;

///We're using AES 128 which has 16-byte (128 bit) blocks.
const BLOCK_SIZE: usize = 16;

fn random_iv() -> [u8; BLOCK_SIZE] {
    let mut iv = [0u8; BLOCK_SIZE];
    let mut rng = rand::thread_rng();
    rng.fill_bytes(&mut iv);
    iv
}

fn xor_arrays(first: [u8; BLOCK_SIZE], second: [u8; BLOCK_SIZE]) -> [u8; BLOCK_SIZE] {
    first
        .iter()
        .zip(second.iter())
        .map(|(f, s)| f ^ s)
        .collect::<Vec<u8>>()
        .try_into()
        .unwrap()
}

fn main() {
    let key = b"Thats my Kung Fu";

    let plaintext = b"This is a very long message that I want to encrypt with AES-ECB 128.";
    let cipher = ecb_encrypt(plaintext.to_vec(), *key);
    let decrypted = ecb_decrypt(cipher, *key);
    println!("{}", std::str::from_utf8(&decrypted).unwrap());
    assert!(plaintext == &decrypted[..plaintext.len()]);

    let initial_vector = random_iv();
    let plaintext = b"This is a very long message that I want to encrypt with AES-CBC 128.";
    let cipher = cbc_encrypt(plaintext.to_vec(), *key, initial_vector);
    let decrypted = cbc_decrypt(cipher, *key, initial_vector);
    println!("{}", std::str::from_utf8(&decrypted).unwrap());
    assert!(plaintext == &decrypted[..plaintext.len()]);

    let initial_vector = random_iv();
    let plaintext = b"This is a very long message that I want to encrypt with AES-CTR 128.";
    let cipher = ctr(plaintext.to_vec(), *key, initial_vector, true);
    let decrypted = ctr(cipher, *key, initial_vector, false);
    println!("{}", std::str::from_utf8(&decrypted).unwrap());
    assert!(plaintext == &decrypted[..plaintext.len()]);
}

/// Simple AES encryption
/// Helper function to make the core AES block cipher easier to understand.
fn aes_encrypt(data: [u8; BLOCK_SIZE], key: &[u8; BLOCK_SIZE]) -> [u8; BLOCK_SIZE] {
    // Convert the inputs to the necessary data type
    let mut block = GenericArray::from(data);
    let key = GenericArray::from(*key);

    let cipher = Aes128::new(&key);

    cipher.encrypt_block(&mut block);

    block.into()
}

/// Simple AES encryption
/// Helper function to make the core AES block cipher easier to understand.
fn aes_decrypt(data: [u8; BLOCK_SIZE], key: &[u8; BLOCK_SIZE]) -> [u8; BLOCK_SIZE] {
    // Convert the inputs to the necessary data type
    let mut block = GenericArray::from(data);
    let key = GenericArray::from(*key);

    let cipher = Aes128::new(&key);

    cipher.decrypt_block(&mut block);

    block.into()
}

/// Before we can begin encrypting our raw data, we need it to be a multiple of the
/// block length which is 16 bytes (128 bits) in AES128.
///
/// The padding algorithm here is actually not trivial. The trouble is that if we just
/// naively throw a bunch of zeros on the end, there is no way to know, later, whether
/// those zeros are padding, or part of the message, or some of each.
///
/// The scheme works like this. If the data is not a multiple of the block length,  we
/// compute how many pad bytes we need, and then write that number into the last several bytes.
/// Later we look at the last byte, and remove that number of bytes.
///
/// But if the data _is_ a multiple of the block length, then we have a problem. We don't want
/// to later look at the last byte and remove part of the data. Instead, in this case, we add
/// another entire block containing the block length in each byte. In our case,
/// [16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16]
fn pad(mut data: Vec<u8>) -> Vec<u8> {
    // When twe have a multiple the second term is 0
    let number_pad_bytes = BLOCK_SIZE - data.len() % BLOCK_SIZE;

    for _ in 0..number_pad_bytes {
        data.push(number_pad_bytes as u8);
    }

    data
}

/// Groups the data into BLOCK_SIZE blocks. Assumes the data is already
/// a multiple of the block size. If this is not the case, call `pad` first.
fn group(data: Vec<u8>) -> Vec<[u8; BLOCK_SIZE]> {
    data.chunks(BLOCK_SIZE)
        .map(|chunk| {
            let mut block = [0u8; BLOCK_SIZE];
            block.copy_from_slice(chunk);
            block
        })
        .collect()
}

/// Does the opposite of the group function
fn un_group(blocks: Vec<[u8; BLOCK_SIZE]>) -> Vec<u8> {
    blocks.flatten().to_vec()
}

/// Does the opposite of the pad function.
fn un_pad(data: Vec<[u8; BLOCK_SIZE]>) -> Vec<u8> {
    let mut arr = un_group(data);
    let bytes_to_remove = arr[arr.len() - 1] as usize;
    arr.truncate(arr.len() - bytes_to_remove);
    arr
}

/// The first mode we will implement is the Electronic Code Book, or ECB mode.
/// Warning: THIS MODE IS NOT SECURE!!!!
///
/// This is probably the first thing you think of when considering how to encrypt
/// large data. In this mode we simply encrypt each block of data under the same key.
/// One good thing about this mode is that it is parallelizable. But to see why it is
/// insecure look at: https://www.ubiqsecurity.com/wp-content/uploads/2022/02/ECB2.png
fn ecb_encrypt(plain_text: Vec<u8>, key: [u8; 16]) -> Vec<u8> {
    let padded_plain_text = pad(plain_text);
    let grouped_plain_text = group(padded_plain_text);
    let cipher_text = grouped_plain_text
        .iter()
        .map(|&block| aes_encrypt(block, &key))
        .collect();
    un_group(cipher_text)
}

/// Opposite of ecb_encrypt.
fn ecb_decrypt(cipher_text: Vec<u8>, key: [u8; BLOCK_SIZE]) -> Vec<u8> {
    let grouped_cipher_text = group(cipher_text);
    let plain_text = grouped_cipher_text
        .iter()
        .map(|block| aes_decrypt(*block, &key))
        .collect();
    un_pad(plain_text)
}

/// The next mode, which you can implement on your own is cipherblock chaining.
/// This mode actually is secure, and it often used in real world applications.
///
/// In this mode, the ciphertext from the first block is XORed with the
/// plaintext of the next block before it is encrypted.
fn cbc_encrypt(
    plain_text: Vec<u8>,
    key: [u8; BLOCK_SIZE],
    initial_vector: [u8; BLOCK_SIZE],
) -> Vec<u8> {
    // Remember to generate a random initialization vector for the first block.
    let mut data = Vec::default();
    let chunks = group(pad(plain_text));
    let mut to_xor: [u8; BLOCK_SIZE] = initial_vector;
    for ch in chunks.iter() {
        let xored = xor_arrays(*ch, to_xor);
        let encrypted = aes_encrypt(xored, &key);
        to_xor = encrypted;
        data.extend_from_slice(&encrypted);
    }
    data
}

fn cbc_decrypt(
    cipher_text: Vec<u8>,
    key: [u8; BLOCK_SIZE],
    initial_vector: [u8; BLOCK_SIZE],
) -> Vec<u8> {
    let chunks: Vec<[u8; BLOCK_SIZE]> = group(cipher_text);
    let mut unencrypted_chunks: Vec<[u8; BLOCK_SIZE]> = vec![];
    let mut to_xor: [u8; BLOCK_SIZE] = initial_vector;
    for ch in chunks.iter() {
        let data = aes_decrypt(*ch, &key);
        let xored = xor_arrays(data, to_xor);
        unencrypted_chunks.push(xored);
        to_xor = *ch;
    }
    un_pad(unencrypted_chunks)
}

fn ctr(
    input_text: Vec<u8>,
    key: [u8; BLOCK_SIZE],
    initial_vector: [u8; BLOCK_SIZE],
    encrypt: bool,
) -> Vec<u8> {
    // Remember to generate a random initialization vector for the first block.
    let mut data = Vec::default();
    let chunks = if encrypt {
        group(pad(input_text))
    } else {
        group(input_text)
    };
    let mut index = 0;
    let mut counter = [index; BLOCK_SIZE];
    for ch in chunks.iter() {
        let init = xor_arrays(initial_vector, counter);
        let encrypted = aes_encrypt(init, &key);
        let cipher_block = xor_arrays(encrypted, *ch);
        data.extend_from_slice(&cipher_block);
        index += 1;
        counter = [index; BLOCK_SIZE];
    }
    if !encrypt {
        data = un_pad(group(data));
    }
    data
}