//! Wrapping of token secrets at rest.
//!
//! Every `token_key` is stored encrypted with a server key, with the `token_id`
//! as additional data. A stolen `.db` therefore yields no working keys, and a
//! wrapped blob cannot be moved from one token row to another.
//!
//! The server refuses to start without a key, and accepts an old one for a
//! one-shot rotation.

use anyhow::{Context, Result, bail};
use chacha20poly1305::aead::{Aead, KeyInit, Payload};
use chacha20poly1305::{ChaCha20Poly1305, Nonce};
use hkdf::Hkdf;
use rand::TryRngCore;
use rand::rngs::OsRng;
use sha2::Sha256;

const KEY_LEN: usize = 32;
const NONCE_LEN: usize = 12;

/// HKDF label separating the revocation master from the wrapping key. They come
/// from the same secret and must never be the same value.
const REVOCATION_MASTER_INFO: &[u8] = b"otp/1/revoke-master";

/// The server's key-wrapping key(s).
pub struct KeyVault {
    current: ChaCha20Poly1305,
    /// Accepted for unwrapping only, so a rotation can re-wrap lazily.
    previous: Option<ChaCha20Poly1305>,
    /// Master for per-token revocation keys. Derived from the same secret rather
    /// than configured separately: one required environment variable is enough,
    /// and an operator who has to manage two will eventually lose one.
    ///
    /// The consequence is that rotating `OT_SECRET_KEY` also invalidates every
    /// `K_rev` already issued. Devices that logged in beforehand stop being able
    /// to verify a revocation notice and fall back to silence until their next
    /// login — the behaviour they would have had anyway without this mechanism.
    revocation_master: otproto::Key,
}

impl KeyVault {
    /// Reads `OT_SECRET_KEY` (base64, 32 bytes) or the file it names, plus the
    /// optional `OT_SECRET_KEY_OLD`.
    pub fn from_env() -> Result<Self> {
        let current = load("OT_SECRET_KEY")?.ok_or_else(|| {
            anyhow::anyhow!(
                "OT_SECRET_KEY is not set. It must be 32 random bytes, base64-encoded, or the path \
                 to a 0600 file containing them. Generate one with:\n  \
                 head -c32 /dev/urandom | base64\n\
                 Without it, token keys would sit in the database in the clear."
            )
        })?;
        let previous = load("OT_SECRET_KEY_OLD")?;
        Ok(Self {
            revocation_master: derive_revocation_master(&current),
            current: ChaCha20Poly1305::new((&current).into()),
            previous: previous.map(|k| ChaCha20Poly1305::new((&k).into())),
        })
    }

    #[cfg(test)]
    pub fn for_test(key: [u8; KEY_LEN]) -> Self {
        Self {
            revocation_master: derive_revocation_master(&key),
            current: ChaCha20Poly1305::new((&key).into()),
            previous: None,
        }
    }

    /// The master from which every `K_rev` is derived.
    ///
    /// Handed to [`crate::ingest::Ingest`] so the hot path can seal a notice for
    /// a `token_id` it has never seen, without a database round trip and without
    /// storing anything per token.
    pub fn revocation_master(&self) -> otproto::Key {
        self.revocation_master
    }

    /// The key sealing revocation notices for one token. Also what the login
    /// response hands the device.
    pub fn revocation_key(&self, token_id: u64) -> otproto::Key {
        otproto::revocation_key(&self.revocation_master, token_id)
    }

    /// `nonce || ciphertext || tag`, with the token id as AAD.
    pub fn wrap(&self, token_id: u64, token_key: &[u8; KEY_LEN]) -> Result<Vec<u8>> {
        let mut nonce = [0u8; NONCE_LEN];
        OsRng.try_fill_bytes(&mut nonce).context("OS RNG failed")?;
        let sealed = self
            .current
            .encrypt(
                Nonce::from_slice(&nonce),
                Payload {
                    msg: token_key,
                    aad: &token_id.to_be_bytes(),
                },
            )
            .map_err(|_| anyhow::anyhow!("wrapping token key failed"))?;
        let mut out = Vec::with_capacity(NONCE_LEN + sealed.len());
        out.extend_from_slice(&nonce);
        out.extend_from_slice(&sealed);
        Ok(out)
    }

    pub fn unwrap(&self, token_id: u64, blob: &[u8]) -> Result<[u8; KEY_LEN]> {
        if blob.len() < NONCE_LEN + 16 {
            bail!(
                "wrapped key for token {token_id} is truncated ({} bytes)",
                blob.len()
            );
        }
        let (nonce, sealed) = blob.split_at(NONCE_LEN);
        let aad = token_id.to_be_bytes();

        for cipher in [Some(&self.current), self.previous.as_ref()]
            .into_iter()
            .flatten()
        {
            if let Ok(plain) = cipher.decrypt(
                Nonce::from_slice(nonce),
                Payload {
                    msg: sealed,
                    aad: &aad,
                },
            ) {
                return plain.try_into().map_err(|v: Vec<u8>| {
                    anyhow::anyhow!("token key is {} bytes, want {KEY_LEN}", v.len())
                });
            }
        }
        bail!(
            "cannot unwrap the key for token {token_id}: neither OT_SECRET_KEY nor \
             OT_SECRET_KEY_OLD decrypts it"
        )
    }
}

