ferric_crypto_lib 0.2.7

use crate::crypto_systems::mix::{Mix, MixKey};
use crate::error::MixError;
use crate::prelude::ALPHABET_LEN;
use crate::Traits::{BruteForce, Decrypt};
use itertools::Itertools;
use rayon::prelude::*;
use std::collections::{HashMap, HashSet};
use std::sync::Mutex;
use strsim::levenshtein;

/// The `PermutationIterator` struct is used to generate all possible permutations of keys for the brute force decryption process.
/// It keeps track of the current permutation, the maximum value a key can have, and whether all permutations have been exhausted.
/// It also stores any known keys, which are not included in the permutations.
pub struct PermutationIterator {
    vec_len: usize,
    current: [Vec<usize>; 3],
    max_value: usize,
    exhausted: bool,
    known_keys: [Option<Vec<usize>>; 3],
}

impl PermutationIterator {
    /// Constructs a new `PermutationIterator`.
    ///
    /// # Arguments
    ///
    /// * `vec_len` - The length of the key vectors.
    /// * `known_keys` - An array of optional vectors representing the known keys.
    ///
    /// # Returns
    ///
    /// * `Self` - A new `PermutationIterator`.
    fn new(vec_len: usize, known_keys: [Option<Vec<usize>>; 3]) -> Self {
        let mut current = [vec![0; vec_len], vec![0; vec_len], vec![0; vec_len]];

        // Initialize with the known keys and set the starting point for unknown keys
        for (i, key) in known_keys.iter().enumerate() {
            if let Some(ref key_val) = key {
                current[i] = key_val.clone();
            }
        }

        Self {
            vec_len,
            current,
            max_value: 27,
            exhausted: false,
            known_keys,
        }
    }

    /// Generates all possible permutations of keys in parallel.
    ///
    /// # Returns
    ///
    /// * `impl ParallelIterator<Item = [Vec<usize>; 3]>` - A parallel iterator over all possible permutations of keys.
    pub fn all_permutations(&self) -> impl ParallelIterator<Item = [Vec<usize>; 3]> {
        let iterators: Vec<_> = self
            .known_keys
            .iter()
            .map(|key| {
                if let Some(ref key_val) = key {
                    vec![key_val.clone()].into_iter()
                } else {
                    Self::generate_permutations(0, self.max_value, self.vec_len).into_iter()
                }
            })
            .collect();

        // Convert the vector of iterators into an iterator and then to a parallel iterator
        iterators
            .into_iter()
            .multi_cartesian_product()
            .par_bridge()
            .map(|perm| [perm[0].clone(), perm[1].clone(), perm[2].clone()])
    }

    /// Generates all permutations of a given length within a range.
    ///
    /// # Arguments
    ///
    /// * `min` - The minimum value a key can have.
    /// * `max` - The maximum value a key can have.
    /// * `length` - The length of the permutations.
    ///
    /// # Returns
    ///
    /// * `Vec<Vec<usize>>` - A vector of all permutations.
    fn generate_permutations(min: usize, max: usize, length: usize) -> Vec<Vec<usize>> {
        if min > max || length == 0 {
            return vec![];
        }

        // Base case: a permutation of length 1 is just a list of elements from min to max
        if length == 1 {
            return (min..=max).map(|x| vec![x]).collect();
        }

        // Recursive case: extend each smaller permutation
        let mut permutations = vec![];
        for perm in Self::generate_permutations(min, max, length - 1) {
            for i in min..=max {
                let mut extended_perm = perm.clone();
                extended_perm.push(i);
                permutations.push(extended_perm);
            }
        }

        permutations
    }

    /// Increments the current permutation.
    ///
    /// # Arguments
    ///
    /// * `idx` - The index of the key vector to increment.
    ///
    /// # Returns
    ///
    /// * `bool` - Whether the increment was successful. If `false`, all permutations have been exhausted.
    fn increment(&mut self, idx: usize) -> bool {
        if idx >= self.current.len() {
            self.exhausted = true;
            return false;
        }

        if self.known_keys[idx].is_some() {
            // Skip incrementing known keys and move to the next
            return self.increment(idx + 1);
        }

        let vec = &mut self.current[idx];
        for i in 0..vec.len() {
            if vec[i] < self.max_value {
                vec[i] += 1;
                return true;
            } else {
                vec[i] = 0; // Reset and carry over
            }
        }

        // If all elements in this vector reached max_value, carry to next
        self.increment(idx + 1)
    }
}

