ferric_crypto_lib 0.2.7

use crate::error::CharacterParseError;
use crate::error::MixError;
use crate::prelude::{encode_string, ALPHABET_LEN};
use crate::utils::BaseString;
#[cfg(feature = "python-integration")]
use pyo3::{pyclass, PyErr};

#[derive(Debug, Default, Clone, Eq, PartialEq, Hash)]
#[cfg_attr(feature = "python-integration", pyclass(get_all))]
pub struct MixKey {
    pub keys: Vec<Vec<usize>>,
    pub len: usize,
}

#[cfg(not(feature = "python-integration"))]
impl MixKey {
    /// Constructs a new `MixKey` instance.
    ///
    /// # Arguments
    ///
    /// * `k1` - A vector of `usize` that represents the first key.
    /// * `k2` - A vector of `usize` that represents the second key.
    /// * `k3` - A vector of `usize` that represents the third key.
    ///
    /// # Returns
    ///
    /// * `Result<Self, MixError>` - A new `Key` instance if the keys are valid, or a `MixError` if they are not.
    ///
    /// # Example
    ///
    /// ```
    /// # use ferric_crypto_lib::crypto_systems::mix::MixKey;
    /// let k1 = vec![1, 2, 3];
    /// let k2 = vec![1, 2, 3];
    /// let k3 = vec![1, 2, 3];
    /// let key = MixKey::new(k1, k2, k3).unwrap();
    /// ```
    pub fn new(k1: Vec<usize>, k2: Vec<usize>, k3: Vec<usize>) -> Result<Self, MixError> {
        let k = vec![k1.clone(), k2, k3];

        let mut key = Self {
            keys: k,
            len: 0usize,
        };

        if !key.is_valid() {
            Err(MixError::InvalidKey)
        } else {
            key.len = k1.len();
            Ok(key)
        }
    }

    pub fn new_with_len(key_len: usize) -> Self {
        let k1 = vec![0; key_len];
        let k2 = vec![0; key_len];
        let k3 = vec![0; key_len];

        Self::new(k1, k2, k3).unwrap()
    }

    /// Checks if the keys in the `Key` instance are valid.
    ///
    /// # Returns
    ///
    /// * `bool` - `true` if the keys are valid, `false` otherwise.
    ///
    /// # Validity Criteria
    ///
    /// A key is considered valid if all of its elements are less than `ALPHABET_LEN`, and all keys have the same length.
    pub fn is_valid(&self) -> bool {
        for (i, key) in self.keys.iter().enumerate() {
            for num in key.iter() {
                if *num >= *ALPHABET_LEN {
                    return false;
                }
            }
        }

        if self.keys[0].len() != self.keys[1].len() || self.keys[0].len() != self.keys[2].len() {
            return false;
        }
        true
    }
}

/// The `Mix` struct represents a symmetric cryptographic system that is a mix of three ciphers: Feistel, Hill, and Vigenère.
/// This cryptographic system was developed by Robert Nyqvist.
///
/// A plaintext block `m` in `Z^(8)28` is divided into two vectors `L0` and `R0` in `Z^(4)28`, so that `m = (L0, R0)`.
/// The ciphertext `c = (L3, R3)` is then determined recursively according to
/// `Ln = Rn-1` and `Rn = Ln-1 + f(Rn-1, Kn) (mod 28)`,
/// where `K1, K2, K3` in `Z428` are the encryption keys.
/// The function `f` is given by `f(R, K) ≡ P R + K`,
/// where `P` is a 4x4 matrix with the values `[1, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 1] (mod 28)`.
/// To perform the calculation `P R + K`, `R` and `K` are considered as column matrices.
///
/// # Fields
///
/// * `key` - A `Key` instance that will be used for the encryption process.
/// * `rounds` - The number of rounds to be performed in the encryption process. Currently it only supports 3 rounds but i wish to make this configurable.
///
/// # Example
///
/// ```
/// # use ferric_crypto_lib::crypto_systems::mix::{Mix, MixKey};
/// # use ferric_crypto_lib::Traits::{Encrypt, Decrypt};
/// let k1 = vec![0, 27, 27, 0];
/// let k2 = vec![1, 0, 0, 1];
/// let k3 = vec![15, 15, 15, 15];
///
/// let clear_text = "rötmånad".to_string();
/// let key = MixKey::new(k1, k2, k3).expect("Key is invalid");
/// let mix = Mix::new(key);
/// let result = mix.encrypt(clear_text);
///
/// assert_eq!(result.unwrap(), "mlbrnbpb");
/// ```
#[derive(Debug, Default)]
#[cfg_attr(feature = "python-integration", pyclass(get_all))]
pub struct Mix {
    pub key: MixKey,
    pub rounds: usize,
}

