Merge pull request #6 from 0xphen/feat/aes-decryption

Implement AesOps decryption

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]
+            ]
+        );
+    }
 }

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<Vec<u8>, 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);
     }
 }

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];

aes/src/definitions.rs

@@ -1,7 +1,8 @@
 use super::error::AesError;
 
 pub trait AesEncryptor {
-    fn encrypt(&mut self, message: &[u8]) -> Result<Vec<[[u8; 4]; 4]>, AesError>;
+    fn encrypt(&mut self, input: &[u8]) -> Result<Vec<[[u8; 4]; 4]>, AesError>;
+    fn decrypt(&mut self, cipher_bytes: &[u8]) -> Result<Vec<u8>, AesError>;
 }
 
 /// Trait for padding processing in cryptographic operations.

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,
 }