impl Iterator for PermutationIterator {
    type Item = [Vec<usize>; 3];

    /// Returns the next permutation.
    ///
    /// # Returns
    ///
    /// * `Option<Self::Item>` - The next permutation, or `None` if all permutations have been exhausted.
    fn next(&mut self) -> Option<Self::Item> {
        if self.exhausted {
            None
        } else {
            let result = self.current.clone();
            self.increment(0);
            Some(result)
        }
    }
}

type BruteForceResult = HashMap<MixKey, String>;

impl BruteForce<BruteForceResult, MixError, PermutationIterator, [Option<Vec<usize>>; 3]> for Mix {
    /// This function is used to brute force the decryption of a given cipher text.
    /// It generates all possible permutations of keys and tries to decrypt the cipher text with each key.
    /// If a clear text is provided, it calculates the similarity between the decrypted text and the clear text.
    /// If the similarity is above a certain threshold (90% by default), it saves the key and the decrypted text in a result map.
    /// If the similarity is exactly 1.0 (i.e., the decrypted text is exactly the same as the clear text), it stops the brute force process.
    /// If no clear text is provided, it saves all keys and their corresponding decrypted texts in the result map.
    /// The function returns the result map.
    ///
    /// # Arguments
    ///
    /// * `cipher_text` - A string representing the text to be decrypted.
    /// * `clear_text` - An optional string representing the clear text. If provided, it is used to calculate the similarity with the decrypted text.
    /// * `known_keys` - An array of optional vectors representing the known keys. If a key is known, it is not included in the permutations.
    ///
    /// # Returns
    ///
    /// * `Result<BruteForceResult, MixError>` - A result containing either a map of keys and their corresponding decrypted texts, or an error.
    #[cfg(not(feature = "parallel"))]
    fn brute_force(
        &mut self,
        cipher_text: String,
        clear_text: Option<String>,
        known_keys: [Option<Vec<usize>>; 3],
    ) -> Result<BruteForceResult, MixError> {
        // Generate all possible permutations of keys
        let mut iterator = self.gen_permutations(known_keys)?;
        let mut result = HashMap::new();
        let mut filtered_result = HashMap::new();
        let similarity_threshold = 0.90; // 90% similarity

        // Try to decrypt the cipher text with each key
        while let Some(permutation) = iterator.next() {
            let a = permutation[0].clone();
            let b = permutation[1].clone();
            let c = permutation[2].clone();

            let key = MixKey::new(a, b, c).map_err(|_| MixError::InvalidKey)?;
            let mix = Mix::new(key.clone());
            let decryption_result = mix.decrypt(cipher_text.clone())?;

            // If a clear text is provided, calculate the similarity with the decrypted text
            if let Some(ref plain_text) = clear_text {
                let similarity = calculate_similarity(&decryption_result, plain_text);
                if similarity >= similarity_threshold {
                    filtered_result.insert(key.clone(), decryption_result.clone());
                    if similarity == 1.0 {
                        println!("Found exact match: {:?}", key);
                        break;
                    }
                }
            } else {
                result.insert(key.clone(), decryption_result);
            }
        }

        // Use the filtered result if it has entries; otherwise, use the full result.
        let final_result = if !filtered_result.is_empty() {
            filtered_result
        } else if clear_text.is_some() {
            println!(
                "No matches found with similarity >= {}.",
                similarity_threshold
            );
            result
        } else {
            HashMap::new()
        };

        Ok(final_result)
    }

