//! RSASSA-PSS signatures (RFC 8017 / PKCS#1 v2.2 §8.1).
//!
//! PSS is the modern RSA signature padding; it supersedes PKCS#1 v1.5
//! signatures in every protocol that cares (TLS 1.3, JWS `PS*`,
//! X.509 `id-RSASSA-PSS`, CMS, ...). For new deployments this should
//! be the default.
//!
//! # Side-channel posture
//!
//! PSS itself is structurally CT (no secret-dependent branches in
//! the EMSA-PSS encode / decode), but it relies on the underlying
//! [`super::rsa::rsa_decrypt_raw`] for the signing direction, which
//! is currently **not protected against the Bellcore single-fault
//! attack** (roadmap item `T1-C` — see
//! `arcana/doc/sca/countermeasures/rsa.rst`). Coron-Mandal (Asiacrypt
//! 2009) showed PSS is **provably secure against random faults** in
//! a separate fault model, but practical Bellcore-class faults still
//! recover the key — the proof assumes the underlying RSA primitive
//! itself is fault-resistant.
//!
//! # Algorithm (EMSA-PSS)
//!
//! For a modulus of `modBits` bits and hash `H` with output `hLen`:
//!
//! 1. `mHash = H(message)`
//! 2. Generate a random `salt` of `sLen` bytes (typically `sLen == hLen`)
//! 3. `M' = (0x00)^8 || mHash || salt`
//! 4. `h = H(M')`
//! 5. `DB = PS || 0x01 || salt`  where `PS` is zero-padding
//! 6. `maskedDB = DB XOR MGF1_H(h, len(DB))`
//! 7. Clear the top `8*emLen - emBits` bits of `maskedDB[0]`
//!    (`emBits = modBits - 1`)
//! 8. `EM = maskedDB || h || 0xbc`
//! 9. `sig = EM^d mod n`  (RSASP1)
//!
//! Verification reverses the process and re-checks `h = H(M')`.
//!
//! # API
//!
//! ```rust,ignore
//! use arcana::rsa::pss::{pss_sign_msg, pss_verify_msg};
//! use arcana::hash::sha256::Sha256;
//!
//! let sig = pss_sign_msg::<Sha256>(&sk, msg, 32, &mut rng).unwrap();
//! assert!(pss_verify_msg::<Sha256>(&pk, msg, 32, &sig));
//! ```
//!
//! Precomputed-digest variants ([`pss_sign`] / [`pss_verify`]) are also
//! exposed, matching the ECDSA API convention -- the common case in
//! protocol implementations is to receive an already-hashed digest
//! from an upstream layer (X.509, CMS, ...).

use super::bigint::BigInt;
use super::rsa::{RsaPublicKey, RsaSecretKey, rsa_decrypt_raw, rsa_encrypt_raw};
use crate::Hasher;

// ============================================================================
// MGF1 (RFC 8017 Appendix B.2.1)
// ============================================================================

/// MGF1 mask generation function using hash `H`.
///
/// MGF1 produces a pseudo-random byte stream of length `len` by iterating
/// `H(seed || counter_be_u32)` for counter = 0, 1, 2, ...
///
/// Generic over `H` so PSS can be instantiated with SHA-256, SHA-384,
/// SHA-512, or any future `Hasher`.
fn mgf1<H: Hasher>(seed: &[u8], len: usize) -> Vec<u8> {
    let h_len = H::OUTPUT_LEN;
    let mut out = Vec::with_capacity(len);
    let mut counter: u32 = 0;
    while out.len() < len {
        let mut hasher = H::new();
        hasher.update(seed);
        hasher.update(&counter.to_be_bytes());
        let block = hasher.finalize();
        let take = (len - out.len()).min(h_len);
        out.extend_from_slice(&block[..take]);
        counter += 1;
    }
    out.truncate(len);
    out
}

// ============================================================================
// EMSA-PSS encode / verify (RFC 8017 §9.1)
// ============================================================================

