ferric_crypto_lib 0.2.7

use crate::decrypt::rsa::*;
use crate::encrypt::rsa::*;
use crate::error::{CharacterParseError, RSAError};
use crate::Traits::{BruteForce, Decrypt, Encrypt};
use nalgebra::one;
use num_bigint::BigUint;
use rug::Integer;
use std::ops::Mul;
use std::str::FromStr;

use crate::utils::factorize;
#[cfg(feature = "python-integration")]
use pyo3::prelude::*;

/// Represents a RSA cipher.
#[derive(Default)]
#[cfg_attr(feature = "python-integration", pyclass)]
pub struct RSA {
    pub n: Integer,
    pub e: Integer,
    pub d: Integer,
    pub p: Integer,
    pub q: Integer,
}

#[cfg(not(feature = "python-integration"))]
impl RSA {
    // create a new RSA cipher from the given p and q values
    pub fn new(p: Integer, q: Integer) -> Self {
        let mut system = Self {
            p,
            q,
            ..Default::default()
        };

        let _ = system.generate_keys();
        system
    }

    /// Calculate the maximum block size for encryption based on the size of 'n'.
    pub fn calculate_block_size(&self) -> usize {
        // Get the bit length of 'n'
        let key_size_in_bits = self.n.significant_bits() as usize;

        // For PKCS#1 v1.5 padding, the padding overhead is 11 bytes.
        let padding_overhead = 11;

        // Convert the key size from bits to bytes and subtract the padding overhead.
        let block_size: isize = key_size_in_bits as isize / 8isize - padding_overhead;
        block_size as usize
    }

    fn generate_keys(&mut self) -> Result<(), RSAError> {
        // calculate n
        self.n = Integer::from(&self.p * &self.q);

        // calculate phi
        let phi = Integer::from(&self.p - 1) * Integer::from(&self.q - 1);

        // calculate e so that gcd(e, phi) = 1
        // calculate e so that gcd(e, phi) = 1
        self.e = Integer::from(1330643366620853071u64);
        while self.e.clone().gcd(&phi) != 1 {
            // increment e by 2 to keep it odd (optimization)
            self.e = Integer::from(&self.e + 2u8);

            // optional: put a limit to prevent an infinite loop
            if self.e > phi {
                return Err(RSAError::Error("Failed to generate keys".to_string()));
            }
        }

        // calculate d
        self.d = match self.e.clone().invert(&phi) {
            Ok(inv) => inv,
            Err(_) => return Err(RSAError::Error("No inverse found".to_string())),
        };

        Ok(())
    }

    pub fn generate_priv_key(&mut self, phi: Integer) -> Option<()> {
        Some(())
    }

    // create from the given n, e values
    pub fn from_public(n: Integer, e: Integer) -> Self {
        Self {
            n,
            e,
            ..Default::default()
        }
    }

    // Function to factorize 'n' into 'p' and 'q'
    /// This only really works for smaller numbers
    pub fn factorize_n(&mut self) -> Result<(Integer, Integer), RSAError> {
        // way to long time to factorize using this method
        let two = Integer::from(2);
        let one = Integer::from(1);

        // Starting from 2, try to divide 'n'
        let mut factor = two.clone();
        while Integer::from(&factor) * Integer::from(&factor) <= self.n {
            if Integer::from(&self.n % &factor) == 0 {
                self.p = factor.clone();
                self.q = Integer::from(&self.n / &factor);
                return Ok((self.p.clone(), self.q.clone())); // returns in this order: (p, q)
            }
            factor += &one; // Correctly increment factor by one
        }

        Err(RSAError::Error("Failed to factorize n".to_string()))
    }

