onepass-seed 0.3.2

use core::fmt::Write;
use std::io::{self, Cursor, Error, Result};

use argon2::{Algorithm, Argon2, Params, Version};
use blake2::{Blake2b256, Digest};
use chacha20::ChaCha20Rng;
use crypto_bigint::{NonZero, RandomBits, RandomMod, U256};
use onepass_base::fmt::DigestWriter;
use rand_core::SeedableRng;
use secrecy::{ExposeSecret, ExposeSecretMut, SecretBox, SecretString};

use crate::{expr::Eval, site::Site};

impl Site<'_> {
    /// Write this site’s password into the passed [`io::Write`] implementation. For security, `W`
    /// should write to a buffer that will be zeroed as soon as possible.
    pub fn write_password_into<W>(&self, w: &mut W, seed_password: &str) -> Result<()>
    where
        W: io::Write,
    {
        let size = self.expr.size();
        let secret = self.secret(seed_password);
        let mut index = secret_uniform(&secret, &size);
        self.expr.write_to(w, &mut index)?;
        Ok(())
    }

    /// Return this site’s unique password for the given `seed_password`.
    pub fn password(&self, seed_password: &str) -> Result<SecretString> {
        // Write to a pre-allocated buffer to prevent reallocations leaking sensitive data.
        let mut buf = SecretBox::from(vec![0u8; 4096]);
        let mut cursor = Cursor::new(buf.expose_secret_mut());
        self.write_password_into(&mut cursor, seed_password)?;
        let pos = usize::try_from(cursor.position()).unwrap();
        let buf = &cursor.into_inner()[..pos];
        let s = str::from_utf8(buf).map_err(Error::other)?;
        Ok(SecretString::from(s))
    }

    /// Return the public salt corresponding to this site’s derivation parameters.
    /// This is just `BLAKE2B256(derivation)`.
    pub fn salt(&self) -> [u8; 32] {
        let mut w = DigestWriter(Blake2b256::new());
        write!(w, "{self}").unwrap();
        w.0.finalize().into()
    }

    /// Return the per-site secret for the given `seed_password`, running [`Argon2`] with the
    /// crate parameters. The parameters are 256MiB memory, 4 iterations, 4 parallelism.
    pub fn secret(&self, seed_password: &str) -> SecretBox<[u8; 32]> {
        // NB. m_cost is measured in KiB.
        let params = Params::new(256 * 1024, 4, 4, None).unwrap();
        let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, params);
        let salt = self.salt();
        SecretBox::init_with_mut(|out: &mut [u8; 32]| {
            argon2
                .hash_password_into(seed_password.as_bytes(), &salt, out)
                .unwrap();
        })
    }
}