#[cfg(not(feature = "python-integration"))]
impl Mix {
    /// Constructs a new `Mix` instance.
    ///
    /// # Arguments
    ///
    /// * `key` - A `MixKey` instance that will be used for the encryption process.
    ///
    /// # Returns
    ///
    /// * A new `Mix` instance.
    ///
    /// # Example
    ///
    /// ```
    /// # use ferric_crypto_lib::crypto_systems::mix::{Mix, MixKey};
    /// let key = MixKey::new(vec![1, 2, 3], vec![1, 2, 3], vec![1, 2, 3]).unwrap();
    /// let mix = Mix::new(key);
    /// ```
    pub fn new(key: MixKey) -> Self {
        Self {
            key,
            rounds: 3usize,
        }
    }

    /// Splits the input string into chunks of equal length based on the length of the key.
    ///
    /// # Arguments
    ///
    /// * `input` - A `String` that will be split into chunks.
    ///
    /// # Returns
    ///
    /// * `Result<Vec<Vec<usize>>, MixError>` - A 2D vector of `usize` if the operation is successful, or a `MixError` if it fails.
    ///
    /// # Errors
    ///
    /// This function will return an error if the length of the input string is not a multiple of the length of the key.
    ///
    /// # Example
    ///
    /// ```
    /// # use ferric_crypto_lib::crypto_systems::mix::{Mix, MixKey};
    /// let key = MixKey::new(vec![1, 2, 3], vec![1, 2, 3], vec![1, 2, 3]).unwrap();
    /// let mix = Mix::new(key);
    /// let input = "abcdef".to_string();
    /// let result = mix.get_split_input(input).unwrap();
    /// assert_eq!(result, vec![vec![0, 1, 2], vec![3, 4, 5]]);
    /// ```
    pub fn get_split_input(&self, input: String) -> Result<Vec<Vec<usize>>, MixError> {
        if input.chars().count() % self.key.keys[0].len() != 0 {
            return Err(MixError::InvalidInput);
        }

        let input = BaseString::from(input);

        match input.encode() {
            Ok(string_vec) => Ok(string_vec
                .data
                .chunks(self.key.len)
                .map(|chunk| chunk.to_vec())
                .collect()),
            Err(e) => Err(MixError::CharacterParseError(e)),
        }
    }
}

impl From<Vec<Vec<usize>>> for Mix {
    fn from(keys: Vec<Vec<usize>>) -> Self {
        // Construct the Key from the provided Vec<Vec<usize>>
        let key = MixKey::new(keys[0].clone(), keys[1].clone(), keys[2].clone())
            .expect("Invalid keys provided");

        // Create a new Mix instance with the constructed Key
        Self {
            key,
            rounds: 3, // you can set the default rounds here
        }
    }
}

#[cfg(feature = "python-integration")]
mod python_inegration {
    use super::*;
    use crate::Traits::{Decrypt, Encrypt};
    use pyo3::{prelude::*, pymethods, PyResult};
    use std::collections::HashMap;

    #[pymethods]
    impl Mix {
        /// Constructs a new `Mix` instance.
        ///
        /// # Arguments
        ///
        /// * `key` - A `MixKey` instance that will be used for the encryption process.
        ///
        /// # Returns
        ///
        /// * A new `Mix` instance.
        ///
        /// # Example
        ///
        /// ```
        /// let key = Key::new(vec![1, 2, 3], vec![1, 2, 3], vec![1, 2, 3]).unwrap();
        /// let mix = Mix::new(key);
        /// ```
        #[new]
        pub fn new(key: MixKey) -> Self {
            Self {
                key,
                rounds: 3usize,
            }
        }

