From 25d3f56120a5231a70675f812e6cdc1f3d9e462d Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@inbox.lv>
Date: Sun, 28 Apr 2024 16:21:51 +0200
Subject: [PATCH] square attack against n random keys

---
 Cargo.toml    |  1 +
 src/aes.rs    |  2 --
 src/main.rs   | 40 ++--------------------------------------
 src/square.rs | 43 +++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 46 insertions(+), 40 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 2a71630..8e5ecf4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,3 +6,4 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+rand = "0.8.5"
diff --git a/src/aes.rs b/src/aes.rs
index 8430cc3..459bb70 100644
--- a/src/aes.rs
+++ b/src/aes.rs
@@ -234,7 +234,6 @@ pub fn add_round_key(state: &mut [[u8; 4]; 4], key: &[[u8; 4]; 4]) {
 
 pub struct Aes {
     n_turn: usize,
-    key: [u8; 16],
     pub expanded_key: [[u8; 4]; 44],
 }
 
@@ -242,7 +241,6 @@ impl Aes {
     pub fn new(&key: &[u8; 16], &n_turn: &usize) -> Self {
         Aes {
             n_turn,
-            key,
             expanded_key: Self::key_schedule(&key),
         }
     }
diff --git a/src/main.rs b/src/main.rs
index 70c7dd3..fad6ca5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,43 +7,7 @@ fn main() {
         0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f,
         0x3c,
     ];
-    let found_key: [u8; 16] = Aes::findroundkey(&key);
-    for &byte in &found_key {
-        print!("{:02x}", byte);
-    }
-    println!();
-    let found_key_origininal = Aes::reverse_key_schedule(&found_key, 4);
-    for &byte in &found_key_origininal {
-        print!("{:02x}", byte);
-    }
-    println!();
-
-    let nturn = 4;
+    let nturn = 11;
     let aescipher = Aes::new(&key, &nturn);
-    let mut output = [[0u8; 16]; 11];
-    for i in 0..11 {
-        for j in 0..4 {
-            output[i][4 * j..4 * (j + 1)].copy_from_slice(&aescipher.expanded_key[4 * i + j]);
-        }
-    }
-    for &key in &output {
-        for &byte in &key {
-            print!("{:02x}", byte);
-        }
-        println!();
-    }
-
-    println!("test");
-    for i in 0..11 {
-        for &byte in &Aes::reverse_key_schedule(&output[i], i) {
-            print!("{:02x}", byte);
-        }
-        println!();
-        println!();
-    }
-    let found = Aes::findkey(&key);
-    for &byte in &found {
-            print!("{:02x}", byte);
-        }
-        println!();
+    Aes::test_square_attack(5);
 }
diff --git a/src/square.rs b/src/square.rs
index f6f6d34..8d40c60 100644
--- a/src/square.rs
+++ b/src/square.rs
@@ -1,4 +1,5 @@
 use crate::aes::*;
+use rand::{Rng, thread_rng};
 
 impl Aes {
     pub fn reverse_key_schedule(original_round_key: &[u8; 16], aes_round: usize) -> [u8; 16] {
@@ -104,6 +105,7 @@ impl Aes {
         let mut more_key_guesses: [Vec<u8>; 16];
         let mut result: [Vec<u8>; 16] = Default::default();
         while key_guesses.iter().any(|vec| vec.len() > 1) {
+            println!("Collisions present, additional key search with a different constant.");
             const_byte += 1_u8;
             ciphertexts = Self::aes_reduced_gen_texts(&key, &const_byte);
             more_key_guesses = Self::guessroundkey(&ciphertexts);
@@ -131,10 +133,51 @@ impl Aes {
         }
         found_key
     }
+
     pub fn findkey(&key: &[u8; 16]) -> [u8; 16] {
         let roundkey: [u8; 16] = Self::findroundkey(&key);
         Self::reverse_key_schedule(&roundkey, 4)
     }
+
+    pub fn generate_random_keys(n: usize) -> Vec<[u8; 16]> {
+        let mut rng = thread_rng();
+        let mut keys_vec = Vec::with_capacity(n);
+        let mut key = [0u8; 16];
+
+        for _ in 0..n {
+            rng.fill(&mut key);
+            keys_vec.push(key);
+        }
+
+        keys_vec
+    }
+
+    pub fn test_square_attack(n: usize) {
+        println!("Square attack against 4-turns AES with {} randomly generated keys", n);
+        let keys: Vec<[u8; 16]> = Self::generate_random_keys(n);
+        let mut found_key: [u8; 16];
+        let mut success: usize = 0;
+        for &key in &keys {
+            print!("\ngenerated key: ");
+            Self::print_key(&key);
+            found_key = Aes::findkey(&key);
+            if found_key == key {
+                print!("found key: ");
+                Self::print_key(&found_key);
+                success += 1;
+            } else {
+                println!("key search failed");
+            }
+        }
+        println!("\n{}/{} key(s) found", success, n);
+    }
+
+    pub fn print_key(&key: &[u8; 16]) {
+        for &byte in &key {
+                    print!("{:02x}", byte);
+                }
+                println!();
+    }
 }
 
 #[cfg(test)]