anubis-age 1.4.0

//! ML-DSA-87 implementation for digital signatures.
//!
//! This module implements NIST Level-5 post-quantum digital signatures using ML-DSA-87
//! (formerly known as Dilithium5), as standardized in FIPS 204.
//!
//! ML-DSA signatures provide:
//! - File authenticity: Verify the file was created by the claimed sender
//! - Integrity: Detect any tampering with the encrypted file
//! - Non-repudiation: Cryptographic proof of file origin

use std::fmt;
use std::io;

use bech32::{ToBase32, Variant};
use oqs::sig::{Algorithm, Sig};
use zeroize::{Zeroize, Zeroizing};

use crate::util::parse_bech32;

const SECRET_KEY_PREFIX: &str = "ANUBIS-MLDSA-87-SECRET";
const PUBLIC_KEY_PREFIX: &str = "anubis1mldsa87";

/// The size in bytes of an ML-DSA-87 public key.
pub const MLDSA87_PUBLIC_KEY_BYTES: usize = 2592;
/// The size in bytes of an ML-DSA-87 secret key.
/// Note: liboqs reports 4896 bytes for ML-DSA-87 secret keys
pub const MLDSA87_SECRET_KEY_BYTES: usize = 4896;
/// The size in bytes of an ML-DSA-87 signature.
/// Note: liboqs reports 4627 bytes for ML-DSA-87 signatures
pub const MLDSA87_SIGNATURE_BYTES: usize = 4627;

fn mldsa87() -> Sig {
    oqs::init();
    Sig::new(Algorithm::MlDsa87).expect("ML-DSA-87 algorithm available")
}

fn invalid_data(msg: &str) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidData, msg)
}

/// An ML-DSA-87 signing key for creating digital signatures.
///
/// This provides NIST Level-5 post-quantum security for digital signatures.
#[derive(Clone)]
pub struct SigningKey {
    secret_key: Zeroizing<Vec<u8>>,
    public_key: Vec<u8>,
}

impl SigningKey {
    /// Generates a new ML-DSA-87 signing key.
    pub fn generate() -> Self {
        let sig = mldsa87();
        let (pk, sk) = sig.keypair().expect("ML-DSA keypair");
        Self {
            secret_key: Zeroizing::new(sk.as_ref().to_vec()),
            public_key: pk.as_ref().to_vec(),
        }
    }

    /// Serializes this signing key to a string.
    ///
    /// Returns a Bech32-encoded string starting with `ANUBIS-MLDSA-87-SECRET`.
    pub fn to_string(&self) -> String {
        let mut material = Vec::with_capacity(MLDSA87_SECRET_KEY_BYTES + MLDSA87_PUBLIC_KEY_BYTES);
        material.extend_from_slice(self.secret_key.as_ref());
        material.extend_from_slice(&self.public_key);

        let encoded = bech32::encode(SECRET_KEY_PREFIX, material.to_base32(), Variant::Bech32)
            .expect("valid HRP");

        material.zeroize();

        encoded.to_uppercase()
    }

    /// Returns the verification key corresponding to this signing key.
    pub fn to_public(&self) -> VerificationKey {
        VerificationKey {
            public_key: self.public_key.clone(),
        }
    }

    /// Signs a message with this signing key.
    ///
    /// Returns an ML-DSA-87 signature over the message.
    pub fn sign(&self, message: &[u8]) -> Result<Signature, io::Error> {
        let sig = mldsa87();

        let sk = sig
            .secret_key_from_bytes(self.secret_key.as_ref())
            .ok_or_else(|| invalid_data("invalid ML-DSA secret key"))?;

        let signature = sig
            .sign(message, sk)
            .map_err(|_| invalid_data("failed to sign message"))?;

        Ok(Signature {
            signature: signature.as_ref().to_vec(),
        })
    }
}

impl std::str::FromStr for SigningKey {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (hrp, bytes) = parse_bech32(s).ok_or("invalid Bech32 encoding")?;
        if !hrp.eq_ignore_ascii_case(SECRET_KEY_PREFIX) {
            return Err("incorrect HRP");
        }
        if bytes.len() != MLDSA87_SECRET_KEY_BYTES + MLDSA87_PUBLIC_KEY_BYTES {
            return Err("incorrect signing key length");
        }

        let secret_key = Zeroizing::new(bytes[..MLDSA87_SECRET_KEY_BYTES].to_vec());
        let public_key = bytes[MLDSA87_SECRET_KEY_BYTES..].to_vec();