fn load(var: &str) -> Result<Option<[u8; KEY_LEN]>> {
    let Ok(raw) = std::env::var(var) else {
        return Ok(None);
    };
    if raw.is_empty() {
        return Ok(None);
    }

    // A path is more likely than base64 to contain a '/', so decide on whether
    // the value names an existing file rather than on its shape.
    let text = if std::path::Path::new(&raw).is_file() {
        std::fs::read_to_string(&raw).with_context(|| format!("{var}: reading {raw}"))?
    } else {
        raw
    };

    use base64::Engine as _;
    let bytes = base64::engine::general_purpose::STANDARD
        .decode(text.trim())
        .with_context(|| format!("{var} is not valid base64"))?;
    if bytes.len() != KEY_LEN {
        bail!("{var} decodes to {} bytes, want {KEY_LEN}", bytes.len());
    }
    Ok(Some(bytes.try_into().expect("length checked")))
}

/// 32 fresh random bytes for a new token secret.
pub fn random_token_key() -> Result<[u8; KEY_LEN]> {
    let mut k = [0u8; KEY_LEN];
    OsRng.try_fill_bytes(&mut k).context("OS RNG failed")?;
    Ok(k)
}

/// A random, non-zero `token_id`.
///
/// Random rather than sequential because the id travels in cleartext in every
/// datagram header: a guessable one would let an attacker enumerate which tokens
/// exist by watching for the absence of a reply.
pub fn random_token_id() -> Result<u64> {
    loop {
        let mut b = [0u8; 8];
        OsRng.try_fill_bytes(&mut b).context("OS RNG failed")?;
        let id = u64::from_be_bytes(b);
        // 0 is reserved as "unset" in a few places; rejecting it costs nothing.
        if id != 0 {
            return Ok(id);
        }
    }
}

fn derive_revocation_master(secret: &[u8; KEY_LEN]) -> otproto::Key {
    let hk = Hkdf::<Sha256>::from_prk(secret).expect("32-byte PRK is valid for HKDF-SHA256");
    let mut out = [0u8; KEY_LEN];
    hk.expand(REVOCATION_MASTER_INFO, &mut out)
        .expect("32 bytes is well under HKDF-SHA256's output limit");
    out
}

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

    #[test]
    fn wrap_then_unwrap_round_trips() {
        let vault = KeyVault::for_test([7; KEY_LEN]);
        let key = [0x42; KEY_LEN];
        let blob = vault.wrap(99, &key).expect("wrap");
        assert_eq!(vault.unwrap(99, &blob).expect("unwrap"), key);
    }

    #[test]
    fn a_blob_cannot_be_moved_to_another_token() {
        // The token id is AAD, so a row-swap in the database is detected rather
        // than silently cloning a credential onto another token.
        let vault = KeyVault::for_test([7; KEY_LEN]);
        let blob = vault.wrap(99, &[0x42; KEY_LEN]).expect("wrap");
        assert!(vault.unwrap(100, &blob).is_err());
    }

    #[test]
    fn a_tampered_blob_is_rejected() {
        let vault = KeyVault::for_test([7; KEY_LEN]);
        let mut blob = vault.wrap(1, &[1; KEY_LEN]).expect("wrap");
        let last = blob.len() - 1;
        blob[last] ^= 1;
        assert!(vault.unwrap(1, &blob).is_err());
    }

    #[test]
    fn a_wrong_server_key_cannot_unwrap() {
        let blob = KeyVault::for_test([7; KEY_LEN])
            .wrap(1, &[1; KEY_LEN])
            .expect("wrap");
        assert!(KeyVault::for_test([8; KEY_LEN]).unwrap(1, &blob).is_err());
    }

    #[test]
    fn truncated_blobs_fail_with_a_clear_error() {
        let vault = KeyVault::for_test([7; KEY_LEN]);
        assert!(vault.unwrap(1, &[]).is_err());
        assert!(vault.unwrap(1, &[0; NONCE_LEN]).is_err());
    }

    #[test]
    fn wrapping_is_randomised() {
        // Same key, same token, different ciphertext: the nonce is fresh each
        // time, so the database never reveals that two tokens share a secret.
        let vault = KeyVault::for_test([7; KEY_LEN]);
        let a = vault.wrap(1, &[1; KEY_LEN]).expect("wrap");
        let b = vault.wrap(1, &[1; KEY_LEN]).expect("wrap");
        assert_ne!(a, b);
    }

    #[test]
    fn token_ids_are_never_zero() {
        for _ in 0..100 {
            assert_ne!(random_token_id().expect("id"), 0);
        }
    }
}