/// EMSA-PSS-Encode (RFC 8017 §9.1.1).
///
/// Produces the encoded message `EM` of length `emLen = ceil(em_bits / 8)`
/// from a precomputed message digest `m_hash` and a caller-supplied `salt`.
///
/// Returns `None` if:
/// - `m_hash.len() != H::OUTPUT_LEN`
/// - `emLen < hLen + sLen + 2` (the modulus is too short for this (H, sLen))
fn emsa_pss_encode<H: Hasher>(m_hash: &[u8], em_bits: usize, salt: &[u8]) -> Option<Vec<u8>> {
    let h_len = H::OUTPUT_LEN;
    let s_len = salt.len();
    let em_len = (em_bits + 7) / 8;

    // Step 1-3: length checks
    if m_hash.len() != h_len {
        return None;
    }
    if em_len < h_len + s_len + 2 {
        return None;
    }

    // Step 4-6: M' = (0x00)^8 || mHash || salt ; h = H(M')
    let mut m_prime = Vec::with_capacity(8 + h_len + s_len);
    m_prime.extend_from_slice(&[0u8; 8]);
    m_prime.extend_from_slice(m_hash);
    m_prime.extend_from_slice(salt);
    let h = H::hash(&m_prime);

    // Step 7-8: DB = PS || 0x01 || salt
    //   PS has length db_len - s_len - 1, all zeros.
    let db_len = em_len - h_len - 1;
    let mut db = vec![0u8; db_len];
    let ps_len = db_len - s_len - 1;
    db[ps_len] = 0x01;
    db[ps_len + 1..].copy_from_slice(salt);

    // Step 9-10: maskedDB = DB XOR MGF(h, db_len)
    let db_mask = mgf1::<H>(&h, db_len);
    for i in 0..db_len {
        db[i] ^= db_mask[i];
    }

    // Step 11: clear the top (8*em_len - em_bits) bits of maskedDB[0].
    // This ensures EM, interpreted as an integer, is strictly less
    // than 2^em_bits < n, so the RSA exponentiation is safe.
    let clear_bits = 8 * em_len - em_bits;
    if clear_bits > 0 {
        db[0] &= 0xff_u8 >> clear_bits;
    }

    // Step 12: EM = maskedDB || h || 0xbc
    let mut em = Vec::with_capacity(em_len);
    em.extend_from_slice(&db);
    em.extend_from_slice(&h);
    em.push(0xbc);

    Some(em)
}

/// EMSA-PSS-Verify (RFC 8017 §9.1.2).
///
/// Checks that `em` is a valid encoding of `m_hash` for the given
/// `em_bits` and `s_len`. Returns `true` iff the encoding is consistent.
fn emsa_pss_verify<H: Hasher>(m_hash: &[u8], em: &[u8], em_bits: usize, s_len: usize) -> bool {
    let h_len = H::OUTPUT_LEN;
    let em_len = (em_bits + 7) / 8;

    // Step 1-3: length checks
    if m_hash.len() != h_len {
        return false;
    }
    if em.len() != em_len {
        return false;
    }
    if em_len < h_len + s_len + 2 {
        return false;
    }

    // Step 4: last byte must be 0xbc
    if em[em_len - 1] != 0xbc {
        return false;
    }

    // Step 5-6: split EM into maskedDB and H
    let db_len = em_len - h_len - 1;
    let masked_db = &em[..db_len];
    let h = &em[db_len..db_len + h_len];

    // Step 7: top (8*em_len - em_bits) bits of maskedDB[0] must be zero.
    let clear_bits = 8 * em_len - em_bits;
    if clear_bits > 0 && (masked_db[0] >> (8 - clear_bits)) != 0 {
        return false;
    }

    // Step 8-9: DB = maskedDB XOR MGF(h, db_len), then clear the top
    // (8*em_len - em_bits) bits of DB[0].
    let db_mask = mgf1::<H>(h, db_len);
    let mut db = vec![0u8; db_len];
    for i in 0..db_len {
        db[i] = masked_db[i] ^ db_mask[i];
    }
    if clear_bits > 0 {
        db[0] &= 0xff_u8 >> clear_bits;
    }

    // Step 10: PS = (0x00)^(db_len - s_len - 1), then 0x01.
    let ps_len = db_len - s_len - 1;
    for byte in &db[..ps_len] {
        if *byte != 0 {
            return false;
        }
    }
    if db[ps_len] != 0x01 {
        return false;
    }

    // Step 11: salt = last s_len bytes of DB
    let salt = &db[ps_len + 1..];

    // Step 12-13: M' = (0x00)^8 || m_hash || salt ; h' = H(M')
    let mut m_prime = Vec::with_capacity(8 + h_len + s_len);
    m_prime.extend_from_slice(&[0u8; 8]);
    m_prime.extend_from_slice(m_hash);
    m_prime.extend_from_slice(salt);
    let h_prime = H::hash(&m_prime);

    // Step 14: h == h'  (constant-time compare, just good hygiene
    //                    even though `h` is not secret here)
    let mut diff = 0u8;
    for (a, b) in h.iter().zip(h_prime.iter()) {
        diff |= a ^ b;
    }
    diff == 0
}