    /// This function is used to brute force the decryption of a given cipher text in parallel.
    /// It generates all possible permutations of keys and tries to decrypt the cipher text with each key.
    /// If a clear text is provided, it calculates the similarity between the decrypted text and the clear text.
    /// If the similarity is above a certain threshold (90% by default), it saves the key and the decrypted text in a result map.
    /// If no clear text is provided, it saves all keys and their corresponding decrypted texts in the result map.
    /// The function returns the result map.
    ///
    /// This function is only available when the "parallel" feature is enabled.
    ///
    /// # Arguments
    ///
    /// * `cipher_text` - A string representing the text to be decrypted.
    /// * `clear_text` - An optional string representing the clear text. If provided, it is used to calculate the similarity with the decrypted text.
    /// * `known_keys` - An array of optional vectors representing the known keys. If a key is known, it is not included in the permutations.
    ///
    /// # Returns
    ///
    /// * `Result<BruteForceResult, MixError>` - A result containing either a map of keys and their corresponding decrypted texts, or an error.
    #[cfg(feature = "parallel")]
    fn brute_force(
        &mut self,
        cipher_text: String,
        clear_text: Option<String>,
        known_keys: [Option<Vec<usize>>; 3],
    ) -> Result<BruteForceResult, MixError> {
        let iterator = self.gen_permutations(known_keys)?;
        let all_permutations = iterator.all_permutations();

        // Use a Mutex to safely collect results from multiple threads
        let results = Mutex::new(HashMap::new());

        all_permutations.for_each(|permutation| {
            let a = permutation[0].clone();
            let b = permutation[1].clone();
            let c = permutation[2].clone();

            let key = match MixKey::new(a, b, c) {
                Ok(k) => k,
                Err(_) => return, // Handle error appropriately
            };
            let mix = Mix::new(key.clone());
            let decryption_result = match mix.decrypt(cipher_text.clone()) {
                Ok(res) => res,
                Err(_) => return, // Handle error appropriately
            };

            if let Some(ref plain_text) = clear_text {
                let similarity = calculate_similarity(&decryption_result, plain_text);
                if similarity >= 0.90 {
                    let mut res = results.lock().unwrap();
                    res.insert(key.clone(), decryption_result.clone());
                }
            }
        });

        Ok(results.into_inner().unwrap())
    }

    /// This function generates a new PermutationIterator based on the known keys.
    ///
    /// # Arguments
    ///
    /// * `known_keys` - An array of optional vectors representing the known keys. If a key is known, it is not included in the permutations.
    ///
    /// # Returns
    ///
    /// * `Result<PermutationIterator, MixError>` - A result containing a new PermutationIterator or an error.
    fn gen_permutations(
        &mut self,
        known_keys: [Option<Vec<usize>>; 3],
    ) -> Result<PermutationIterator, MixError> {
        Ok(PermutationIterator::new(self.key.len, known_keys))
    }
}

/// This function calculates the similarity between two strings using the Levenshtein distance.
/// The Levenshtein distance is a string metric for measuring the difference between two sequences.
/// It is calculated as the minimum number of single-character edits (insertions, deletions or substitutions) required to change one word into the other.
/// The function returns a similarity score between 0.0 and 1.0, where 1.0 means the strings are identical.
///
/// # Arguments
///
/// * `text1` - The first string to compare.
/// * `text2` - The second string to compare.
///
/// # Returns
///
/// * `f64` - The similarity score between the two strings.
fn calculate_similarity(text1: &str, text2: &str) -> f64 {
    let distance = levenshtein(text1, text2);
    let max_len = text1.len().max(text2.len());
    if max_len == 0 {
        return 1.0; // both strings are empty
    }
    1.0 - (distance as f64 / max_len as f64)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    #[test]
    fn brute_force_returns_correct_decryption_for_known_key() {
        let k1 = vec![13, 5, 1, 0];
        let k2 = vec![12, 4, 16, 8];
        let k3 = vec![1, 24, 2, 21];

        let clear_text = "orosmoln".to_string();
        let key = MixKey::new(k1.clone(), k2.clone(), k3.clone()).expect("Key is invalid");
        let mut mix = Mix::new(key.clone());
        let cipher_text = String::from("HJUMTKLC");
        let known_keys: [Option<Vec<usize>>; 3] = [Some(k1), None, Some(k3)];

        let result = mix
            .brute_force(cipher_text, Some(clear_text.clone()), known_keys)
            .unwrap();
        // look up the known key in the result map
        let result_key = result.keys().find(|&k| k == &key).unwrap();
        assert_eq!(result_key, &key);
        assert_eq!(result.values().next().unwrap(), &clear_text);
    }
}