aes key schedule + test

src/aes.rs

@@ -1,5 +1,9 @@
 use std::io::{Error, ErrorKind};
 
+// AES specs
+// https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.197.pdf
+
+
 // S box, 16*16 Matrix with fixed values
 static SBOX: [[u8; 16]; 16] = [
     [
@@ -135,6 +139,11 @@ static INVERSE_SBOX: [[u8; 16]; 16] = [
     ],
 ];
 
+// for key schedule
+static RC: [u8; 11] = [
+    0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36,
+];
+
 // Mix Columns
 static MIXCOLUMNS: [[u8; 4]; 4] = [
     [0x02, 0x03, 0x01, 0x01],
@@ -143,22 +152,79 @@ static MIXCOLUMNS: [[u8; 4]; 4] = [
     [0x03, 0x01, 0x01, 0x02],
 ];
 
+fn xor(bits1: &[u8; 4], bits2: &[u8; 4]) -> [u8; 4] {
+    let mut result = [0u8; 4];
+    for i in 0..4 {
+        result[i] = bits1[i] ^ bits2[i];
+    }
+    result
+}
+
+fn substitute_bytes(bytes: &[u8; 4]) -> [u8; 4] {
+    let mut result = [0u8; 4];
+    for i in 0..4 {
+        result[i] = substitute(bytes[i], true);
+    }
+    return result;
+}
+fn substitute_state(state: &mut [[u8; 4]; 4]) {
+    for i in 0..4 {
+        for j in 0..4 {
+            state[i][j] = substitute(state[i][j], true);
+        }
+    }
+}
+fn inverse_substitute_state(state: &mut [[u8; 4]; 4]) {
+    for i in 0..4 {
+        for j in 0..4 {
+            state[i][j] = substitute(state[i][j], false);
+        }
+    }
+}
+
+fn substitute(byte: u8, encryption: bool) -> u8 {
+    let i: usize;
+    let j: usize;
+    i = ((byte >> 4) & 0xF).into();
+    j = (byte & 0xF).into();
+    if encryption == true {
+        SBOX[i][j]
+    } else {
+        INVERSE_SBOX[i][j]
+    }
+}
+
+fn shift_bytes(bytes: &[u8; 4]) -> [u8; 4] {
+    let mut result = [0u8; 4];
+    for i in 0..4 {
+        result[i] = bytes[(i + 1) % 4];
+    }
+    result
+}
+
+fn add_round_key(state: &mut [[u8; 4]; 4], key: &[[u8; 4]; 4]) {
+    for i in 0..4 {
+        for j in 0..4 {
+            state[i][j] = state[i][j] ^ key[j][i];
+        }
+    }
+}
+
 pub struct Aes {
     n_turn: usize,
-    n_key: usize,
-    key: Vec<u8>,
+    key: [u8; 16],
+    expanded_key: [[u8; 4]; 44],
     state: [[u8; 4]; 4],
 }
 
 impl Aes {
-    fn new(block: [[u8; 4]; 4], key: Vec<u8>, n_turn: usize) -> std::io::Result<Self> {
-        let n_key: usize = match key.len() {
-            128 => 128,
-            192 => 192,
-            256 => 256,
-            _ => return Err(Error::new(ErrorKind::Other, "unsupported key size")),
-        };
-        Ok(Aes {n_turn, n_key, key, state: block})
+    fn new(block: [[u8; 4]; 4], key: [u8; 16], n_turn: usize) -> Self {
+        Aes {
+            n_turn,
+            key,
+            expanded_key: Self::key_schedule(&key),
+            state: block,
+        }
     }
     fn sub_bytes(&mut self) {
         unimplemented!();
@@ -181,6 +247,64 @@ impl Aes {
     }
 
     fn encrypt_block(block: [[u8; 4]; 4], key: Vec<u8>, n_turn: usize) {
-        a = new(block, key, n_turn)
+        unimplemented!();
+    }
+    pub fn key_schedule(key_bytes: &[u8; 16]) -> [[u8; 4]; 44] {
+        let mut original_key = [[0u8; 4]; 4];
+        let mut expanded_key = [[0u8; 4]; 44];
+        let n = 4;
+
+        for i in 0..16 {
+            original_key[i / 4][i % 4] = key_bytes[i];
+        }
+
+        // AES128, 11 rounds, i ∈ [0;4*rounds-1]
+        for i in 0..44 {
+            if i < n {
+                expanded_key[i] = original_key[i];
+            } else if i >= n && i % n == 0 {
+                let mut rcon = [0u8; 4];
+                rcon[0] = RC[i / n];
+                expanded_key[i] = xor(
+                    &xor(
+                        &expanded_key[i - n],
+                        &substitute_bytes(&shift_bytes(&expanded_key[i - 1])),
+                    ),
+                    &rcon,
+                );
+            } else {
+                expanded_key[i] = xor(&expanded_key[i - n], &expanded_key[i - 1]);
+            }
+        }
+
+        expanded_key
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn key_schedule_test() {
+        let key: [u8; 16] = [
+        0x2b, 0x7e, 0x15, 0x16,
+        0x28, 0xae, 0xd2, 0xa6,
+        0xab, 0xf7, 0x15, 0x88,
+        0x09, 0xcf, 0x4f, 0x3c,
+        ];
+        let expanded_key = Aes::key_schedule(&key);
+        let expected: [u32; 44] = [
+            0x2b7e1516, 0x28aed2a6, 0xabf71588, 0x09cf4f3c, 0xa0fafe17, 0x88542cb1, 0x23a33939, 0x2a6c7605, 0xf2c295f2, 0x7a96b943, 0x5935807a, 0x7359f67f, 0x3d80477d, 0x4716fe3e, 0x1e237e44, 0x6d7a883b, 0xef44a541, 0xa8525b7f, 0xb671253b, 0xdb0bad00, 0xd4d1c6f8, 0x7c839d87, 0xcaf2b8bc, 0x11f915bc, 0x6d88a37a, 0x110b3efd, 0xdbf98641, 0xca0093fd, 0x4e54f70e, 0x5f5fc9f3, 0x84a64fb2, 0x4ea6dc4f, 0xead27321, 0xb58dbad2, 0x312bf560, 0x7f8d292f, 0xac7766f3, 0x19fadc21, 0x28d12941, 0x575c006e, 0xd014f9a8, 0xc9ee2589, 0xe13f0cc8, 0xb6630ca6,
+        ];
+        let mut expected_bytes: [u8; 4];
+        for i in 0..11 {
+            for j in 0..4 {
+                let index = i * 4 + j;
+                expected_bytes = expected[index].to_be_bytes();
+                for (&byte, &expected_byte) in expanded_key[index].iter().zip(expected_bytes.iter()) {
+                    assert_eq!(byte, expected_byte);
+                }
+            }
+        }
     }
 }