// ============================================================================
// Public API -- RSASSA-PSS sign / verify
// ============================================================================

/// RSASSA-PSS sign with a **caller-supplied salt** (RFC 8017 §8.1.1).
///
/// Primarily useful for tests (pinning against external vectors that
/// report a specific salt) and for reproducible signatures in
/// deterministic / audit contexts. Production callers should use
/// [`pss_sign`] with a fresh random salt.
///
/// `m_hash` is the precomputed `H(message)` digest. `H` is the hash
/// function used for both the message digest and the internal
/// MGF1 / `H(M')` -- PSS requires the same hash throughout.
pub fn pss_sign_with_salt<H: Hasher>(sk: &RsaSecretKey, m_hash: &[u8], salt: &[u8]) -> Option<Vec<u8>> {
    let k = sk.modulus_byte_len();
    let mod_bits = sk.n.bit_len();
    let em_bits = mod_bits - 1;
    let em = emsa_pss_encode::<H>(m_hash, em_bits, salt)?;

    // RSASP1: sign by exponentiating EM with the secret exponent.
    let m = BigInt::from_be_bytes(&em);
    let s = rsa_decrypt_raw(sk, &m);
    Some(s.to_be_bytes(k))
}

/// RSASSA-PSS sign of a precomputed digest (RFC 8017 §8.1.1).
///
/// Draws a fresh `s_len`-byte random salt from `rng` and calls
/// [`pss_sign_with_salt`]. Each invocation produces a different
/// signature even for the same `(sk, m_hash)`.
///
/// Recommended salt length: `s_len = H::OUTPUT_LEN` (the same length
/// as the hash). Setting `s_len = 0` produces a deterministic
/// signature ("no salt") -- the same `(sk, m_hash, H)` will always
/// yield the same bytes; this is allowed by the spec but gives up
/// the randomized-signature security property.
pub fn pss_sign<H: Hasher>(
    sk: &RsaSecretKey,
    m_hash: &[u8],
    s_len: usize,
    rng: &mut dyn FnMut(&mut [u8]),
) -> Option<Vec<u8>> {
    let mut salt = vec![0u8; s_len];
    if s_len > 0 {
        rng(&mut salt);
    }
    pss_sign_with_salt::<H>(sk, m_hash, &salt)
}

/// Convenience: hash `msg` with `H`, then sign with a random salt.
pub fn pss_sign_msg<H: Hasher>(
    sk: &RsaSecretKey,
    msg: &[u8],
    s_len: usize,
    rng: &mut dyn FnMut(&mut [u8]),
) -> Option<Vec<u8>> {
    let digest = H::hash(msg);
    pss_sign::<H>(sk, &digest, s_len, rng)
}