    pub fn factorize_and_set_d(&mut self) -> Result<(), RSAError> {
        // i really dont like this solution, but it works...

        let tmp = self.n.clone().to_string();
        let big_num = BigUint::from_str(&tmp).unwrap();

        let factors = factorize(big_num);

        if factors.len() != 2 {
            return Err(RSAError::Error("Failed to factorize n".to_string()));
        }

        // extract p and q from the factors
        let p = factors[0].clone();
        let q = factors[1].clone();

        // put them into the struct
        self.p = Integer::from_str_radix(&p, 10).unwrap();
        self.q = Integer::from_str_radix(&q, 10).unwrap();

        let phi = Integer::from(&self.p - 1) * Integer::from(&self.q - 1);

        // generate private keys
        self.d = match self.e.clone().invert(&phi) {
            Ok(inv) => inv,
            Err(_) => return Err(RSAError::Error("No inverse found".to_string())),
        };

        Ok(())
    }

    fn is_valid(&self) -> Result<(), RSAError> {
        Ok(())
    }

    pub fn print_vals(&self) {
        println!(
            "n: {}, e: {}, d: {}, p: {}, q: {}",
            self.n, self.e, self.d, self.p, self.q
        );
    }
}

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

    #[pymethods]
    impl RSA {
        #[new]
        pub fn new(p: String, q: String, e: Option<String>) -> PyResult<Self> {
            let p = Integer::from_str_radix(p.as_str(), 10).map_err(|e| {
                PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Invalid p: {}", e))
            })?;
            let q = Integer::from_str_radix(q.as_str(), 10).map_err(|e| {
                PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Invalid q: {}", e))
            })?;

            let mut system = Self {
                p,
                q,
                ..Default::default()
            };
            system.generate_keys(e).map_err(|e| {
                PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
                    "Failed to generate keys: {}",
                    e
                ))
            })?;

            Ok(system)
        }

        #[pyo3(signature = (e=None))]
        fn generate_keys(&mut self, e: Option<String>) -> PyResult<()> {
            // calculate n
            self.n = (&self.p * &self.q).into();

            // calculate phi
            let phi = Integer::from(&self.p - 1) * Integer::from(&self.q - 1);

            // calculate e so that gcd(e, phi) = 1
            if let Some(e) = e {
                self.e = Integer::from_str_radix(e.as_str(), 10).map_err(|e| {
                    PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Invalid e: {}", e))
                })?;
            } else {
                self.e = Integer::from(65537);
            }
            while Integer::from(&self.e).gcd(&phi) != Integer::from(1) {
                // increment e by 2 to keep it odd (optimization)
                self.e += 2;

                // optional: put a limit to prevent an infinite loop
                if self.e > phi {
                    return Err(RSAError::Error("Failed to generate keys".to_string()).into());
                }
            }

            // calculate d
            self.d = self
                .e
                .clone()
                .invert(&phi)
                .map_err(|_| RSAError::Error("No inverse found".to_string()))?;

            Ok(())
        }

        // create from the given n, e values
        #[pyo3(signature = (n = "".to_string(), e = "".to_string()))]
        #[staticmethod]
        pub fn from_public(n: String, e: String) -> PyResult<Self> {
            let n = Integer::from_str_radix(&n, 10).map_err(|e| {
                PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Invalid n: {}", e))
            })?;
            let e = Integer::from_str_radix(&e, 10).map_err(|e| {
                PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Invalid e: {}", e))
            })?;

            Ok(Self {
                n,
                e,
                d: Integer::new(), // Default value as 'd' is not used
                p: Integer::new(), // Default value
                q: Integer::new(), // Default value
            })
        }

        // Function to factorize 'n' into 'p' and 'q'
        pub fn factorize_n(&mut self) -> PyResult<(String, String)> {
            let two = Integer::from(2);
            let one = Integer::from(1);

            // Starting from 2, try to divide 'n'
            let mut factor = two.clone();
            while Integer::from(&factor) * Integer::from(&factor) <= self.n {
                if Integer::from(&self.n % &factor) == Integer::from(0) {
                    self.p = factor.clone();
                    self.q = Integer::from(&self.n / &factor);
                    return Ok((self.p.clone().to_string(), self.q.clone().to_string()));
                    // returns in this order: (p, q)
                }
                factor = factor + &one; // Correctly increment factor by one
            }

            Err(RSAError::Error("Failed to factorize n".to_string()).into())
        }

        pub fn factorize_and_set_d(&mut self) -> PyResult<()> {
            // i really dont like this solution, but it works...

            let tmp = self.n.clone().to_string();
            let big_num = BigUint::from_str(&tmp).unwrap();

            let factors = factorize(big_num);

            if factors.len() != 2 {
                return Err(RSAError::Error("Failed to factorize n".to_string()).into());
            }

            // extract p and q from the factors
            let p = factors[0].clone();
            let q = factors[1].clone();

            // put them into the struct
            self.p = Integer::from_str_radix(&p, 10).unwrap();
            self.q = Integer::from_str_radix(&q, 10).unwrap();

            let phi = Integer::from(&self.p - 1) * Integer::from(&self.q - 1);

            // generate private keys
            self.d = match self.e.clone().invert(&phi) {
                Ok(inv) => inv,
                Err(_) => return Err(RSAError::Error("No inverse found".to_string()).into()),
            };

            Ok(())
        }

        fn is_valid(&self) -> Result<(), RSAError> {
            Ok(())
        }

        pub fn encrypt(&self, input: String) -> PyResult<String> {
            match Encrypt::encrypt(self, input.into()) {
                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))),
            }
        }

        pub fn brute_force(
            &mut self,
            input: String,
            clear_text: Option<String>,
        ) -> PyResult<HashMap<usize, String>> {
            match BruteForce::brute_force(self, input, clear_text, None) {
                Ok(s) => Ok(s),
                Err(e) => Err(pyo3::exceptions::PyException::new_err(format!("{:?}", e))),
            }
        }

        // Define Python specific methods here, static methods require the pyo3 decorator `#[staticmethod]`

        fn __str__(&self) -> PyResult<String> {
            let mut s = String::new();
            s.push_str(&format!("n: {}, e: {}", self.n, self.e));

            // print size of n and e
            s.push_str(&format!(
                "\nSize of n: {} bits, Size of e: {} bits",
                self.n.significant_bits(),
                self.e.significant_bits()
            ));

            Ok(s)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::utils::EncodedString;
    use crate::Traits::{Decrypt, Encrypt};
    use itertools::assert_equal;

    #[test]
    fn test_factorize() {
        let mut rsa_system = RSA::from_public(Integer::from(14317), Integer::from(7777));

        let excpected_factors = (Integer::from(103u32), Integer::from(139u32));

        match rsa_system.factorize_n() {
            Ok((p, q)) => {
                assert_eq!(p, excpected_factors.0);
                assert_eq!(q, excpected_factors.1);
                println!("p: {}, q: {}", p, q);
            }
            Err(e) => panic!("{}", e.to_string()),
        }
    }

    #[test]
    fn test_full_sys() {
        let rsa_system = RSA::new(Integer::from(853), Integer::from(857));

        let encrypted = rsa_system.encrypt("hello".into()).unwrap();
        println!("encrypted: {}", encrypted);
        let decrypted = rsa_system.decrypt(encrypted).unwrap();
        println!("decrypted: {}", decrypted);
    }

    #[test]
    fn test_larger_rsa_sys() {
        let mut rsa_system = RSA::from_public(
            36134934063919959150141797353966441u128.into(),
            1330643366620853071u128.into(),
        );

        let starting_word = "spillintebönorna".to_string();

        // test a encryption
        let encrypted = rsa_system.encrypt(starting_word.clone().into()).unwrap();
        let excpected_encrypted = "32107833669138743416991214827014308".to_string();
        assert_eq!(&encrypted, &excpected_encrypted);
        let excpected_factors = (
            Integer::from(4177248169415681u64),
            Integer::from(8650415919381337961u64),
        );

        match rsa_system.factorize_and_set_d() {
            Ok(_) => {}
            Err(e) => panic!("{}", e.to_string()),
        }

        assert_eq!(rsa_system.p, excpected_factors.0);
        assert_eq!(rsa_system.q, excpected_factors.1);

        // test a decryption
        let decrypted = rsa_system.decrypt(encrypted).unwrap();
        let decoded_decrypted = EncodedString::from('A'.to_string() + &*decrypted)
            .decode()
            .unwrap()
            .data;
        assert_eq!(decoded_decrypted, starting_word);
    }
}