From 5ba752d95bd0c304d7335adadf0fd8188a860fd2 Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@inbox.lv>
Date: Sun, 28 Apr 2024 15:21:16 +0200
Subject: [PATCH] find key + test

---
 src/main.rs   |  5 +++++
 src/square.rs | 15 ++++++++++++---
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 7f409e9..70c7dd3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -41,4 +41,9 @@ fn main() {
         println!();
         println!();
     }
+    let found = Aes::findkey(&key);
+    for &byte in &found {
+            print!("{:02x}", byte);
+        }
+        println!();
 }
diff --git a/src/square.rs b/src/square.rs
index 7787acb..f6f6d34 100644
--- a/src/square.rs
+++ b/src/square.rs
@@ -131,9 +131,10 @@ impl Aes {
         }
         found_key
     }
-    // pub fn findkey(&key: &[u8; 16]) -> [u8; 16] {
-    //
-    // }
+    pub fn findkey(&key: &[u8; 16]) -> [u8; 16] {
+        let roundkey: [u8; 16] = Self::findroundkey(&key);
+        Self::reverse_key_schedule(&roundkey, 4)
+    }
 }
 
 #[cfg(test)]
@@ -176,4 +177,12 @@ mod tests {
             assert_eq!(Aes::reverse_key_schedule(&round_keys[i], i), key);
         }
     }
+    #[test]
+    fn find_key_test() {
+        let key: [u8; 16] = [
+            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
+            0x4f, 0x3c,
+        ];
+        assert_eq!(key, Aes::findkey(&key));
+    }
 }