plain-aes 0.1.1

use super::key_expansion::ExpandedKey;
use super::RIJNDAEL_S_BOX;
/// Lookup table for the Galois field GF(2^8) multiplication by 2 used in mix columns step.
const MULTIPLE_2: [u8; 256] = [
    0x00, 0x02, 0x04, 0x06, 0x08, 0x0a, 0x0c, 0x0e, 0x10, 0x12, 0x14, 0x16, 0x18, 0x1a, 0x1c, 0x1e,
    0x20, 0x22, 0x24, 0x26, 0x28, 0x2a, 0x2c, 0x2e, 0x30, 0x32, 0x34, 0x36, 0x38, 0x3a, 0x3c, 0x3e,
    0x40, 0x42, 0x44, 0x46, 0x48, 0x4a, 0x4c, 0x4e, 0x50, 0x52, 0x54, 0x56, 0x58, 0x5a, 0x5c, 0x5e,
    0x60, 0x62, 0x64, 0x66, 0x68, 0x6a, 0x6c, 0x6e, 0x70, 0x72, 0x74, 0x76, 0x78, 0x7a, 0x7c, 0x7e,
    0x80, 0x82, 0x84, 0x86, 0x88, 0x8a, 0x8c, 0x8e, 0x90, 0x92, 0x94, 0x96, 0x98, 0x9a, 0x9c, 0x9e,
    0xa0, 0xa2, 0xa4, 0xa6, 0xa8, 0xaa, 0xac, 0xae, 0xb0, 0xb2, 0xb4, 0xb6, 0xb8, 0xba, 0xbc, 0xbe,
    0xc0, 0xc2, 0xc4, 0xc6, 0xc8, 0xca, 0xcc, 0xce, 0xd0, 0xd2, 0xd4, 0xd6, 0xd8, 0xda, 0xdc, 0xde,
    0xe0, 0xe2, 0xe4, 0xe6, 0xe8, 0xea, 0xec, 0xee, 0xf0, 0xf2, 0xf4, 0xf6, 0xf8, 0xfa, 0xfc, 0xfe,
    0x1b, 0x19, 0x1f, 0x1d, 0x13, 0x11, 0x17, 0x15, 0x0b, 0x09, 0x0f, 0x0d, 0x03, 0x01, 0x07, 0x05,
    0x3b, 0x39, 0x3f, 0x3d, 0x33, 0x31, 0x37, 0x35, 0x2b, 0x29, 0x2f, 0x2d, 0x23, 0x21, 0x27, 0x25,
    0x5b, 0x59, 0x5f, 0x5d, 0x53, 0x51, 0x57, 0x55, 0x4b, 0x49, 0x4f, 0x4d, 0x43, 0x41, 0x47, 0x45,
    0x7b, 0x79, 0x7f, 0x7d, 0x73, 0x71, 0x77, 0x75, 0x6b, 0x69, 0x6f, 0x6d, 0x63, 0x61, 0x67, 0x65,
    0x9b, 0x99, 0x9f, 0x9d, 0x93, 0x91, 0x97, 0x95, 0x8b, 0x89, 0x8f, 0x8d, 0x83, 0x81, 0x87, 0x85,
    0xbb, 0xb9, 0xbf, 0xbd, 0xb3, 0xb1, 0xb7, 0xb5, 0xab, 0xa9, 0xaf, 0xad, 0xa3, 0xa1, 0xa7, 0xa5,
    0xdb, 0xd9, 0xdf, 0xdd, 0xd3, 0xd1, 0xd7, 0xd5, 0xcb, 0xc9, 0xcf, 0xcd, 0xc3, 0xc1, 0xc7, 0xc5,
    0xfb, 0xf9, 0xff, 0xfd, 0xf3, 0xf1, 0xf7, 0xf5, 0xeb, 0xe9, 0xef, 0xed, 0xe3, 0xe1, 0xe7, 0xe5,
];
/// Lookup table for the Galois field GF(2^8) multiplication by 3 used in mix columns step.
const MULTIPLE_3: [u8; 256] = [
    0x00, 0x03, 0x06, 0x05, 0x0c, 0x0f, 0x0a, 0x09, 0x18, 0x1b, 0x1e, 0x1d, 0x14, 0x17, 0x12, 0x11,
    0x30, 0x33, 0x36, 0x35, 0x3c, 0x3f, 0x3a, 0x39, 0x28, 0x2b, 0x2e, 0x2d, 0x24, 0x27, 0x22, 0x21,
    0x60, 0x63, 0x66, 0x65, 0x6c, 0x6f, 0x6a, 0x69, 0x78, 0x7b, 0x7e, 0x7d, 0x74, 0x77, 0x72, 0x71,
    0x50, 0x53, 0x56, 0x55, 0x5c, 0x5f, 0x5a, 0x59, 0x48, 0x4b, 0x4e, 0x4d, 0x44, 0x47, 0x42, 0x41,
    0xc0, 0xc3, 0xc6, 0xc5, 0xcc, 0xcf, 0xca, 0xc9, 0xd8, 0xdb, 0xde, 0xdd, 0xd4, 0xd7, 0xd2, 0xd1,
    0xf0, 0xf3, 0xf6, 0xf5, 0xfc, 0xff, 0xfa, 0xf9, 0xe8, 0xeb, 0xee, 0xed, 0xe4, 0xe7, 0xe2, 0xe1,
    0xa0, 0xa3, 0xa6, 0xa5, 0xac, 0xaf, 0xaa, 0xa9, 0xb8, 0xbb, 0xbe, 0xbd, 0xb4, 0xb7, 0xb2, 0xb1,
    0x90, 0x93, 0x96, 0x95, 0x9c, 0x9f, 0x9a, 0x99, 0x88, 0x8b, 0x8e, 0x8d, 0x84, 0x87, 0x82, 0x81,
    0x9b, 0x98, 0x9d, 0x9e, 0x97, 0x94, 0x91, 0x92, 0x83, 0x80, 0x85, 0x86, 0x8f, 0x8c, 0x89, 0x8a,
    0xab, 0xa8, 0xad, 0xae, 0xa7, 0xa4, 0xa1, 0xa2, 0xb3, 0xb0, 0xb5, 0xb6, 0xbf, 0xbc, 0xb9, 0xba,
    0xfb, 0xf8, 0xfd, 0xfe, 0xf7, 0xf4, 0xf1, 0xf2, 0xe3, 0xe0, 0xe5, 0xe6, 0xef, 0xec, 0xe9, 0xea,
    0xcb, 0xc8, 0xcd, 0xce, 0xc7, 0xc4, 0xc1, 0xc2, 0xd3, 0xd0, 0xd5, 0xd6, 0xdf, 0xdc, 0xd9, 0xda,
    0x5b, 0x58, 0x5d, 0x5e, 0x57, 0x54, 0x51, 0x52, 0x43, 0x40, 0x45, 0x46, 0x4f, 0x4c, 0x49, 0x4a,
    0x6b, 0x68, 0x6d, 0x6e, 0x67, 0x64, 0x61, 0x62, 0x73, 0x70, 0x75, 0x76, 0x7f, 0x7c, 0x79, 0x7a,
    0x3b, 0x38, 0x3d, 0x3e, 0x37, 0x34, 0x31, 0x32, 0x23, 0x20, 0x25, 0x26, 0x2f, 0x2c, 0x29, 0x2a,
    0x0b, 0x08, 0x0d, 0x0e, 0x07, 0x04, 0x01, 0x02, 0x13, 0x10, 0x15, 0x16, 0x1f, 0x1c, 0x19, 0x1a,
];

