From c936c9d82d52bc5fca195a52f1ee512aa665ed9a Mon Sep 17 00:00:00 2001 From: 0xphen Date: Sun, 10 Dec 2023 23:20:28 +0100 Subject: [PATCH] Implement AesOps encryption --- aes/src/aes_ops.rs | 179 +++++++++++++++++++++++++++++++++++------ aes/src/block_modes.rs | 37 ++++++++- aes/src/constants.rs | 26 ++++++ aes/src/definitions.rs | 3 +- aes/src/error.rs | 3 + 5 files changed, 222 insertions(+), 26 deletions(-) diff --git a/aes/src/aes_ops.rs b/aes/src/aes_ops.rs index dd02a70..a2586d2 100644 --- a/aes/src/aes_ops.rs +++ b/aes/src/aes_ops.rs @@ -1,5 +1,7 @@ use super::{ - constants::{AES_S_BOX, TRANSFORMATION_MATRIX}, + constants::{ + AES_INVERSE_S_BOX, AES_S_BOX, INVERSE_TRANSFORMATION_MATRIX, TRANSFORMATION_MATRIX, + }, key_schedule::KeySchedule, util::{galois_mul, xor_matrices}, }; @@ -38,18 +40,35 @@ impl AesOps { // Main encryption rounds for round in 1..(rounds) { - Self::sub_bytes(state); + Self::sub_bytes(state, AES_S_BOX); Self::shift_rows(state); - Self::mix_columns(state); + Self::mix_columns(state, TRANSFORMATION_MATRIX); Self::add_round_key(state, keys.round_key(round as usize)); } //Final round without mixing columns - Self::sub_bytes(state); + Self::sub_bytes(state, AES_S_BOX); Self::shift_rows(state); Self::add_round_key(state, keys.round_key(rounds as usize)); } + pub fn decrypt(cipher_bytes: &mut [[u8; 4]; 4], keys: &KeySchedule) { + let rounds = keys.rounds; + + Self::add_round_key(cipher_bytes, keys.round_key(rounds as usize)); + + for round in (1..(rounds)).rev() { + Self::inv_shift_rows(cipher_bytes); + Self::sub_bytes(cipher_bytes, AES_INVERSE_S_BOX); + Self::add_round_key(cipher_bytes, keys.round_key(round as usize)); + Self::mix_columns(cipher_bytes, INVERSE_TRANSFORMATION_MATRIX); + } + + Self::inv_shift_rows(cipher_bytes); + Self::sub_bytes(cipher_bytes, AES_INVERSE_S_BOX); + Self::add_round_key(cipher_bytes, keys.round_key(0)); + } + /// Performs the AddRoundKey step, a crucial part of the AES encryption algorithm. /// /// This method XORs each byte of the AES state with the corresponding byte of the given round key. @@ -60,16 +79,23 @@ impl AesOps { *state = xor_matrices(*state, key); } - /// Performs the SubBytes transformation on the AES state. - /// This is a non-linear byte substitution step where each byte is replaced - /// with another according to a lookup table (S-box). - /// It mutates the current AES state by updating each byte with its substituted value. - fn sub_bytes(state: &mut [[u8; 4]; 4]) { + /// Performs the SubBytes or InvSubBytes transformation on the AES state. + /// + /// This function executes a non-linear byte substitution step where each byte + /// in the state is replaced with another according to the provided lookup table (S-box). + /// This lookup table can be either the standard S-box for encryption or the inverse S-box + /// for decryption, allowing this function to be used for both SubBytes in encryption + /// and InvSubBytes in decryption. + /// + /// # Arguments + /// * `state` - A mutable reference to the 4x4 state matrix. + /// * `s_box` - The S-box used for the transformation, either standard or inverse. + fn sub_bytes(state: &mut [[u8; 4]; 4], s_box: [u8; 256]) { // Iterate over each byte of the state matrix for (i, row) in state.iter_mut().enumerate() { for (j, e) in row.iter_mut().enumerate() { // Apply the S-box transformation and store in `new_state` - *e = AES_S_BOX[*e as usize]; + *e = s_box[*e as usize]; } } } @@ -98,24 +124,60 @@ impl AesOps { } } - /// Performs the MixColumns transformation on the AES state. + /// Performs the inverse shift rows step in AES decryption. /// - /// This function applies the MixColumns step to each column of the AES state matrix. - /// It uses the Galois Field multiplication (`galois_mul`) for the transformation. + /// Each row of the state is shifted cyclically to the right. + /// The first row remains unchanged, the second row is shifted by one position, + /// the third row by two positions, and the fourth row by three positions. /// /// # Arguments - /// * `&mut self` - A mutable reference to the AES structure, containing the state matrix. - fn mix_columns(state: &mut [[u8; 4]; 4]) { + /// * `state` - A mutable reference to the 4x4 state matrix. + fn inv_shift_rows(state: &mut [[u8; 4]; 4]) { + let mut temp: [u8; 4]; + + // Second row: shift one position to the right + temp = [state[3][1], state[0][1], state[1][1], state[2][1]]; + for j in 0..4 { + state[j][1] = temp[j]; + } + + // Third row: shift two positions to the right + temp = [state[2][2], state[3][2], state[0][2], state[1][2]]; + for j in 0..4 { + state[j][2] = temp[j]; + } + + // Fourth row: shift three positions to the right + temp = [state[1][3], state[2][3], state[3][3], state[0][3]]; + for j in 0..4 { + state[j][3] = temp[j]; + } + } + + /// Performs the MixColumns or InvMixColumns transformation on the AES state. + /// + /// This function applies either the MixColumns or InvMixColumns step to each column + /// of the AES state matrix, depending on the provided transformation matrix. + /// It uses Galois Field multiplication (`galois_mul`) for the transformation. + /// + /// The transformation matrix should be the standard matrix for MixColumns in AES encryption, + /// or the inverse matrix for InvMixColumns in AES decryption. + /// + /// # Arguments + /// * `state` - A mutable reference to the 4x4 state matrix. + /// * `transformation_matrix` - The matrix used for the transformation, either for + /// MixColumns or InvMixColumns. + fn mix_columns(state: &mut [[u8; 4]; 4], transformation_matrix: [[u8; 4]; 4]) { for col in 0..4 { // Temporary storage for the column being processed let mut temp_column = [0u8; 4]; // Transform the current column using Galois Field multiplication for i in 0..4 { - temp_column[i] = galois_mul(TRANSFORMATION_MATRIX[i][0], state[col][0]) - ^ galois_mul(TRANSFORMATION_MATRIX[i][1], state[col][1]) - ^ galois_mul(TRANSFORMATION_MATRIX[i][2], state[col][2]) - ^ galois_mul(TRANSFORMATION_MATRIX[i][3], state[col][3]); + temp_column[i] = galois_mul(transformation_matrix[i][0], state[col][0]) + ^ galois_mul(transformation_matrix[i][1], state[col][1]) + ^ galois_mul(transformation_matrix[i][2], state[col][2]) + ^ galois_mul(transformation_matrix[i][3], state[col][3]); } // Update the state matrix with the transformed column @@ -131,7 +193,7 @@ mod tests { use super::*; #[test] - fn aes_ops_encrypt_test() { + fn aes_ops_encrypt_decrypt_test() { let mut state: [[u8; 4]; 4] = [ [0, 17, 34, 51], [68, 85, 102, 119], @@ -143,7 +205,6 @@ mod tests { KeySchedule::new(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]).unwrap(); AesOps::encrypt(&mut state, &key_schedule); - assert_eq!( state, [ @@ -153,10 +214,21 @@ mod tests { [112, 180, 197, 90] ] ); + + AesOps::decrypt(&mut state, &key_schedule); + assert_eq!( + state, + [ + [0, 17, 34, 51], + [68, 85, 102, 119], + [136, 153, 170, 187], + [204, 221, 238, 255] + ] + ); } #[test] - fn initial_round_key_and_one_round_test() { + fn one_round_encryption_test() { let mut state: [[u8; 4]; 4] = [ [0, 17, 34, 51], [68, 85, 102, 119], @@ -178,7 +250,7 @@ mod tests { ] ); - AesOps::sub_bytes(&mut state); + AesOps::sub_bytes(&mut state, AES_S_BOX); assert_eq!( state, [ @@ -200,7 +272,7 @@ mod tests { ] ); - AesOps::mix_columns(&mut state); + AesOps::mix_columns(&mut state, TRANSFORMATION_MATRIX); assert_eq!( state, [ @@ -211,4 +283,63 @@ mod tests { ] ); } + + #[test] + fn one_round_decryption_test() { + let key_schedule = + KeySchedule::new(&[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]).unwrap(); + + let mut state: [[u8; 4]; 4] = [ + [105, 196, 224, 216], + [106, 123, 4, 48], + [216, 205, 183, 128], + [112, 180, 197, 90], + ]; + + AesOps::add_round_key(&mut state, key_schedule.round_key(10)); + assert_eq!( + state, + [ + [122, 213, 253, 167], + [137, 239, 78, 39], + [43, 202, 16, 11], + [61, 159, 245, 159] + ] + ); + + AesOps::inv_shift_rows(&mut state); + assert_eq!( + state, + [ + [122, 159, 16, 39], + [137, 213, 245, 11], + [43, 239, 253, 159], + [61, 202, 78, 167] + ] + ); + + AesOps::sub_bytes(&mut state, AES_INVERSE_S_BOX); + assert_eq!( + state, + [ + [189, 110, 124, 61], + [242, 181, 119, 158], + [11, 97, 33, 110], + [139, 16, 182, 137] + ] + ); + + AesOps::mix_columns(&mut state, TRANSFORMATION_MATRIX); + + AesOps::mix_columns(&mut state, INVERSE_TRANSFORMATION_MATRIX); + assert_eq!( + state, + [ + [189, 110, 124, 61], + [242, 181, 119, 158], + [11, 97, 33, 110], + [139, 16, 182, 137] + ] + ); + } } diff --git a/aes/src/block_modes.rs b/aes/src/block_modes.rs index 7fd5fe2..3635276 100644 --- a/aes/src/block_modes.rs +++ b/aes/src/block_modes.rs @@ -82,13 +82,39 @@ impl<'k> AesEncryptor for CbcEncryptor<'k> { let mut encrypted_blocks = Vec::with_capacity(input_blocks.len()); for block in input_blocks { - AesOps::encrypt(&mut working_state, &self.keys); + AesOps::encrypt(&mut working_state, self.keys); encrypted_blocks.push(working_state); working_state = xor_matrices(working_state, block); } Ok(encrypted_blocks) } + + fn decrypt(&mut self, cipher_bytes: &[u8]) -> Result, AesError> { + if cipher_bytes.len() % 16 != 0 { + return Err(AesError::InvalidCipherText); + } + + let input_blocks = chunk_bytes_into_4x4_matrices(&cipher_bytes.to_vec()); + let mut decrypted_blocks: Vec<[[u8; 4]; 4]> = Vec::with_capacity(input_blocks.len()); + + let mut working_block = self.iv; + + for block in input_blocks { + let mut cipher_block = block; + AesOps::decrypt(&mut cipher_block, self.keys); + + cipher_block = xor_matrices(cipher_block, working_block); + decrypted_blocks.push(cipher_block); + working_block = block; + } + + Ok(decrypted_blocks + .into_iter() + .flat_map(|block| block.into_iter()) + .flat_map(|row| row.into_iter()) + .collect()) + } } #[cfg(test)] @@ -119,6 +145,15 @@ mod tests { ]]; let result = cbc_ops.encrypt(&INPUT).unwrap(); + println!("result: {:?}", result); assert!(result.as_slice().starts_with(&start_cipher_bytes)); + + let plain_bytes = cbc_ops + .decrypt(&[ + 59, 67, 136, 134, 79, 78, 189, 114, 137, 150, 207, 148, 186, 117, 130, 178, 17, + 210, 7, 174, 109, 178, 129, 201, 24, 52, 14, 108, 136, 148, 142, 63, + ]) + .unwrap(); + println!("plain_bytes: {:?}", plain_bytes); } } diff --git a/aes/src/constants.rs b/aes/src/constants.rs index 3e27737..6511538 100644 --- a/aes/src/constants.rs +++ b/aes/src/constants.rs @@ -17,7 +17,33 @@ pub const AES_S_BOX: [u8; 256] = [ 0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16, ]; +pub const AES_INVERSE_S_BOX: [u8; 256] = [ + 0x52, 0x09, 0x6A, 0xD5, 0x30, 0x36, 0xA5, 0x38, 0xBF, 0x40, 0xA3, 0x9E, 0x81, 0xF3, 0xD7, 0xFB, + 0x7C, 0xE3, 0x39, 0x82, 0x9B, 0x2F, 0xFF, 0x87, 0x34, 0x8E, 0x43, 0x44, 0xC4, 0xDE, 0xE9, 0xCB, + 0x54, 0x7B, 0x94, 0x32, 0xA6, 0xC2, 0x23, 0x3D, 0xEE, 0x4C, 0x95, 0x0B, 0x42, 0xFA, 0xC3, 0x4E, + 0x08, 0x2E, 0xA1, 0x66, 0x28, 0xD9, 0x24, 0xB2, 0x76, 0x5B, 0xA2, 0x49, 0x6D, 0x8B, 0xD1, 0x25, + 0x72, 0xF8, 0xF6, 0x64, 0x86, 0x68, 0x98, 0x16, 0xD4, 0xA4, 0x5C, 0xCC, 0x5D, 0x65, 0xB6, 0x92, + 0x6C, 0x70, 0x48, 0x50, 0xFD, 0xED, 0xB9, 0xDA, 0x5E, 0x15, 0x46, 0x57, 0xA7, 0x8D, 0x9D, 0x84, + 0x90, 0xD8, 0xAB, 0x00, 0x8C, 0xBC, 0xD3, 0x0A, 0xF7, 0xE4, 0x58, 0x05, 0xB8, 0xB3, 0x45, 0x06, + 0xD0, 0x2C, 0x1E, 0x8F, 0xCA, 0x3F, 0x0F, 0x02, 0xC1, 0xAF, 0xBD, 0x03, 0x01, 0x13, 0x8A, 0x6B, + 0x3A, 0x91, 0x11, 0x41, 0x4F, 0x67, 0xDC, 0xEA, 0x97, 0xF2, 0xCF, 0xCE, 0xF0, 0xB4, 0xE6, 0x73, + 0x96, 0xAC, 0x74, 0x22, 0xE7, 0xAD, 0x35, 0x85, 0xE2, 0xF9, 0x37, 0xE8, 0x1C, 0x75, 0xDF, 0x6E, + 0x47, 0xF1, 0x1A, 0x71, 0x1D, 0x29, 0xC5, 0x89, 0x6F, 0xB7, 0x62, 0x0E, 0xAA, 0x18, 0xBE, 0x1B, + 0xFC, 0x56, 0x3E, 0x4B, 0xC6, 0xD2, 0x79, 0x20, 0x9A, 0xDB, 0xC0, 0xFE, 0x78, 0xCD, 0x5A, 0xF4, + 0x1F, 0xDD, 0xA8, 0x33, 0x88, 0x07, 0xC7, 0x31, 0xB1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xEC, 0x5F, + 0x60, 0x51, 0x7F, 0xA9, 0x19, 0xB5, 0x4A, 0x0D, 0x2D, 0xE5, 0x7A, 0x9F, 0x93, 0xC9, 0x9C, 0xEF, + 0xA0, 0xE0, 0x3B, 0x4D, 0xAE, 0x2A, 0xF5, 0xB0, 0xC8, 0xEB, 0xBB, 0x3C, 0x83, 0x53, 0x99, 0x61, + 0x17, 0x2B, 0x04, 0x7E, 0xBA, 0x77, 0xD6, 0x26, 0xE1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0C, 0x7D, +]; + pub const TRANSFORMATION_MATRIX: [[u8; 4]; 4] = [[2, 3, 1, 1], [1, 2, 3, 1], [1, 1, 2, 3], [3, 1, 1, 2]]; +pub const INVERSE_TRANSFORMATION_MATRIX: [[u8; 4]; 4] = [ + [14, 11, 13, 9], + [9, 14, 11, 13], + [13, 9, 14, 11], + [11, 13, 9, 14], +]; + pub const ROUND_CONSTANT_128: [u8; 10] = [1, 2, 4, 8, 16, 32, 64, 128, 27, 54]; diff --git a/aes/src/definitions.rs b/aes/src/definitions.rs index f335754..50b0758 100644 --- a/aes/src/definitions.rs +++ b/aes/src/definitions.rs @@ -1,7 +1,8 @@ use super::error::AesError; pub trait AesEncryptor { - fn encrypt(&mut self, message: &[u8]) -> Result, AesError>; + fn encrypt(&mut self, input: &[u8]) -> Result, AesError>; + fn decrypt(&mut self, cipher_bytes: &[u8]) -> Result, AesError>; } /// Trait for padding processing in cryptographic operations. diff --git a/aes/src/error.rs b/aes/src/error.rs index 6ebfa2c..f98edd3 100644 --- a/aes/src/error.rs +++ b/aes/src/error.rs @@ -16,4 +16,7 @@ pub enum AesError { #[error("Failed to parse slice to matrix: {0}")] FailedToParseSliceToMatrix(String), + + #[error("Invalid cipher text")] + InvalidCipherText, }