/// RSASSA-PSS verify of a precomputed digest (RFC 8017 §8.1.2).
///
/// Returns `true` iff `sig` is a valid PSS signature of `m_hash` under
/// `pk` with the given `s_len`. The caller must supply the same hash
/// `H` and the same salt length the signer used (both are protocol
/// parameters and are typically negotiated out-of-band or fixed by
/// the signature algorithm OID).
pub fn pss_verify<H: Hasher>(pk: &RsaPublicKey, m_hash: &[u8], s_len: usize, sig: &[u8]) -> bool {
    let k = pk.modulus_byte_len();
    if sig.len() != k {
        return false;
    }
    let mod_bits = pk.n.bit_len();
    if mod_bits == 0 {
        return false;
    }
    let em_bits = mod_bits - 1;
    let em_len = (em_bits + 7) / 8;

    // RSAVP1: verify by exponentiating sig with the public exponent.
    let s = BigInt::from_be_bytes(sig);
    let m = rsa_encrypt_raw(pk, &s);
    let em = m.to_be_bytes(em_len);
    if em.len() != em_len {
        // m was too large to fit in em_len bytes. For well-formed
        // signatures this cannot happen because m < n and
        // em_bits = modBits - 1 guarantees m fits in em_len.
        return false;
    }

    emsa_pss_verify::<H>(m_hash, &em, em_bits, s_len)
}

