#![forbid(unsafe_code)]

//! SP 800-90A Rev.1 HMAC_DRBG(SHA-512) implementation.

extern crate alloc;

/// Pure Rust, zeroizes state on drop, and enforces a reseed interval.
pub mod rand_adapter;

use alloc::vec::Vec;
use core::{cmp::min, fmt};
use getrandom::getrandom;
use hmac::{Hmac, Mac};
use sha2::Sha512;
use subtle::ConstantTimeEq;
use zeroize::{Zeroize, ZeroizeOnDrop};

/// Errors that can occur during DRBG operation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Error {
    /// Generated output exceeds SP 800-90A per-request cap (64 KiB).
    RequestTooLarge,
    /// Reseed interval exhausted and new entropy is required.
    ReseedRequired,
    /// OS entropy source failed.
    EntropyUnavailable,
    /// Entropy health test failed (SP 800-90B).
    EntropyHealthFailed,
}

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

#[cfg(feature = "std")]
impl std::error::Error for Error {}

/// Default reseed interval recommended by SP 800-90A (Section 10.1.2).
const DEFAULT_RESEED_INTERVAL: u64 = 1u64 << 48;
const BLOCK_LEN: usize = 64;
const MAX_REQUEST: usize = 65536; // 64 KiB cap per generate call.
const DEFAULT_MAX_BYTES_BETWEEN_RESEED: u128 = 1u128 << 20;

/// HMAC_DRBG with HMAC-SHA-512 backbone.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct HmacDrbg {
    k: [u8; BLOCK_LEN],
    v: [u8; BLOCK_LEN],
    reseed_counter: u64,
    reseed_interval: u64,
    generated_bytes: u128,
    max_bytes_between_reseed: u128,
    last_entropy: Vec<u8>,
}

impl HmacDrbg {
    /// Instantiate DRBG from entropy, nonce, and optional personalization string.
    pub fn new(
        entropy: &[u8],
        nonce: &[u8],
        personalization: Option<&[u8]>,
    ) -> Result<Self, Error> {
        Self::validate_entropy(entropy)?;

        let mut seed = Vec::with_capacity(
            entropy.len() + nonce.len() + personalization.map_or(0, |p| p.len()),
        );
        seed.extend_from_slice(entropy);
        seed.extend_from_slice(nonce);
        if let Some(pers) = personalization {
            seed.extend_from_slice(pers);
        }

        let mut drbg = Self {
            k: [0u8; BLOCK_LEN],
            v: [0x01u8; BLOCK_LEN],
            reseed_counter: 1,
            reseed_interval: DEFAULT_RESEED_INTERVAL,
            generated_bytes: 0,
            max_bytes_between_reseed: DEFAULT_MAX_BYTES_BETWEEN_RESEED,
            last_entropy: entropy.to_vec(),
        };
        drbg.update(Some(&seed));
        seed.zeroize();
        Ok(drbg)
    }

    /// Instantiate by sampling entropy+nonce from the operating system.
    pub fn from_os(personalization: Option<&[u8]>) -> Result<Self, Error> {
        let mut entropy = [0u8; 48];
        let mut nonce = [0u8; 16];
        getrandom(&mut entropy).map_err(|_| Error::EntropyUnavailable)?;
        getrandom(&mut nonce).map_err(|_| Error::EntropyUnavailable)?;
        Self::validate_entropy(&nonce)?;
        let drbg = Self::new(&entropy, &nonce, personalization)?;
        entropy.zeroize();
        nonce.zeroize();
        Ok(drbg)
    }

    /// Override the reseed interval (useful for tests or hardened policies).
    pub fn set_reseed_interval(&mut self, interval: u64) {
        self.reseed_interval = interval.max(1);
    }

    /// Override the byte budget requiring reseed.
    pub fn set_max_bytes_between_reseed(&mut self, max_bytes: u128) {
        self.max_bytes_between_reseed = max_bytes.max(1);
    }