        /// Splits the input string into chunks of equal length based on the length of the key.
        ///
        /// # Arguments
        ///
        /// * `input` - A `String` that will be split into chunks.
        ///
        /// # Returns
        ///
        /// * `Result<Vec<Vec<usize>>, MixError>` - A 2D vector of `usize` if the operation is successful, or a `MixError` if it fails.
        ///
        /// # Errors
        ///
        /// This function will return an error if the length of the input string is not a multiple of the length of the key.
        ///
        /// # Example
        ///
        /// ```
        /// let key = Key::new(vec![1, 2, 3], vec![1, 2, 3], vec![1, 2, 3]).unwrap();
        /// let mix = Mix::new(key);
        /// let input = "abcdef".to_string();
        /// let result = mix.get_split_input(input).unwrap();
        /// ```
        pub fn get_split_input(&self, input: String) -> Result<Vec<Vec<usize>>, MixError> {
            if input.chars().count() % self.key.keys[0].len() != 0 {
                return Err(MixError::InvalidInput);
            }

            match encode_string(&input) {
                Ok(string_vec) => Ok(string_vec
                    .chunks(self.key.len)
                    .map(|chunk| chunk.to_vec())
                    .collect()),
                Err(e) => Err(MixError::CharacterParseError(e)),
            }
        }

        pub fn print(&self) {
            println!("{:?}", self);
        }

        pub fn encrypt(&self, input: String) -> PyResult<String> {
            match Encrypt::encrypt(self, input) {
                Ok(s) => Ok(s),
                Err(e) => Err(pyo3::exceptions::PyException::new_err(format!("{:?}", e))),
            }
        }

        pub fn decrypt(&self, input: String) -> PyResult<String> {
            match Decrypt::decrypt(self, input) {
                Ok(s) => Ok(s),
                Err(e) => Err(pyo3::exceptions::PyException::new_err(format!("{:?}", e))),
            }
        }

        #[pyo3(signature = (input="".to_string(), clear_text=None, key_info=[None, None, None]))]
        pub fn brute_force(
            &mut self,
            input: String,
            clear_text: Option<String>,
            key_info: [Option<Vec<usize>>; 3],
        ) -> PyResult<HashMap<MixKey, String>> {
            match crate::Traits::BruteForce::brute_force(self, input, clear_text, key_info) {
                Ok(s) => Ok(s),
                Err(e) => Err(pyo3::exceptions::PyException::new_err(format!("{:?}", e))),
            }
        }

        #[staticmethod]
        pub fn create_from_keys(keys: Vec<Vec<usize>>) -> PyResult<Self> {
            if keys.len() != 3 {
                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                    "Expected exactly three keys",
                ));
            }

