ferric_crypto_lib 0.2.7

use crate::error::RSAError;
use crate::utils::{BaseString, EncodedString, StringType};
use pyo3::{prelude::*, pymethods, types::PyString, PyResult};
use rug::Integer;
use std::fmt::Display;

// annoyingly we need biguint for the factorization :(
use num_bigint::BigUint;
use num_prime::nt_funcs::factorize;
use std::str::FromStr;

#[derive(FromPyObject)]
pub enum StringOrBaseString {
    #[pyo3(annotation = "String")]
    StdString(String),
    #[pyo3(annotation = "BaseString")]
    MyBaseString(BaseString),
}

impl From<StringOrBaseString> for BaseString {
    fn from(s: StringOrBaseString) -> Self {
        match s {
            StringOrBaseString::StdString(s) => BaseString::new(s),
            StringOrBaseString::MyBaseString(s) => s,
        }
    }
}

impl From<String> for StringOrBaseString {
    fn from(s: String) -> Self {
        StringOrBaseString::StdString(s)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
#[pyclass]
pub struct PyBaseString {
    pub base_string: BaseString,
}

#[derive(Debug, Clone, PartialEq, Eq)]
#[pyclass]
pub struct PyEncodedString {
    pub encoded_string: EncodedString,
}

#[pymethods]
impl PyBaseString {
    #[new]
    pub fn new(data: String) -> Self {
        PyBaseString {
            base_string: BaseString::new(data),
        }
    }

    pub fn encode(&self) -> PyResult<PyEncodedString> {
        match self.base_string.encode() {
            Ok(encoded) => Ok(PyEncodedString {
                encoded_string: encoded,
            }),
            Err(e) => Err(pyo3::exceptions::PyValueError::new_err(format!("{:?}", e))),
        }
    }

    pub fn encode_asym(&self) -> PyResult<PyEncodedString> {
        match self.base_string.encode_asym() {
            Ok(encoded) => Ok(PyEncodedString {
                encoded_string: encoded,
            }),
            Err(e) => Err(pyo3::exceptions::PyValueError::new_err(format!("{:?}", e))),
        }
    }

    pub fn to_uppercase(&mut self) {
        <BaseString as Clone>::clone(&self.base_string).to_uppercase();
    }

    pub fn get_data(&self) -> String {
        self.base_string.data.clone()
    }

    fn __repr__(&self) -> PyResult<String> {
        Ok(format!("BaseString({})", self.base_string.data))
    }

    fn __str__(&self) -> PyResult<String> {
        Ok(format!("{}", self.base_string.data))
    }
}

#[pymethods]
impl PyEncodedString {
    #[new]
    pub fn new(data: Vec<usize>, str_type: String) -> PyResult<Self> {
        let stype = match str_type.as_str() {
            "Standard" => StringType::Standard,
            "Asymetric" => StringType::Assymetric,
            _ => {
                return Err(pyo3::exceptions::PyValueError::new_err(
                    "Invalid StringType",
                ))
            }
        };
        Ok(Self {
            encoded_string: EncodedString::new(data, stype),
        })
    }

    pub fn decode(&self) -> PyResult<String> {
        match self.encoded_string.decode() {
            Ok(base_string) => Ok(base_string.data),
            Err(e) => Err(pyo3::exceptions::PyValueError::new_err(format!("{:?}", e))),
        }
    }

    pub fn flatten(&self) -> String {
        self.encoded_string.flatten()
    }

    pub fn get_data(&self) -> Vec<usize> {
        self.encoded_string.data.clone()
    }

    pub fn get_str_type(&self) -> String {
        match self.encoded_string.str_type {
            StringType::Standard => "Standard".to_string(),
            StringType::Assymetric => "Asymetric".to_string(),
            _ => "Unknown".to_string(),
        }
    }

    /// String must start with either 'S' or 'A' to be valid internally
    #[staticmethod]
    pub fn decode_from_str(input: String) -> PyResult<String> {
        let encoded = EncodedString::from(input);
        match encoded.decode() {
            Ok(base_string) => Ok(base_string.data),
            Err(e) => Err(pyo3::exceptions::PyValueError::new_err(format!("{:?}", e))),
        }
    }

    fn __repr__(&self) -> PyResult<String> {
        Ok(format!(
            "EncodedString({:?}, {:?})",
            self.encoded_string.data, self.encoded_string.str_type
        ))
    }

    fn __str__(&self) -> PyResult<String> {
        Ok(format!(
            "Data: {:?}, Type: {:?}",
            self.encoded_string.data, self.encoded_string.str_type
        ))
    }
}

impl Display for PyBaseString {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.base_string.data)
    }
}

impl Display for PyEncodedString {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{:?} {:?}",
            self.encoded_string.data, self.encoded_string.str_type
        )
    }
}

#[pyfunction]
pub fn factorize_n(n: String) -> PyResult<Vec<String>> {
    let n = BigUint::from_str(&n).unwrap();

    let factors = factorize(n);

    // convert the factors to strings and put in a vector
    let mut factors_str = Vec::new();
    for factor in factors {
        factors_str.push(format!("{:?}", factor.0));
    }

    Ok(factors_str)
}

impl From<PyBaseString> for BaseString {
    fn from(py_base_string: PyBaseString) -> Self {
        py_base_string.base_string
    }
}

impl From<PyEncodedString> for EncodedString {
    fn from(py_encoded_string: PyEncodedString) -> Self {
        py_encoded_string.encoded_string
    }
}

impl From<PyBaseString> for String {
    fn from(py_base_string: PyBaseString) -> Self {
        py_base_string.base_string.data
    }
}

impl From<PyEncodedString> for String {
    fn from(py_encoded_string: PyEncodedString) -> Self {
        py_encoded_string.encoded_string.flatten()
    }
}

impl From<BaseString> for PyBaseString {
    fn from(base_string: BaseString) -> Self {
        PyBaseString { base_string }
    }
}

impl From<EncodedString> for PyEncodedString {
    fn from(encoded_string: EncodedString) -> Self {
        PyEncodedString { encoded_string }
    }
}