    /// Reseed with fresh entropy and optional additional input.
    pub fn reseed(&mut self, entropy: &[u8], additional_input: Option<&[u8]>) -> Result<(), Error> {
        self.check_new_entropy(entropy)?;
        let mut seed = Vec::with_capacity(entropy.len() + additional_input.map_or(0, |a| a.len()));
        seed.extend_from_slice(entropy);
        if let Some(ai) = additional_input {
            seed.extend_from_slice(ai);
        }
        self.update(Some(&seed));
        seed.zeroize();
        self.last_entropy.clear();
        self.last_entropy.extend_from_slice(entropy);
        self.reseed_counter = 1;
        self.generated_bytes = 0;
        Ok(())
    }

    /// Generate output, optionally mixing in additional input (per SP 800-90A Sect. 10.1.2.5).
    pub fn generate(&mut self, out: &mut [u8], additional_input: Option<&[u8]>) -> Result<(), Error> {
        if out.len() > MAX_REQUEST {
            return Err(Error::RequestTooLarge);
        }
        if self.reseed_counter > self.reseed_interval || self.generated_bytes > self.max_bytes_between_reseed
        {
            return Err(Error::ReseedRequired);
        }
        if let Some(ai) = additional_input {
            self.update(Some(ai));
        }
        let mut generated = 0usize;
        while generated < out.len() {
            let mut mac = Hmac::<Sha512>::new_from_slice(&self.k).expect("hmac key len");
            mac.update(&self.v);
            self.v = mac.finalize().into_bytes().into();
            let take = min(out.len() - generated, BLOCK_LEN);
            out[generated..generated + take].copy_from_slice(&self.v[..take]);
            generated += take;
        }
        self.update(additional_input);
        self.reseed_counter = self.reseed_counter.saturating_add(1);
        self.generated_bytes = self.generated_bytes.saturating_add(out.len() as u128);
        Ok(())
    }

    fn update(&mut self, provided_data: Option<&[u8]>) {
        let zero = [0x00u8];
        let one = [0x01u8];

        let mut mac = Hmac::<Sha512>::new_from_slice(&self.k).expect("hmac key len");
        mac.update(&self.v);
        mac.update(&zero);
        if let Some(data) = provided_data {
            mac.update(data);
        }
        self.k = mac.finalize().into_bytes().into();
        let mut mac = Hmac::<Sha512>::new_from_slice(&self.k).expect("hmac key len");
        mac.update(&self.v);
        self.v = mac.finalize().into_bytes().into();

        if provided_data.is_some() {
            let mut mac = Hmac::<Sha512>::new_from_slice(&self.k).expect("hmac key len");
            mac.update(&self.v);
            mac.update(&one);
            if let Some(data) = provided_data {
                mac.update(data);
            }
            self.k = mac.finalize().into_bytes().into();
            let mut mac = Hmac::<Sha512>::new_from_slice(&self.k).expect("hmac key len");
            mac.update(&self.v);
            self.v = mac.finalize().into_bytes().into();
        }
    }

    fn validate_entropy(entropy: &[u8]) -> Result<(), Error> {
        if entropy.len() < 16 {
            return Err(Error::EntropyHealthFailed);
        }
        Ok(())
    }

    fn check_new_entropy(&self, new_entropy: &[u8]) -> Result<(), Error> {
        if new_entropy.len() < 16 {
            return Err(Error::EntropyHealthFailed);
        }
        let prefix = &self.last_entropy[..min(16, self.last_entropy.len())];
        let new_prefix = &new_entropy[..min(16, new_entropy.len())];
        if bool::from(prefix.ct_eq(new_prefix)) {
            return Err(Error::EntropyHealthFailed);
        }
        Ok(())
    }
}

#[cfg(test)]
mod aes_ctr {
    use super::Error;
    use aes::Aes256;
    use ctr::cipher::{KeyIvInit, StreamCipher};
    use zeroize::{Zeroize, ZeroizeOnDrop};

    type Ctr128 = ctr::Ctr128BE<Aes256>;

    /// AES-256-CTR_DRBG scaffold (update/derivation per SP 800-90A still TODO).
    #[derive(Zeroize, ZeroizeOnDrop)]
    pub struct AesCtrDrbg {
        key: [u8; 32],
        v: [u8; 16],
        reseed_counter: u64,
        reseed_interval: u64,
    }