            let key = MixKey::new(keys[0].clone(), keys[1].clone(), keys[2].clone())
                .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("{:?}", e)))?;

            Ok(Mix::new(key))
        }
    }

    #[pymethods]
    impl MixKey {
        /// Constructs a new `MixKey` instance.
        ///
        /// # Arguments
        ///
        /// * `k1` - A vector of `usize` that represents the first key.
        /// * `k2` - A vector of `usize` that represents the second key.
        /// * `k3` - A vector of `usize` that represents the third key.
        ///
        /// # Returns
        ///
        /// * `Result<Self, MixError>` - A new `Key` instance if the keys are valid, or a `MixError` if they are not.
        ///
        /// # Example
        ///
        /// ```
        /// let k1 = vec![1, 2, 3];
        /// let k2 = vec![1, 2, 3];
        /// let k3 = vec![1, 2, 3];
        /// let key = Key::new(k1, k2, k3).unwrap();
        /// ```
        #[new]
        pub fn new(k1: Vec<usize>, k2: Vec<usize>, k3: Vec<usize>) -> Result<Self, MixError> {
            let k = vec![k1.clone(), k2, k3];

            let mut key = Self {
                keys: k,
                len: 0usize,
            };

            if !key.is_valid() {
                Err(MixError::InvalidKey)
            } else {
                key.len = k1.len();
                Ok(key)
            }
        }

        /// Checks if the keys in the `Key` instance are valid.
        ///
        /// # Returns
        ///
        /// * `bool` - `true` if the keys are valid, `false` otherwise.
        ///
        /// # Validity Criteria
        ///
        /// A key is considered valid if all of its elements are less than `ALPHABET_LEN`,
        /// and all keys have the same length.
        pub fn is_valid(&self) -> bool {
            for (i, key) in self.keys.iter().enumerate() {
                for num in key.iter() {
                    if *num >= *ALPHABET_LEN {
                        return false;
                    }
                }
            }

            if self.keys[0].len() != self.keys[1].len() || self.keys[0].len() != self.keys[2].len()
            {
                return false;
            }
            true
        }

        pub fn print_keys(&self) {
            for key in self.keys.iter() {
                println!("{:?}", key);
            }
        }
    }
}

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

    #[test]
    fn test_from() {
        let keys = vec![vec![1, 2, 3], vec![1, 2, 3], vec![1, 2, 3]];

        let mix = Mix::from(keys);

        assert_eq!(mix.key.keys[0], vec![1, 2, 3]);
        assert_eq!(mix.key.keys[1], vec![1, 2, 3]);
        assert_eq!(mix.key.keys[2], vec![1, 2, 3]);
    }

    #[test]
    fn key_creation_with_valid_values() {
        let k1 = vec![1, 2, 3];
        let k2 = vec![1, 2, 3];
        let k3 = vec![1, 2, 3];
        let key = MixKey::new(k1, k2, k3);
        assert!(key.is_ok());
    }

    #[test]
    fn key_creation_with_invalid_values() {
        let k1 = vec![1, 2, 3];
        let k2 = vec![1, 2, 3];
        let k3 = vec![1, 2, 3, 4]; // Extra element
        let key = MixKey::new(k1, k2, k3);
        assert!(key.is_err());
    }

    #[test]
    fn key_creation_with_values_exceeding_alphabet_len() {
        let k1 = vec![*ALPHABET_LEN, 2, 3]; // First element exceeds ALPHABET_LEN
        let k2 = vec![1, 2, 3];
        let k3 = vec![1, 2, 3];
        let key = MixKey::new(k1, k2, k3);
        assert!(key.is_err());
    }

    #[test]
    fn key_validation_with_valid_values() {
        let k1 = vec![1, 2, 3];
        let k2 = vec![1, 2, 3];
        let k3 = vec![1, 2, 3];
        let key = MixKey::new(k1, k2, k3).unwrap();
        assert!(key.is_valid());
    }

    #[test]
    fn key_validation_with_invalid_values() {
        let k1 = vec![1, 2, 3];
        let k2 = vec![1, 2, 3];
        let k3 = vec![1, 2, 3, 4]; // Extra element
        let key = MixKey::new(k1, k2, k3);
        assert!(key.is_err());
    }

    #[test]
    fn key_validation_with_values_exceeding_alphabet_len() {
        let k1 = vec![*ALPHABET_LEN, 2, 3]; // First element exceeds ALPHABET_LEN
        let k2 = vec![1, 2, 3];
        let k3 = vec![1, 2, 3];
        let key = MixKey::new(k1, k2, k3);
        assert!(key.is_err());
    }

    #[test]
    fn test_get_split_input() {
        let k1 = vec![1, 2, 3];
        let k2 = vec![1, 2, 3];
        let k3 = vec![1, 2, 3];
        let key = MixKey::new(k1, k2, k3).unwrap();
        let mix = Mix::new(key);
        let input = "abcdef".to_string();
        let result = mix.get_split_input(input).unwrap();
        let expected = vec![vec![0, 1, 2], vec![3, 4, 5]];
        assert_eq!(result, expected);
    }

    use crate::crypto_systems::mix::{Mix, MixKey};
    use crate::Traits::{Decrypt, Encrypt};

    fn brute_force() -> Vec<usize> {
        let k1 = vec![13, 5, 1, 0];
        let k2 = vec![12, 4, 16, 8];

        let crypto_text = String::from("HJUMTKLC");

        for i in 0..28 {
            for j in 0..28 {
                for k in 0..28 {
                    for l in 0..28 {
                        let k3 = vec![i, j, k, l];
                        let key = match MixKey::new(k1.clone(), k2.clone(), k3.clone()) {
                            Ok(k) => k,
                            Err(e) => {
                                println!("{:?}", e);
                                continue;
                            }
                        };
                        let mix = Mix::new(key);

                        match mix.decrypt(crypto_text.clone()) {
                            Ok(s) => {
                                if s == "orosmoln" {
                                    return k3.clone();
                                }
                            }
                            Err(e) => println!("{:?}", e),
                        };
                    }
                }
            }
        }
        vec![]
    }

    #[test]
    fn test_brute_force() {
        let k1 = vec![13, 5, 1, 0];
        let k2 = vec![12, 4, 16, 8];
        let k3 = brute_force();
        let key = match MixKey::new(k1, k2, k3) {
            Ok(k) => k,
            Err(e) => {
                println!("{:?}", e);
                return;
            }
        };

        let mix = Mix::new(key);

        let input = String::from("orosmoln");

        match mix.encrypt(input) {
            Ok(s) => println!("{}", s),
            Err(e) => println!("{:?}", e),
        };

        let crypto_text = String::from("YTVRVNLR");

        match mix.decrypt(crypto_text) {
            Ok(s) => println!("{}", s),
            Err(e) => println!("{:?}", e),
        };
    }
}