command line clap, command: encrypt, findkey, square

Cargo.toml

@@ -6,4 +6,5 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+clap = { version = "4.5.4", features = ["derive"] }
 rand = "0.8.5"

src/aes.rs

@@ -1,4 +1,3 @@
-
 // AES specs
 // https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.197.pdf
 
@@ -280,6 +279,19 @@ impl Aes {
         result
     }
 
+    pub fn encrypt_string(&self, text: &String) -> [u8; 16] {
+        let mut block = [0u8; 16];
+        let bytes = text.as_bytes();
+        if bytes.len() > 16 {
+            println!("string too long, will be cut.")
+        }
+        let len = bytes.len().min(16);
+
+        block[..len].copy_from_slice(&bytes[..len]);
+
+        self.encrypt_block(&block)
+    }
+
     pub fn decrypt_block(&self, block: &[u8; 16]) -> [u8; 16] {
         let mut result = [0u8; 16];
 
@@ -316,6 +328,19 @@ impl Aes {
         result
     }
 
+    pub fn decrypt_string(&self, text: &String) -> [u8; 16] {
+        let mut block = [0u8; 16];
+        let bytes = text.as_bytes();
+        if bytes.len() > 16 {
+            println!("string too long, will be cut.")
+        }
+        let len = bytes.len().min(16);
+
+        block[..len].copy_from_slice(&bytes[..len]);
+
+        self.decrypt_block(&block)
+    }
+
     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];

src/main.rs

@@ -1,13 +1,117 @@
 mod aes;
 mod square;
+mod utils;
 use crate::aes::Aes;
 
-fn main() {
-    let key: [u8; 16] = [
-        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f,
-        0x3c,
-    ];
-    let nturn = 11;
-    let aescipher = Aes::new(&key, &nturn);
-    Aes::test_square_attack(5);
+use clap::{Parser, Subcommand};
+
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+#[command(propagate_version = true)]
+struct Cli {
+    #[command(subcommand)]
+    command: Commands,
+}
+
+#[derive(Subcommand)]
+enum Commands {
+    /// Ecnrypt
+    Encrypt {
+        /// (default Some('11'))
+        #[arg(long = "nturn", short = 'n')]
+        nturn: Option<String>,
+        /// (default "")
+        #[arg(long = "text", short = 't')]
+        text: String,
+        /// (default True)
+        #[arg(long = "hextext", short = 'T')]
+        hextext: bool,
+        /// (default "")
+        #[arg(long = "key", short = 'k')]
+        key: String,
+        /// (default True)
+        #[arg(long = "hexkey", short = 'K')]
+        hexkey: bool,
+        /// (default False)
+        #[arg(long = "decrypt", short = 'd')]
+        decrypt: bool,
+    },
+    Findkey {
+        /// (default "")
+        key: String,
+        /// (default True)
+        #[arg(long = "hexkey", short = 'K')]
+        hexkey: bool,
+    },
+    Square {
+        /// (default "1")
+        number: String,
+    },
+}
+
+fn main() -> std::io::Result<()> {
+    let cli = Cli::parse();
+
+    match &cli.command {
+        Commands::Encrypt {
+            nturn,
+            text,
+            hextext,
+            key,
+            hexkey,
+            decrypt,
+        } => {
+            println!("original text: {}", &text);
+            let aeskey: [u8; 16];
+            let result: [u8; 16];
+            let aestext: [u8; 16];
+            let mut aesnturn: usize = 10;
+
+            match hexkey {
+                false => aeskey = utils::parse_hex_string(&key),
+                true => aeskey = utils::string_to_u8_16(&key),
+            }
+            match nturn {
+                Some(number) => aesnturn = number.parse::<usize>().unwrap(),
+                _ => {}
+            }
+
+            let aescipher = Aes::new(&aeskey, &aesnturn);
+
+            match hextext {
+                true => aestext = utils::parse_hex_string(&text),
+                false => aestext = utils::string_to_u8_16(&text),
+            }
+
+            match decrypt {
+                true => {
+                    result = aescipher.decrypt_block(&aestext);
+                    print!("Decrypted text, hex-encoded: ");
+                }
+                false => {
+                    result = aescipher.encrypt_block(&aestext);
+                    print!("Encrypted text, hex-encoded: ");
+                }
+            }
+            Aes::print_key(&result);
+            println!();
+        }
+        Commands::Findkey { key, hexkey } => {
+            let aeskey: [u8; 16];
+            match hexkey {
+                false => aeskey = utils::parse_hex_string(&key),
+                true => aeskey = utils::string_to_u8_16(&key),
+            }
+            let keyfound: [u8; 16] = Aes::findkey(&aeskey);
+            if keyfound == aeskey {
+                println!("key found.");
+            } else {
+                println!("key not found.");
+            }
+        }
+        Commands::Square { number } => {
+            Aes::test_square_attack(number.parse::<usize>().unwrap());
+        }
+    }
+    Ok(())
 }

src/square.rs

@@ -1,5 +1,5 @@
 use crate::aes::*;
-use rand::{Rng, thread_rng};
+use rand::{thread_rng, Rng};
 
 impl Aes {
     pub fn reverse_key_schedule(original_round_key: &[u8; 16], aes_round: usize) -> [u8; 16] {
@@ -153,7 +153,10 @@ impl Aes {
     }
 
     pub fn test_square_attack(n: usize) {
-        println!("Square attack against 4-turns AES with {} randomly generated keys", n);
+        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;
@@ -174,9 +177,9 @@ impl Aes {
 
     pub fn print_key(&key: &[u8; 16]) {
         for &byte in &key {
-                    print!("{:02x}", byte);
-                }
-                println!();
+            print!("{:02x}", byte);
+        }
+        println!();
     }
 }
 

src/utils.rs

@@ -0,0 +1,32 @@
+use std::str;
+
+pub fn parse_hex_string(input: &String) -> [u8; 16] {
+    let mut result = [0u8; 16];
+    // Iterate over the string 2 characters each time
+    for (index, byte) in input.as_bytes().chunks(2).enumerate() {
+        if let Ok(byte_str) = str::from_utf8(byte) {
+            if let Ok(byte_value) = u8::from_str_radix(byte_str, 16) {
+                result[index] = byte_value;
+            } else {
+                panic!("Failed to parse byte");
+            }
+        } else {
+            panic!("Invalid byte string");
+        }
+    }
+    result
+}
+
+pub fn string_to_u8_16(input: &String) -> [u8; 16] {
+    let mut result = [0u8; 16];
+    let bytes = input.as_bytes();
+    let len = bytes.len().min(16);
+    result[..len].copy_from_slice(&bytes[..len]);
+    result
+}
+
+pub fn u8_16_to_string(input: &[u8; 16]) -> String {
+    let slice = &input[..];
+    let string = str::from_utf8(slice).expect("Invalid UTF-8");
+    string.to_string()
+}