From 7de7e084290199dc1fe9d75577c642d48e75bb5d Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@inbox.lv>
Date: Fri, 26 Apr 2024 17:03:03 +0200
Subject: [PATCH] more test

---
 src/aes.rs | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/src/aes.rs b/src/aes.rs
index eaacc51..89d744c 100644
--- a/src/aes.rs
+++ b/src/aes.rs
@@ -160,7 +160,7 @@ fn xor(bits1: &[u8; 4], bits2: &[u8; 4]) -> [u8; 4] {
     result
 }
 
-fn substitute_bytes(bytes: &[u8; 4]) -> [u8; 4] {
+fn substitute_word(bytes: &[u8; 4]) -> [u8; 4] {
     let mut result = [0u8; 4];
     for i in 0..4 {
         result[i] = substitute(bytes[i], true);
@@ -194,7 +194,7 @@ fn substitute(byte: u8, encryption: bool) -> u8 {
     }
 }
 
-fn shift_bytes(bytes: &[u8; 4]) -> [u8; 4] {
+fn shift_word(bytes: &[u8; 4]) -> [u8; 4] {
     let mut result = [0u8; 4];
     for i in 0..4 {
         result[i] = bytes[(i + 1) % 4];
@@ -268,7 +268,7 @@ impl Aes {
                 expanded_key[i] = xor(
                     &xor(
                         &expanded_key[i - n],
-                        &substitute_bytes(&shift_bytes(&expanded_key[i - 1])),
+                        &substitute_word(&shift_word(&expanded_key[i - 1])),
                     ),
                     &rcon,
                 );
@@ -285,6 +285,29 @@ impl Aes {
 mod tests {
     use super::*;
     #[test]
+    fn shift_word_test() {
+        let test: u32 = 0x09cf4f3c;
+        let test_bytes: [u8; 4] = test.to_be_bytes();
+        let expected: u32 = 0xcf4f3c09;
+        let expected_bytes: [u8; 4] = expected.to_be_bytes();
+        assert_eq!(shift_word(&test_bytes), expected_bytes);
+    }
+    #[test]
+    fn substitute_word_test() {
+        let test: u32 = 0xcf4f3c09;
+        let test_bytes: [u8; 4] = test.to_be_bytes();
+        let expected: u32 = 0x8a84eb01;
+        let expected_bytes: [u8; 4] = expected.to_be_bytes();
+        assert_eq!(substitute_word(&test_bytes), expected_bytes);
+    }
+    #[test]
+    fn xor_test() {
+        let w1: [u8; 4] = [0x00, 0x11, 0x10, 0x10];
+        let w2: [u8; 4] = [0x00, 0x11, 0x01, 0x10];
+        let expected: [u8; 4] = [0x00, 0x00, 0x11, 0x00];
+        assert_eq!(xor(&w1, &w2), expected);
+    }
+    #[test]
     fn key_schedule_test() {
         let key: [u8; 16] = [
         0x2b, 0x7e, 0x15, 0x16,