    impl AesCtrDrbg {
        pub fn new(seed_key: &[u8; 32], seed_v: &[u8; 16]) -> Self {
            Self {
                key: *seed_key,
                v: *seed_v,
                reseed_counter: 1,
                reseed_interval: super::DEFAULT_RESEED_INTERVAL,
            }
        }

        pub fn generate(&mut self, out: &mut [u8]) -> Result<(), Error> {
            if self.reseed_counter > self.reseed_interval {
                return Err(Error::ReseedRequired);
            }
            let mut cipher = Ctr128::new((&self.key).into(), (&self.v).into());
            cipher.apply_keystream(out);
            // Increment counter by number of blocks consumed.
            let blocks = (out.len() as u128 + 15) / 16;
            let ctr = u128::from_be_bytes(self.v).wrapping_add(blocks);
            self.v = ctr.to_be_bytes();
            self.reseed_counter = self.reseed_counter.saturating_add(1);
            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn entropy(seed: u8) -> [u8; 48] {
        let mut out = [0u8; 48];
        for (i, byte) in out.iter_mut().enumerate() {
            *byte = seed.wrapping_add(i as u8);
        }
        out
    }

    fn nonce(seed: u8) -> [u8; 16] {
        let mut out = [0u8; 16];
        for (i, byte) in out.iter_mut().enumerate() {
            *byte = seed.wrapping_add((i * 3) as u8);
        }
        out
    }

    #[test]
    fn identical_seed_produces_identical_stream() {
        let mut a = HmacDrbg::new(&entropy(1), &nonce(2), Some(b"p")).unwrap();
        let mut b = HmacDrbg::new(&entropy(1), &nonce(2), Some(b"p")).unwrap();
        let mut out_a = [0u8; 96];
        let mut out_b = [0u8; 96];
        a.generate(&mut out_a, None).unwrap();
        b.generate(&mut out_b, None).unwrap();
        assert_eq!(out_a, out_b);
        assert!(bool::from(a.ct_eq(&b)));
    }

    #[test]
    fn reseed_changes_output() {
        let mut drbg = HmacDrbg::new(&entropy(3), &nonce(4), None).unwrap();
        let mut first = [0u8; 64];
        drbg.generate(&mut first, None).unwrap();
        drbg.reseed(&entropy(9), None).unwrap();
        let mut second = [0u8; 64];
        drbg.generate(&mut second, None).unwrap();
        assert_ne!(first, second);
    }

    #[test]
    fn additional_input_affects_stream() {
        let mut drbg = HmacDrbg::new(&entropy(5), &nonce(6), None).unwrap();
        let mut buf1 = [0u8; 64];
        let mut buf2 = [0u8; 64];
        drbg.generate(&mut buf1, Some(b"ai1")).unwrap();
        drbg.generate(&mut buf2, Some(b"ai2")).unwrap();
        assert_ne!(buf1, buf2);
    }

    #[test]
    fn request_too_large_fails() {
        let mut drbg = HmacDrbg::new(&entropy(7), &nonce(8), None).unwrap();
        let mut buf = vec![0u8; MAX_REQUEST + 1];
        assert_eq!(drbg.generate(&mut buf, None), Err(Error::RequestTooLarge));
    }

    #[test]
    fn reseed_interval_enforced() {
        let mut drbg = HmacDrbg::new(&entropy(9), &nonce(10), None).unwrap();
        drbg.set_reseed_interval(1);
        let mut buf = [0u8; 32];
        drbg.generate(&mut buf, None).unwrap();
        assert_eq!(drbg.generate(&mut buf, None), Err(Error::ReseedRequired));
    }
    #[test]
    fn repeated_entropy_fails_health() {
        let mut drbg = HmacDrbg::new(&entropy(1), &nonce(2), None).unwrap();
        assert!(matches!(
            drbg.reseed(&entropy(1), None),
            Err(Error::EntropyHealthFailed)
        ));
    }

    #[test]
    fn byte_budget_enforced() {
        let mut drbg = HmacDrbg::new(&entropy(11), &nonce(12), None).unwrap();
        drbg.set_max_bytes_between_reseed(64);
        let mut buf = [0u8; 32];
        drbg.generate(&mut buf, None).unwrap();
        assert!(matches!(
            drbg.generate(&mut buf, None),
            Err(Error::ReseedRequired)
        ));
    }
}