/// Randomly sample a [`U256`] from the given 256-bit secret. Uses rejection sampling to prevent
/// bias in the results.
fn secret_uniform(secret: &dyn ExposeSecret<[u8; 32]>, n: &NonZero<U256>) -> SecretBox<U256> {
    let n_bits = n.bits_vartime();
    if n_bits == 1 {
        return SecretBox::default();
    }

    let mut rng = ChaCha20Rng::from_seed(*secret.expose_secret());
    SecretBox::init_with(|| {
        if n.trailing_zeros_vartime() == n_bits - 1 {
            // For powers of 2, we do not need rejection-sampling.
            // We can simply generate `n_bits - 1` random bits.
            RandomBits::random_bits(&mut rng, n_bits - 1)
        } else {
            RandomMod::random_mod_vartime(&mut rng, n)
        }
    })
}

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

    fn test_site() -> Site<'static> {
        Site::new("google.com", None, "{words}", 0).unwrap()
    }

    #[test]
    fn derivation_works() {
        assert_eq!(
            "v3/priv\thttps://google.com/\t\t{words|323606b363ebdedff9f562cb84c50df1a21cbd4b597ff4566df92bb9f2cefdfd}\t0",
            &format!("{}", &test_site())
        );
    }

    #[test]
    fn salt_works() {
        assert_eq!(
            "1cef595ae1bd08b50d1fd457eab716961891b8e6244f377bea6f18cc171f749d",
            hex::encode(test_site().salt())
        );
        let mut site2 = test_site();
        site2.username = Some("me@example.com".into());
        assert_eq!(
            "05801ea917cb4da2045b9db3bff6b7421dc93f1d3708bfd34a73100f16460bb3",
            hex::encode(site2.salt())
        );
    }

    #[test]
    #[ignore] // too slow in debug
    fn secret() {
        assert_eq!(
            "b9d8aeffbcf60b4054d399be576648e1a058d3b61f194a5fab73126362b3a301",
            hex::encode(test_site().secret("testpass").expose_secret())
        );
        assert_eq!(
            "92ea6c4c44a3cb223e601ade213bbcfdf55c2758997c8657631e7dd33ffe0ed2",
            hex::encode(test_site().secret("testpass2").expose_secret())
        );
        let mut site2 = test_site();
        site2.increment = 1;
        site2.username = Some("you@example.com".into());
        assert_eq!(
            "d76fbab456c845b005aa8527781197a8691a763984d05e75450d266bc6f1cd27",
            hex::encode(*site2.secret("testpass").expose_secret())
        );
    }

    #[test]
    fn secret_uniform_short() {
        let tests = [(1, 0x3c5), (2, 0xf6a), (3, 0x180), (4, 0x390), (5, 0x19d)];
        for (seed, want) in tests {
            let secret = SecretBox::init_with(|| U256::from_u32(seed).to_le_bytes().into());
            let n = NonZero::new(U256::from_u32(0x1000)).unwrap();
            assert_eq!(
                U256::from_u32(want),
                *secret_uniform(&secret, &n).expose_secret()
            );
        }
    }

    #[test]
    fn secret_uniform_vectors() {
        let tests: [(&str, &str, &str); _] = [
            (
                "0000000000000000000000000000000000000000000000000000000000000000",
                "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF",
                "C70D778BCCEF36A81AED8DA0B819D2BD28BD8653E56A5D40903DF1A0ADE0B876",
            ),
            (
                "0123456789abcdeffedcba98765432100123456789abcdeffedcba9876543210",
                "0000000000000000000000000000000000000000000000000000000000100000",
                "000000000000000000000000000000000000000000000000000000000005D415",
            ),
            (
                "0123456789abcdeffedcba98765432100123456789abcdeffedcba9876543210",
                "295A7969D28101E13473A8DD15E68D28CCD4F578591D8994008C5D999F85D416",
                "295A7969D28101E13473A8DD15E68D28CCD4F578591D8994008C5D999F85D415",
            ),
            (
                "0123456789abcdeffedcba98765432100123456789abcdeffedcba9876543210",
                "295A7969D28101E13473A8DD15E68D28CCD4F578591D8994008C5D999F85D415",
                "0D313C0A2DDB1AE37A6EF3ECC18F8588FB946C5BE4A31B39784D7C9530E31D51",
            ),
            (
                "0000000000000000000000000000000000000000000000000000000000000000",
                "0000000000000000000000000000000000000000000000000000000000000001",
                "0000000000000000000000000000000000000000000000000000000000000000",
            ),
            (
                "a96d610f969d8befcc5a8f7db635976eeb5c83718a2a0d9974a4bb1b6423fac9",
                "00000000000000001FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF",
                "00000000000000000B6CE7C37CBAA4C1133D97B36751CCE9AA56B264F9E8698D",
            ),
        ];
        for (sec, siz, want) in tests {
            let sec = SecretBox::init_with(|| U256::from_be_hex(sec).to_be_bytes().into());
            assert_eq!(
                U256::from_be_hex(want),
                *secret_uniform(&sec, &NonZero::new(U256::from_be_hex(siz)).unwrap())
                    .expose_secret(),
            );
        }
    }

    #[test]
    #[ignore]
    fn password_e2e() {
        assert_eq!(
            "parasitic prompter dimmer overdrive designer",
            &*test_site().password("testpass").unwrap().expose_secret(),
        );
    }
}