        Ok(Self {
            secret_key,
            public_key,
        })
    }
}

/// An ML-DSA-87 verification key for verifying digital signatures.
///
/// This provides NIST Level-5 post-quantum security for signature verification.
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct VerificationKey {
    public_key: Vec<u8>,
}

impl VerificationKey {
    fn ensure_length(bytes: &[u8]) -> Result<(), &'static str> {
        if bytes.len() == MLDSA87_PUBLIC_KEY_BYTES {
            Ok(())
        } else {
            Err("incorrect verification key length")
        }
    }

    /// Verifies a signature over a message.
    ///
    /// Returns `Ok(())` if the signature is valid, `Err` otherwise.
    pub fn verify(&self, message: &[u8], signature: &Signature) -> Result<(), io::Error> {
        let sig = mldsa87();

        let pk = sig
            .public_key_from_bytes(&self.public_key)
            .ok_or_else(|| invalid_data("invalid ML-DSA public key"))?;

        let sig_bytes = sig
            .signature_from_bytes(&signature.signature)
            .ok_or_else(|| invalid_data("invalid ML-DSA signature"))?;

        sig.verify(message, sig_bytes, pk)
            .map_err(|_| invalid_data("signature verification failed"))
    }
}

impl fmt::Display for VerificationKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}",
            bech32::encode(
                PUBLIC_KEY_PREFIX,
                self.public_key.to_base32(),
                Variant::Bech32
            )
            .expect("valid HRP")
        )
    }
}

impl fmt::Debug for VerificationKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self)
    }
}

impl std::str::FromStr for VerificationKey {
    type Err = &'static str;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let (hrp, bytes) = parse_bech32(s).ok_or("invalid Bech32 encoding")?;
        if !hrp.eq_ignore_ascii_case(PUBLIC_KEY_PREFIX) {
            return Err("incorrect HRP");
        }
        Self::ensure_length(&bytes)?;
        Ok(Self { public_key: bytes })
    }
}

/// An ML-DSA-87 signature.
#[derive(Clone, PartialEq, Eq)]
pub struct Signature {
    signature: Vec<u8>,
}

impl Signature {
    /// Returns the signature bytes.
    pub fn as_bytes(&self) -> &[u8] {
        &self.signature
    }

    /// Creates a signature from bytes.
    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self, &'static str> {
        if bytes.len() == MLDSA87_SIGNATURE_BYTES {
            Ok(Self { signature: bytes })
        } else {
            Err("incorrect signature length")
        }
    }
}

impl fmt::Debug for Signature {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Signature({} bytes)", self.signature.len())
    }
}

#[cfg(test)]
mod tests {
    use super::{SigningKey, VerificationKey};

    #[test]
    fn sign_and_verify() {
        let signing_key = SigningKey::generate();
        let verification_key = signing_key.to_public();
        let message = b"Test message for ML-DSA-87 signature";

        let signature = signing_key.sign(message).unwrap();
        assert!(verification_key.verify(message, &signature).is_ok());

        // Wrong message should fail
        let wrong_message = b"Different message";
        assert!(verification_key.verify(wrong_message, &signature).is_err());
    }

    #[test]
    fn bech32_round_trip() {
        let signing_key = SigningKey::generate();

        // Debug: print actual sizes
        println!("Secret key size: {}", signing_key.secret_key.len());
        println!("Public key size: {}", signing_key.public_key.len());
        println!(
            "Total size: {}",
            signing_key.secret_key.len() + signing_key.public_key.len()
        );
        println!(
            "Expected total: {}",
            super::MLDSA87_SECRET_KEY_BYTES + super::MLDSA87_PUBLIC_KEY_BYTES
        );

        let encoded = signing_key.to_string();
        let reparsed: SigningKey = encoded.parse().unwrap();
        assert_eq!(signing_key.public_key, reparsed.public_key);

        let verification_key = signing_key.to_public();
        let encoded_vk = verification_key.to_string();
        let reparsed_vk: VerificationKey = encoded_vk.parse().unwrap();
        assert_eq!(verification_key, reparsed_vk);
    }

    #[test]
    fn signature_size() {
        let signing_key = SigningKey::generate();
        let message = b"Test";
        let signature = signing_key.sign(message).unwrap();
        assert_eq!(signature.as_bytes().len(), super::MLDSA87_SIGNATURE_BYTES);
    }
}