/// Encrypt a single block of data using an expanded key.
///
/// This function takes the expanded key, and not the original key,
///  because it is used to implement [crate::encrypt] for each [crate::ModeOfOperation].
/// In order not to re-expand the key for every block, the expanded key is used instead.
/// # Warning
/// This function does not enforce the block size, and as such, passing a block of size other than 16 bytes produces undefined behavior.
///
/// Ideally, you should only use this to implement your own mode of operation, for most use cases, you should consider using the higher-level [crate::encrypt] instead.
pub fn encrypt_block(block: &[u8], expanded_key: &ExpandedKey) -> [u8; 16] {
    use crate::add_round_key;
    let mut cypher_block = [0; 16];
    for i in 0..16 {
        cypher_block[i] = block[i];
    }
    let expanded_key_content = expanded_key.content();
    let rounds_num = expanded_key.round_num();
    // First round
    add_round_key(&mut cypher_block, expanded_key.original_key());
    for i in 0..rounds_num - 1 {
        sub_bytes(&mut cypher_block);
        shift_rows(&mut cypher_block);
        mix_columns(&mut cypher_block);
        add_round_key(
            &mut cypher_block,
            &expanded_key_content[(i + 1) * 16..(i + 2) * 16], // Skip the first 16 bytes, they're the same as the original key.
        );
    }
    // Final round
    sub_bytes(&mut cypher_block);
    shift_rows(&mut cypher_block);
    add_round_key(
        &mut cypher_block,
        &expanded_key_content[expanded_key_content.len() - 16..],
    ); // Add the final key.
    cypher_block
}
/// Substitute each byte in the state with another byte according to [RIJNDAEL_S_BOX].
fn sub_bytes(state: &mut [u8; 16]) {
    for i in 0..16 {
        let byte: usize = state[i].into();
        state[i] = RIJNDAEL_S_BOX[byte];
    }
}
/// Applies a left circular byte shift for each row of the state by it's index.
/// You can imagine the state as a 4x4 matrix.
/// The first row is unchanged, the second row is shifted by 1 byte, etc.
fn shift_rows(state: &mut [u8; 16]) {
    // Since the shift is pretty much consistent, we can just hardcode every case instead of using a loop.
    let orig_state = state.clone();
    // First column
    state[1] = orig_state[5];
    state[2] = orig_state[10];
    state[3] = orig_state[15];
    // Second column
    state[5] = orig_state[9];
    state[6] = orig_state[14];
    state[7] = orig_state[3];
    // Third column
    state[9] = orig_state[13];
    state[10] = orig_state[2];
    state[11] = orig_state[7];
    // Forth column
    state[13] = orig_state[1];
    state[14] = orig_state[6];
    state[15] = orig_state[11];
}

