stage

problems/p6/src/main.rs

@@ -0,0 +1,148 @@
+use core::f32;
+use std::fs;
+use std::{cmp::Ordering, str::from_utf8};
+
+use std::collections::HashMap;
+
+use anyhow::{Result, anyhow};
+use base64::{Engine, engine::general_purpose::STANDARD};
+use common::xor_with_key;
+
+fn base64_to_bytes(input: &str) -> Result<Vec<u8>> {
+    Ok(STANDARD.decode(input)?)
+}
+
+fn hamming_distance(input_1: &[u8], input_2: &[u8]) -> Result<u32> {
+    if input_1.len() != input_2.len() {
+        return Err(anyhow!("Can't count hamming distance"));
+    }
+
+    let hamming_distance = input_1
+        .iter()
+        .zip(input_2.iter())
+        .map(|(a, b)| (a ^ b).count_ones())
+        .sum::<u32>();
+    Ok(hamming_distance)
+}
+
+fn get_posible_keysize(cipher_bytes: &[u8]) -> Result<usize> {
+    let mut posible_keysizes: HashMap<usize, f32> = HashMap::new();
+    for keysize in 2..=40 {
+        if cipher_bytes.len() < 4 * keysize {
+            return Err(anyhow!("Wrong keysize"));
+        }
+        let mut total_distance = 0;
+
+        for i in 0..=(cipher_bytes.len() - 4 * keysize) {
+            let slice_1 = &cipher_bytes[i..i + keysize];
+            let slice_2 = &cipher_bytes[i + keysize..i + 2 * keysize];
+            let slice_3 = &cipher_bytes[i + 2 * keysize..i + 3 * keysize];
+            let slice_4 = &cipher_bytes[i + 3 * keysize..i + 4 * keysize];
+            total_distance += hamming_distance(slice_1, slice_2)?;
+            total_distance += hamming_distance(slice_1, slice_3)?;
+            total_distance += hamming_distance(slice_1, slice_4)?;
+            total_distance += hamming_distance(slice_2, slice_3)?;
+            total_distance += hamming_distance(slice_2, slice_4)?;
+            total_distance += hamming_distance(slice_3, slice_4)?;
+        }
+        let score = total_distance as f32 / keysize as f32;
+        posible_keysizes.insert(keysize, score);
+    }
+    let mut posible_keysizes: Vec<_> = posible_keysizes.into_iter().collect();
+    posible_keysizes.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal));
+    Ok(posible_keysizes.first().unwrap().0)
+}
+
+fn score_english_text(input: &[u8]) -> f32 {
+    use std::collections::HashMap;
+    let mut score = 0_f32;
+    let freq: HashMap<u8, f32> = [
+        (b'a', 0.0817),
+        (b'b', 0.0150),
+        (b'c', 0.0278),
+        (b'd', 0.0425),
+        (b'e', 0.1270),
+        (b'f', 0.0223),
+        (b'g', 0.0202),
+        (b'h', 0.0609),
+        (b'i', 0.0697),
+        (b'j', 0.0015),
+        (b'k', 0.0077),
+        (b'l', 0.0403),
+        (b'm', 0.0241),
+        (b'n', 0.0675),
+        (b'o', 0.0751),
+        (b'p', 0.0193),
+        (b'q', 0.0010),
+        (b'r', 0.0599),
+        (b's', 0.0633),
+        (b't', 0.0906),
+        (b'u', 0.0276),
+        (b'v', 0.00981),
+        (b'w', 0.0236),
+        (b'x', 0.0015),
+        (b'y', 0.0197),
+        (b'z', 0.0007),
+        (b' ', 0.1300),
+    ]
+    .iter()
+    .cloned()
+    .collect();
+
+    for c in input.iter().map(|b| b.to_ascii_lowercase()) {
+        if let Some(f) = freq.get(&c) {
+            score += f;
+        }
+    }
+    score
+}
+
+fn main() -> Result<()> {
+    let cipher = fs::read_to_string("./problems/p6/6.txt")?;
+    let cipher: String = cipher.lines().collect();
+    let cipher_bytes = base64_to_bytes(&cipher)?;
+
+    let posible_keysize = get_posible_keysize(&cipher_bytes)?;
+    // posible_keysize = 29
+
+    // 按密钥长度分组
+    let mut blocks: Vec<Vec<u8>> = vec![Vec::new(); posible_keysize];
+    for (i, &byte) in cipher_bytes.iter().enumerate() {
+        blocks[i % posible_keysize].push(byte);
+    }
+    let mut key_list: Vec<u8> = Vec::new();
+    for block in blocks {
+        let mut high_score = 0_f32;
+        let mut possible_key = 0_u8;
+        for key in 0..=255 {
+            let decrypted: Vec<u8> = block.iter().map(|&b| b ^ key).collect();
+            let score = score_english_text(&decrypted);
+            if score > high_score {
+                high_score = score;
+                possible_key = key;
+            }
+        }
+        key_list.push(possible_key);
+    }
+    let str_key = String::from_utf8(key_list.clone())?;
+    println!("key: {str_key}");
+    println!();
+    let plaintext = xor_with_key(&cipher_bytes, &key_list)?;
+    let str_plaintext = from_utf8(&plaintext)?;
+    println!("plaintext: {str_plaintext}");
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    // 导入外部作用域中的所有内容
+    use super::*;
+
+    #[test]
+    fn test_hamming_distance() {
+        let str_1 = "this is a test".as_bytes();
+        let str_2 = "wokka wokka!!!".as_bytes();
+        let result = hamming_distance(str_1, str_2).unwrap();
+        assert_eq!(result, 37)
+    }
+}