//! HMAC-DRBG — deterministic random bit generator (NIST SP 800-90A), generic
//! over any [`Digest`].

use super::{CryptoRng, RngCore};
use crate::hash::{Digest, Hmac};

/// HMAC-DRBG instantiated with hash function `D`.
///
/// Seed it with [`new`](HmacDrbg::new) from a high-entropy source (e.g.
/// [`OsRng`](super::OsRng)), then draw bytes via [`RngCore`]. Given the same
/// seed it is fully deterministic, which is what makes it testable against the
/// NIST known-answer vectors.
///
/// SP 800-90A §10.1.2.4 mandates reseeding before the counter exceeds `2^48`
/// (the "reseed_interval"). This implementation enforces the limit: once
/// reached, [`generate`](Self::generate) panics until [`reseed`](Self::reseed)
/// is called. Per SP 800-90A the panic-vs-error choice is implementation
/// defined; panic is the conservative behaviour for a CSPRNG.
#[derive(Clone)]
pub struct HmacDrbg<D: Digest> {
    /// HMAC key.
    k: D::Output,
    /// Internal value updated on each output block.
    v: D::Output,
    reseed_counter: u64,
}

/// SP 800-90A §10.1.2.4 reseed interval for HMAC-DRBG: `2^48` requests.
const RESEED_INTERVAL: u64 = 1u64 << 48;

impl<D: Digest> HmacDrbg<D> {
    /// Instantiates the DRBG from `entropy`, a `nonce`, and an optional
    /// `personalization` string (any may be empty, though entropy should not
    /// be).
    pub fn new(entropy: &[u8], nonce: &[u8], personalization: &[u8]) -> Self {
        // K = 0x00…, V = 0x01…
        let k = D::zeroed_output();
        let mut v = D::zeroed_output();
        for b in v.as_mut() {
            *b = 0x01;
        }
        let mut drbg = HmacDrbg {
            k,
            v,
            reseed_counter: 1,
        };
        drbg.update(&[entropy, nonce, personalization]);
        drbg
    }

    /// The SP 800-90A `HMAC_DRBG_Update` step over the concatenation of
    /// `provided`.
    fn update(&mut self, provided: &[&[u8]]) {
        // K = HMAC(K, V || 0x00 || provided)
        let mut mac = Hmac::<D>::new(self.k.as_ref());
        mac.update(self.v.as_ref());
        mac.update(&[0x00]);
        for p in provided {
            mac.update(p);
        }
        self.k = mac.finalize();
        // V = HMAC(K, V)
        self.v = Hmac::<D>::mac(self.k.as_ref(), self.v.as_ref());

        // The second pass runs only when there was provided data.
        if provided.iter().any(|p| !p.is_empty()) {
            let mut mac = Hmac::<D>::new(self.k.as_ref());
            mac.update(self.v.as_ref());
            mac.update(&[0x01]);
            for p in provided {
                mac.update(p);
            }
            self.k = mac.finalize();
            self.v = Hmac::<D>::mac(self.k.as_ref(), self.v.as_ref());
        }
    }

    /// Reseeds with fresh `entropy` and optional `additional` input.
    pub fn reseed(&mut self, entropy: &[u8], additional: &[u8]) {
        self.update(&[entropy, additional]);
        self.reseed_counter = 1;
    }

    /// Fills `out` with generated bytes, optionally mixing in `additional`
    /// input.
    ///
    /// # Panics
    /// Panics if the per-instance request count would exceed SP 800-90A's
    /// reseed interval (`2^48`). Call [`reseed`](Self::reseed) before that
    /// happens.
    pub fn generate(&mut self, out: &mut [u8], additional: &[u8]) {
        assert!(
            self.reseed_counter < RESEED_INTERVAL,
            "HMAC-DRBG reseed interval exceeded (SP 800-90A §10.1.2.4); call reseed() first"
        );
        if !additional.is_empty() {
            self.update(&[additional]);
        }

        let mut filled = 0;
        while filled < out.len() {
            self.v = Hmac::<D>::mac(self.k.as_ref(), self.v.as_ref());
            let block = self.v.as_ref();
            let n = (out.len() - filled).min(block.len());
            out[filled..filled + n].copy_from_slice(&block[..n]);
            filled += n;
        }

        self.update(&[additional]);
        self.reseed_counter += 1;
    }
}

impl<D: Digest> RngCore for HmacDrbg<D> {
    #[inline]
    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.generate(dest, &[]);
    }
}

impl<D: Digest> CryptoRng for HmacDrbg<D> {}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::hash::Sha256;
    use crate::test_util::from_hex;

    #[test]
    fn nist_hmac_drbg_sha256() {
        // NIST CAVP HMAC_DRBG SHA-256, no reseed / no prediction resistance,
        // empty personalization & additional input, COUNT 0. Two 1024-bit
        // generate calls; the second is the known answer.
        let entropy =
            from_hex::<32>("ca851911349384bffe89de1cbdc46e6831e44d34a4fb935ee285dd14b71a7488");
        let nonce = from_hex::<16>("659ba96c601dc69fc902940805ec0ca8");
        let expected = from_hex::<128>(
            "e528e9abf2dece54d47c7e75e5fe302149f817ea9fb4bee6f4199697d04d5b89\
             d54fbb978a15b5c443c9ec21036d2460b6f73ebad0dc2aba6e624abf07745bc1\
             07694bb7547bb0995f70de25d6b29e2d3011bb19d27676c07162c8b5ccde0668\
             961df86803482cb37ed6d5c0bb8d50cf1f50d476aa0458bdaba806f48be9dcb8",
        );

        let mut drbg = HmacDrbg::<Sha256>::new(&entropy, &nonce, &[]);
        let mut buf = [0u8; 128];
        drbg.generate(&mut buf, &[]); // discarded
        drbg.generate(&mut buf, &[]);
        assert_eq!(buf, expected);
    }

    #[test]
    fn deterministic_and_seed_sensitive() {
        let mk = |seed: &[u8]| {
            let mut d = HmacDrbg::<Sha256>::new(seed, b"nonce", &[]);
            let mut o = [0u8; 48];
            d.fill_bytes(&mut o);
            o
        };
        assert_eq!(mk(b"seed-a"), mk(b"seed-a"));
        assert_ne!(mk(b"seed-a"), mk(b"seed-b"));
    }

    #[test]
    fn multiblock_generate_is_deterministic() {
        // 100 bytes spans four SHA-256 output blocks, exercising block
        // stitching within a single request; two equally seeded generators
        // must agree byte-for-byte.
        let mut a = HmacDrbg::<Sha256>::new(b"entropy-source", b"nonce", &[]);
        let mut b = HmacDrbg::<Sha256>::new(b"entropy-source", b"nonce", &[]);
        let mut whole = [0u8; 100];
        let mut whole2 = [0u8; 100];
        a.fill_bytes(&mut whole);
        b.fill_bytes(&mut whole2);
        assert_ne!(whole, [0u8; 100]);
        assert_eq!(whole, whole2);
    }

    #[test]
    fn reseed_changes_stream() {
        let mut a = HmacDrbg::<Sha256>::new(b"seed", b"nonce", &[]);
        let mut b = HmacDrbg::<Sha256>::new(b"seed", b"nonce", &[]);
        b.reseed(b"more-entropy", &[]);
        let (mut x, mut y) = ([0u8; 32], [0u8; 32]);
        a.fill_bytes(&mut x);
        b.fill_bytes(&mut y);
        assert_ne!(x, y);
    }
}