/// Performs a Galois field GF(2^8) matrix multiplication of each byte from the state and a circulant MDS matrix ( 2 3 1 1; 1 2 3 1; 1 1 2 3; 3 1 1 2 ).
/// This involves a dot product of each column and row, a sum (XOR) of two byte products reduced by the Galois field x^8 + x^4 + x^3 + x + 1.
/// Since the MDS matrix has only 2 and 3 as non-neutral elements, we can precompute the results inside two lookup tables instead of re-calculating the results each time.
fn mix_columns(state: &mut [u8; 16]) {
    let orig_state = state.clone();
    // First column
    state[0] = MULTIPLE_2[orig_state[0] as usize]
        ^ MULTIPLE_3[orig_state[1] as usize]
        ^ orig_state[2]
        ^ orig_state[3];
    state[1] = orig_state[0]
        ^ MULTIPLE_2[orig_state[1] as usize]
        ^ MULTIPLE_3[orig_state[2] as usize]
        ^ orig_state[3];
    state[2] = orig_state[0]
        ^ orig_state[1]
        ^ MULTIPLE_2[orig_state[2] as usize]
        ^ MULTIPLE_3[orig_state[3] as usize];
    state[3] = MULTIPLE_3[orig_state[0] as usize]
        ^ orig_state[1]
        ^ orig_state[2]
        ^ MULTIPLE_2[orig_state[3] as usize];
    // Second column
    state[4] = MULTIPLE_2[orig_state[4] as usize]
        ^ MULTIPLE_3[orig_state[5] as usize]
        ^ orig_state[6]
        ^ orig_state[7];
    state[5] = orig_state[4]
        ^ MULTIPLE_2[orig_state[5] as usize]
        ^ MULTIPLE_3[orig_state[6] as usize]
        ^ orig_state[7];
    state[6] = orig_state[4]
        ^ orig_state[5]
        ^ MULTIPLE_2[orig_state[6] as usize]
        ^ MULTIPLE_3[orig_state[7] as usize];
    state[7] = MULTIPLE_3[orig_state[4] as usize]
        ^ orig_state[5]
        ^ orig_state[6]
        ^ MULTIPLE_2[orig_state[7] as usize];

    // Third column
    state[8] = MULTIPLE_2[orig_state[8] as usize]
        ^ MULTIPLE_3[orig_state[9] as usize]
        ^ orig_state[10]
        ^ orig_state[11];
    state[9] = orig_state[8]
        ^ MULTIPLE_2[orig_state[9] as usize]
        ^ MULTIPLE_3[orig_state[10] as usize]
        ^ orig_state[11];
    state[10] = orig_state[8]
        ^ orig_state[9]
        ^ MULTIPLE_2[orig_state[10] as usize]
        ^ MULTIPLE_3[orig_state[11] as usize];
    state[11] = MULTIPLE_3[orig_state[8] as usize]
        ^ orig_state[9]
        ^ orig_state[10]
        ^ MULTIPLE_2[orig_state[11] as usize];

    // Forth column
    state[12] = MULTIPLE_2[orig_state[12] as usize]
        ^ MULTIPLE_3[orig_state[13] as usize]
        ^ orig_state[14]
        ^ orig_state[15];
    state[13] = orig_state[12]
        ^ MULTIPLE_2[orig_state[13] as usize]
        ^ MULTIPLE_3[orig_state[14] as usize]
        ^ orig_state[15];
    state[14] = orig_state[12]
        ^ orig_state[13]
        ^ MULTIPLE_2[orig_state[14] as usize]
        ^ MULTIPLE_3[orig_state[15] as usize];
    state[15] = MULTIPLE_3[orig_state[12] as usize]
        ^ orig_state[13]
        ^ orig_state[14]
        ^ MULTIPLE_2[orig_state[15] as usize];
}