/// Convenience: hash `msg` with `H`, then verify.
pub fn pss_verify_msg<H: Hasher>(pk: &RsaPublicKey, msg: &[u8], s_len: usize, sig: &[u8]) -> bool {
    let digest = H::hash(msg);
    pss_verify::<H>(pk, &digest, s_len, sig)
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::hash::sha256::Sha256;
    use crate::hash::sha384::Sha384;
    use crate::hash::sha512::Sha512;

    /// Deterministic PRNG for reproducible tests (NOT cryptographic).
    fn test_rng() -> impl FnMut(&mut [u8]) {
        let mut state: u64 = 0xdeadbeefcafebabe;
        move |buf: &mut [u8]| {
            for b in buf.iter_mut() {
                state = state
                    .wrapping_mul(6364136223846793005)
                    .wrapping_add(1442695040888963407);
                *b = (state >> 33) as u8;
            }
        }
    }

    // ----------------------------------------------------------------------
    // EMSA-PSS self-consistency (no RSA involved)
    //
    // Encode with a known salt, then verify the resulting EM. This
    // exercises the padding layer independently of the RSA primitive
    // so any failure points straight at the padding code.
    // ----------------------------------------------------------------------

    #[test]
    fn test_emsa_pss_encode_verify_roundtrip_sha256() {
        let m_hash = Sha256::hash(b"hello PSS");
        let salt = [0x5a; 32];
        let em_bits = 2047; // emulates a 2048-bit modulus
        let em = emsa_pss_encode::<Sha256>(&m_hash, em_bits, &salt).expect("encode");
        assert_eq!(em.len(), (em_bits + 7) / 8);
        assert!(emsa_pss_verify::<Sha256>(&m_hash, &em, em_bits, salt.len()));
    }

    #[test]
    fn test_emsa_pss_encode_verify_roundtrip_sha384() {
        let m_hash = Sha384::hash(b"hello PSS sha384");
        let salt = [0x17; 48];
        let em_bits = 3071; // 3072-bit modulus
        let em = emsa_pss_encode::<Sha384>(&m_hash, em_bits, &salt).unwrap();
        assert_eq!(em.len(), 384);
        assert!(emsa_pss_verify::<Sha384>(&m_hash, &em, em_bits, salt.len()));
    }

    #[test]
    fn test_emsa_pss_encode_verify_roundtrip_sha512() {
        let m_hash = Sha512::hash(b"hello PSS sha512");
        let salt = [0x00; 64];
        let em_bits = 4095; // 4096-bit modulus
        let em = emsa_pss_encode::<Sha512>(&m_hash, em_bits, &salt).unwrap();
        assert_eq!(em.len(), 512);
        assert!(emsa_pss_verify::<Sha512>(&m_hash, &em, em_bits, salt.len()));
    }

    /// Salt length must match between encode and verify. Changing it
    /// on verify must reject the encoding.
    #[test]
    fn test_emsa_pss_verify_wrong_salt_length_rejects() {
        let m_hash = Sha256::hash(b"msg");
        let salt = [0xAB; 32];
        let em = emsa_pss_encode::<Sha256>(&m_hash, 2047, &salt).unwrap();
        // Ask verify to expect a DIFFERENT salt length.
        assert!(!emsa_pss_verify::<Sha256>(&m_hash, &em, 2047, 16));
    }

    /// Tampering with the encoded message in the masked-DB region
    /// must be detected.
    #[test]
    fn test_emsa_pss_verify_tampered_rejects() {
        let m_hash = Sha256::hash(b"msg");
        let salt = [0x12; 32];
        let mut em = emsa_pss_encode::<Sha256>(&m_hash, 2047, &salt).unwrap();
        em[0] ^= 0x01;
        assert!(!emsa_pss_verify::<Sha256>(&m_hash, &em, 2047, salt.len()));
    }

    /// The trailing 0xbc byte is mandatory.
    #[test]
    fn test_emsa_pss_verify_missing_bc_byte_rejects() {
        let m_hash = Sha256::hash(b"msg");
        let salt = [0x12; 32];
        let mut em = emsa_pss_encode::<Sha256>(&m_hash, 2047, &salt).unwrap();
        let last = em.len() - 1;
        em[last] = 0xbd;
        assert!(!emsa_pss_verify::<Sha256>(&m_hash, &em, 2047, salt.len()));
    }

    // ----------------------------------------------------------------------
    // Full RSA + PSS roundtrips
    // ----------------------------------------------------------------------

    /// Happy path: sign and verify a message with a random salt.
    #[test]
    fn test_pss_sign_verify_roundtrip_sha256() {
        let mut rng = test_rng();
        let (pk, sk) = super::super::rsa::rsa_keygen(1024, &mut rng);
        let msg = b"PSS end-to-end with SHA-256";
        let sig = pss_sign_msg::<Sha256>(&sk, msg, 32, &mut rng).expect("sign");
        assert_eq!(sig.len(), pk.modulus_byte_len());
        assert!(pss_verify_msg::<Sha256>(&pk, msg, 32, &sig));
    }

    /// Deterministic PSS (s_len = 0) must produce byte-identical
    /// signatures across calls. This is a much sharper self-test of
    /// the whole pipeline than the randomized roundtrip because it
    /// proves the sign path is pure-functional once the salt is fixed.
    #[test]
    fn test_pss_deterministic_signatures_agree() {
        let mut rng = test_rng();
        let (pk, sk) = super::super::rsa::rsa_keygen(1024, &mut rng);
        let m_hash = Sha256::hash(b"determinism");
        let sig1 = pss_sign_with_salt::<Sha256>(&sk, &m_hash, &[]).expect("sign 1");
        let sig2 = pss_sign_with_salt::<Sha256>(&sk, &m_hash, &[]).expect("sign 2");
        assert_eq!(sig1, sig2);
        assert!(pss_verify::<Sha256>(&pk, &m_hash, 0, &sig1));
    }

    /// Two random-salt signatures of the same message must differ
    /// (the salt must actually be consumed).
    #[test]
    fn test_pss_random_signatures_differ() {
        let mut rng = test_rng();
        let (_pk, sk) = super::super::rsa::rsa_keygen(1024, &mut rng);
        let msg = b"randomisation";
        let sig1 = pss_sign_msg::<Sha256>(&sk, msg, 32, &mut rng).unwrap();
        let sig2 = pss_sign_msg::<Sha256>(&sk, msg, 32, &mut rng).unwrap();
        assert_ne!(sig1, sig2);
    }

    /// Verify must reject a signature on a different message.
    #[test]
    fn test_pss_verify_rejects_wrong_message() {
        let mut rng = test_rng();
        let (pk, sk) = super::super::rsa::rsa_keygen(1024, &mut rng);
        let sig = pss_sign_msg::<Sha256>(&sk, b"original", 32, &mut rng).unwrap();
        assert!(!pss_verify_msg::<Sha256>(&pk, b"tampered", 32, &sig));
    }

    /// Verify must reject a tampered signature byte.
    #[test]
    fn test_pss_verify_rejects_tampered_signature() {
        let mut rng = test_rng();
        let (pk, sk) = super::super::rsa::rsa_keygen(1024, &mut rng);
        let msg = b"msg";
        let mut sig = pss_sign_msg::<Sha256>(&sk, msg, 32, &mut rng).unwrap();
        sig[0] ^= 0x01;
        assert!(!pss_verify_msg::<Sha256>(&pk, msg, 32, &sig));
    }

    /// Verify must reject a signature under a different public key.
    #[test]
    fn test_pss_verify_rejects_wrong_key() {
        let mut rng = test_rng();
        let (_pk_a, sk_a) = super::super::rsa::rsa_keygen(1024, &mut rng);
        let (pk_b, _sk_b) = super::super::rsa::rsa_keygen(1024, &mut rng);
        let sig = pss_sign_msg::<Sha256>(&sk_a, b"msg", 32, &mut rng).unwrap();
        assert!(!pss_verify_msg::<Sha256>(&pk_b, b"msg", 32, &sig));
    }

    /// Verify must reject when the declared salt length doesn't match
    /// what the signer used.
    #[test]
    fn test_pss_verify_rejects_wrong_salt_length() {
        let mut rng = test_rng();
        let (pk, sk) = super::super::rsa::rsa_keygen(1024, &mut rng);
        let sig = pss_sign_msg::<Sha256>(&sk, b"msg", 32, &mut rng).unwrap();
        assert!(!pss_verify_msg::<Sha256>(&pk, b"msg", 16, &sig));
    }

    /// Signing with a modulus too small to fit `H(M) || salt || 2`
    /// bytes must return None.
    #[test]
    fn test_pss_sign_rejects_too_small_modulus_for_salt() {
        let mut rng = test_rng();
        // 512-bit modulus = 64 bytes. SHA-256 hLen=32, sLen=32
        // requires emLen >= 32+32+2 = 66 bytes. 64 < 66 -> reject.
        let (_pk, sk) = super::super::rsa::rsa_keygen(512, &mut rng);
        let m_hash = Sha256::hash(b"msg");
        let result = pss_sign::<Sha256>(&sk, &m_hash, 32, &mut rng);
        assert!(result.is_none());
    }

    /// But the same key with a shorter salt works.
    #[test]
    fn test_pss_sign_accepts_short_salt_on_small_modulus() {
        let mut rng = test_rng();
        let (pk, sk) = super::super::rsa::rsa_keygen(512, &mut rng);
        // emLen = 64, hLen = 32 -> max sLen = 64 - 32 - 2 = 30
        let msg = b"small-modulus PSS";
        let sig = pss_sign_msg::<Sha256>(&sk, msg, 16, &mut rng).expect("sign");
        assert!(pss_verify_msg::<Sha256>(&pk, msg, 16, &sig));
    }

    /// Cross-hash: signing and verifying with different hashes must
    /// fail even if both are otherwise valid.
    #[test]
    fn test_pss_hash_mismatch_rejected() {
        let mut rng = test_rng();
        let (pk, sk) = super::super::rsa::rsa_keygen(2048, &mut rng);
        let msg = b"hash flexibility matters";
        let sig = pss_sign_msg::<Sha256>(&sk, msg, 32, &mut rng).unwrap();
        // Same key, same message, but verifying as if it were SHA-384.
        assert!(!pss_verify_msg::<Sha384>(&pk, msg, 32, &sig));
    }

    /// PSS with SHA-384 on a 3072-bit key (NIST recommended pairing).
    #[test]
    fn test_pss_sha384_3072_roundtrip() {
        let mut rng = test_rng();
        let (pk, sk) = super::super::rsa::rsa_keygen(3072, &mut rng);
        let msg = b"PSS SHA-384 / RSA-3072";
        let sig = pss_sign_msg::<Sha384>(&sk, msg, 48, &mut rng).unwrap();
        assert!(pss_verify_msg::<Sha384>(&pk, msg, 48, &sig));
    }
}
