Documentation
use anyhow::{anyhow, Result};
use base64::URL_SAFE_NO_PAD;
use rsa::pkcs8::{
    DecodePrivateKey, DecodePublicKey, EncodePrivateKey, EncodePublicKey, LineEnding,
};
use rsa::PublicKey;
use sha2::Sha256;

use crate::cryptos;

pub struct CryptoRsaPrivateKey {
    pri_key: rsa::RsaPrivateKey,
}

pub struct CryptoRsaPublicKey {
    pub_key: rsa::RsaPublicKey,
}

impl CryptoRsaPrivateKey {
    pub fn new(bits: usize) -> Result<Self> {
        let mut rand = rand::rngs::OsRng;
        Ok(CryptoRsaPrivateKey {
            pri_key: rsa::RsaPrivateKey::new(&mut rand, bits)?,
        })
    }

    pub fn from_pri_key_str(private_key_str: &str) -> Result<Self> {
        let vec = base64::decode(private_key_str)?;
        Ok(CryptoRsaPrivateKey {
            pri_key: rsa::RsaPrivateKey::from_pkcs8_der(&vec)?,
        })
    }
    pub fn from(private_key_pem: &str) -> Result<Self> {
        Ok(CryptoRsaPrivateKey {
            pri_key: rsa::RsaPrivateKey::from_pkcs8_pem(private_key_pem)?,
        })
    }

    pub fn serialize(&self) -> Result<String> {
        Ok(self.pri_key.to_pkcs8_pem(LineEnding::LF)?.to_string())
    }

    pub fn decrypt(&self, encrypted_data: &str) -> Result<String> {
        let encrypted_data = hex::decode(encrypted_data)?;
        let data = self.pri_key.decrypt(
            rsa::PaddingScheme::PKCS1v15Encrypt,
            encrypted_data.as_slice(),
        )?;
        Ok(String::from_utf8(data)?)
    }

    pub fn sign(&self, data: &str) -> Result<String> {
        let data = cryptos::sha256(data)?;
        let signed_data = self.pri_key.sign(
            rsa::PaddingScheme::new_pkcs1v15_sign::<Sha256>(),
            hex::decode(data)?.as_slice(),
        )?;
        Ok(base64::encode_config(signed_data, URL_SAFE_NO_PAD))
    }
}

impl CryptoRsaPublicKey {
    pub fn from_private_key(private_key: &CryptoRsaPrivateKey) -> Result<Self> {
        let public_key = rsa::RsaPublicKey::from(&private_key.pri_key);
        Ok(CryptoRsaPublicKey {
            pub_key: public_key,
        })
    }

    pub fn from_private_key_str(private_key_pem: &str) -> Result<Self> {
        let private_key = rsa::RsaPrivateKey::from_pkcs8_pem(private_key_pem)?;
        let public_key = rsa::RsaPublicKey::from(private_key);
        Ok(CryptoRsaPublicKey {
            pub_key: public_key,
        })
    }

    pub fn from_public_key_str(public_key_str: &str) -> Result<Self> {
        let vec = base64::decode(public_key_str)?;
        Ok(CryptoRsaPublicKey {
            pub_key: rsa::RsaPublicKey::from_public_key_der(&vec)?,
        })
    }

    pub fn serialize(&self) -> Result<String> {
        match self.pub_key.to_public_key_pem(LineEnding::LF) {
            Ok(s) => Ok(s),
            Err(e) => Err(anyhow!("{e}")),
        }
    }

    pub fn encrypt(&self, data: &str) -> Result<String> {
        let mut rand = rand::rngs::OsRng;
        let encrypted_data = self.pub_key.encrypt(
            &mut rand,
            rsa::PaddingScheme::PKCS1v15Encrypt,
            data.as_bytes(),
        )?;
        Ok(hex::encode(encrypted_data))
    }

    pub fn verify(&self, data: &str, signed_data: &str) -> Result<bool> {
        let data = cryptos::sha256(data)?;
        let signed_data = base64::decode_config(signed_data, URL_SAFE_NO_PAD)?;
        let result = self.pub_key.verify(
            rsa::PaddingScheme::new_pkcs1v15_sign::<Sha256>(),
            hex::decode(data)?.as_slice(),
            signed_data.as_slice(),
        );
        match result {
            Ok(_) => Ok(true),
            Err(_) => Ok(false),
        }
    }
}