#[cfg(test)]
mod tests {

    use super::*;
    use crate::{CipherVersion, ModeOfOperation};
    #[test]
    fn sub_bytes_test() {
        let mut mock_state: [u8; 16] = [
            0x54, 0x68, 0x69, 0x73, 0x20, 0x6C, 0x69, 0x62, 0x20, 0x69, 0x73, 0x20, 0x63, 0x6F,
            0x6F, 0x6C,
        ]; // This lib is cool
        let substitution_expected: [u8; 16] = [
            0x20, 0x45, 0xf9, 0x8f, 0xb7, 0x50, 0xf9, 0xaa, 0xb7, 0xf9, 0x8f, 0xb7, 0xfb, 0xa8,
            0xa8, 0x50,
        ];
        sub_bytes(&mut mock_state);
        assert!(mock_state.iter().eq(&substitution_expected));
    }
    #[test]
    fn shift_rows_test() {
        let mut mock_state: [u8; 16] = [
            0x20, 0x45, 0xf9, 0x8f, 0xb7, 0x50, 0xf9, 0xaa, 0xb7, 0xf9, 0x8f, 0xb7, 0xfb, 0xa8,
            0xa8, 0x50,
        ];
        let shift_expected: [u8; 16] = [
            0x20, 0x50, 0x8f, 0x50, 0xb7, 0xf9, 0xa8, 0x8f, 0xb7, 0xa8, 0xf9, 0xaa, 0xfb, 0x45,
            0xf9, 0xb7,
        ];
        shift_rows(&mut mock_state);
        assert!(shift_expected.iter().eq(&mock_state));
    }
    #[test]
    fn mix_columns_test() {
        let mut mock_state: [u8; 16] = [
            0x20, 0x50, 0x8f, 0x50, 0xb7, 0xf9, 0xa8, 0x8f, 0xb7, 0xa8, 0xf9, 0xaa, 0xfb, 0x45,
            0xf9, 0xb7,
        ];
        let mix_expected: [u8; 16] = [
            0x6f, 0x5a, 0x85, 0x1f, 0x42, 0x32, 0x8f, 0x96, 0xc5, 0x46, 0x13, 0xdc, 0x6c, 0xd6,
            0x95, 0xdf,
        ];
        mix_columns(&mut mock_state);
        assert!(mix_expected.iter().eq(&mock_state));
    }
    #[test]
    fn encrypt_block_aes128_test() {
        let key = [
            0x54, 0x68, 0x69, 0x73, 0x20, 0x6C, 0x69, 0x62, 0x20, 0x69, 0x73, 0x20, 0x63, 0x6F,
            0x6F, 0x6C,
        ]; // This lib is cool
        let cipher_version = CipherVersion::Aes128(&key[..], ModeOfOperation::ECB);

        let expanded_key = ExpandedKey::new(&cipher_version);
        let mock_block = [
            0x54, 0x48, 0x69, 0x73, 0x20, 0x61, 0x6E, 0x20, 0x65, 0x78, 0x61, 0x6D, 0x70, 0x6C,
            0x65, 0x31,
        ]; // THis an example1
        let encrypt_expected = [
            0x75, 0x92, 0xce, 0xdb, 0x29, 0xf9, 0x3a, 0x1f, 0x9b, 0x98, 0xcb, 0x08, 0x86, 0xd5,
            0x82, 0x01,
        ];
        let encrypt_actual = encrypt_block(&mock_block, &expanded_key.unwrap());
        assert!(encrypt_expected.iter().eq(&encrypt_actual));
    }
    #[test]
    /// FIPS test vector.
    fn encrypt_block_aes192_test() {
        let key = [
            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
            0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
        ];
        let cipher_version = CipherVersion::Aes192(&key[..], ModeOfOperation::ECB);

        let expanded_key = ExpandedKey::new(&cipher_version);
        let mock_block = [
            0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd,
            0xee, 0xff,
        ];
        let encrypt_expected = [
            0xdd, 0xa9, 0x7c, 0xa4, 0x86, 0x4c, 0xdf, 0xe0, 0x6e, 0xaf, 0x70, 0xa0, 0xec, 0x0d,
            0x71, 0x91,
        ];
        let encrypt_actual = encrypt_block(&mock_block, &expanded_key.unwrap());
        assert!(encrypt_expected.iter().eq(&encrypt_actual));
    }
}