wecanencrypt 0.1.0

use rand::{CryptoRng, Rng};
use signature::{Signer as _, Verifier};
use slh_dsa::Shake256s;
use zeroize::ZeroizeOnDrop;

use crate::pgp::{
    crypto::{hash::HashAlgorithm, Signer},
    errors::{ensure, ensure_eq, format_err, Result},
    ser::Serialize,
    types::{SignatureBytes, SlhDsaShake256sPublicParams},
};

/// Size in bytes of the serialized secret key.
pub const KEY_LEN: usize = 128;

/// Secret key for SLH DSA Shake256s
#[derive(Clone, PartialEq, derive_more::Debug)]
pub struct SecretKey {
    /// The secret key.
    #[debug("..")]
    key: slh_dsa::SigningKey<Shake256s>,
}

impl Eq for SecretKey {}

impl ZeroizeOnDrop for SecretKey {}

impl Drop for SecretKey {
    fn drop(&mut self) {
        // TODO: zeroize, fixed in latest `master`
    }
}

impl From<&SecretKey> for SlhDsaShake256sPublicParams {
    fn from(value: &SecretKey) -> Self {
        Self {
            key: value.key.as_ref().clone(),
        }
    }
}

impl SecretKey {
    /// Generate an Ed448 `SecretKey`.
    pub fn generate<R: Rng + CryptoRng>(mut rng: R) -> Self {
        let key = slh_dsa::SigningKey::new(&mut rng);

        SecretKey { key }
    }

    /// Create a key from the raw byte values
    pub fn try_from_bytes(seed: [u8; KEY_LEN]) -> Result<Self> {
        let key = slh_dsa::SigningKey::<Shake256s>::try_from(&seed[..])
            .map_err(|e| format_err!("invalid key {:?}", e))?;
        Ok(Self { key })
    }

    /// Returns the secret key in their raw byte level representation.
    pub fn to_bytes(&self) -> [u8; KEY_LEN] {
        self.key.to_bytes().into()
    }
}

impl Signer for SecretKey {
    fn sign(&self, hash: HashAlgorithm, digest: &[u8]) -> Result<SignatureBytes> {
        ensure!(
            ![
                HashAlgorithm::Md5,
                HashAlgorithm::Ripemd160,
                HashAlgorithm::Sha1
            ]
            .contains(&hash),
            "invalid hash algorithm"
        );

        let sig = self.key.sign(digest);
        let bytes = sig.to_bytes();

        Ok(SignatureBytes::Native(bytes.to_vec().into()))
    }
}

impl Serialize for SecretKey {
    fn to_writer<W: std::io::Write>(&self, writer: &mut W) -> Result<()> {
        let key = self.to_bytes();
        writer.write_all(&key)?;
        Ok(())
    }

    fn write_len(&self) -> usize {
        KEY_LEN
    }
}

/// Verify an EdDSA signature.
pub fn verify(
    key: &slh_dsa::VerifyingKey<Shake256s>,
    hash: HashAlgorithm,
    hashed: &[u8],
    sig_bytes: &[u8],
) -> Result<()> {
    ensure!(
        ![
            HashAlgorithm::Md5,
            HashAlgorithm::Ripemd160,
            HashAlgorithm::Sha1
        ]
        .contains(&hash),
        "invalid hash algorithm"
    );
    ensure_eq!(sig_bytes.len(), 29792, "invalid signature length");

    let sig = slh_dsa::Signature::<Shake256s>::try_from(sig_bytes)?;
    key.verify(hashed, &sig)?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use proptest::prelude::*;
    use rand::SeedableRng;
    use rand_chacha::ChaCha8Rng;

    use super::*;

    impl Arbitrary for SecretKey {
        type Parameters = ();
        type Strategy = BoxedStrategy<Self>;

        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
            any::<u64>()
                .prop_map(|seed| {
                    let mut rng = ChaCha8Rng::seed_from_u64(seed);
                    SecretKey::generate(&mut rng)
                })
                .boxed()
        }
    }
}