//! HPKE base mode (RFC 9180 §5.1, §6.1) for one suite: DHKEM(X25519, HKDF-SHA256), HKDF-SHA256
//! and AES-128-GCM, the suite `draft-ietf-ppm-dap-20` §7 requires (`telemetry_aggregate`) and the
//! one a remote desktop controller seals its key, text and clipboard actions with
//! (`remote_desktop`, extensions/remote-desktop.md §5.5). Single-shot SealBase and OpenBase, and
//! DeriveKeyPair (§7.1.3), which this module also uses for the sender's ephemeral key: a caller
//! gives fresh random bytes, and a corpus fixed seeds, so that it is written byte for byte.

use aes_gcm::aead::{Aead, KeyInit, Payload};
use aes_gcm::{Aes128Gcm, Nonce};
use curve25519_dalek::montgomery::MontgomeryPoint;
use hkdf::Hkdf;
use sha2::Sha256;

pub const KEM_X25519_HKDF_SHA256: u16 = 0x0020;
pub const KDF_HKDF_SHA256: u16 = 0x0001;
pub const AEAD_AES_128_GCM: u16 = 0x0001;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HpkeError;

fn kem_suite() -> Vec<u8> {
    let mut s = b"KEM".to_vec();
    s.extend_from_slice(&KEM_X25519_HKDF_SHA256.to_be_bytes());
    s
}

fn hpke_suite() -> Vec<u8> {
    let mut s = b"HPKE".to_vec();
    s.extend_from_slice(&KEM_X25519_HKDF_SHA256.to_be_bytes());
    s.extend_from_slice(&KDF_HKDF_SHA256.to_be_bytes());
    s.extend_from_slice(&AEAD_AES_128_GCM.to_be_bytes());
    s
}

fn labeled_extract(suite: &[u8], salt: &[u8], label: &[u8], ikm: &[u8]) -> [u8; 32] {
    let mut m = b"HPKE-v1".to_vec();
    m.extend_from_slice(suite);
    m.extend_from_slice(label);
    m.extend_from_slice(ikm);
    let (prk, _) = Hkdf::<Sha256>::extract(Some(salt), &m);
    prk.into()
}

fn labeled_expand(suite: &[u8], prk: &[u8], label: &[u8], info: &[u8], len: usize) -> Vec<u8> {
    let mut m = (len as u16).to_be_bytes().to_vec();
    m.extend_from_slice(b"HPKE-v1");
    m.extend_from_slice(suite);
    m.extend_from_slice(label);
    m.extend_from_slice(info);
    let mut out = vec![0u8; len];
    Hkdf::<Sha256>::from_prk(prk).expect("PRK length").expand(&m, &mut out).expect("expand length");
    out
}

fn x25519(sk: &[u8; 32], pk: &[u8; 32]) -> Result<[u8; 32], HpkeError> {
    let dh = MontgomeryPoint(*pk).mul_clamped(*sk).to_bytes();
    // RFC 9180 §7.1.4: the all-zero output is refused.
    if dh == [0u8; 32] { Err(HpkeError) } else { Ok(dh) }
}

/// The public key of an X25519 private key.
pub fn public_key(sk: &[u8; 32]) -> [u8; 32] {
    MontgomeryPoint::mul_base_clamped(*sk).to_bytes()
}

/// DeriveKeyPair (§7.1.3) for X25519: (private key, public key).
pub fn derive_key_pair(ikm: &[u8]) -> ([u8; 32], [u8; 32]) {
    let s = kem_suite();
    let prk = labeled_extract(&s, b"", b"dkp_prk", ikm);
    let sk: [u8; 32] = labeled_expand(&s, &prk, b"sk", b"", 32).try_into().unwrap();
    (sk, public_key(&sk))
}

fn extract_and_expand(dh: &[u8], kem_context: &[u8]) -> Vec<u8> {
    let s = kem_suite();
    let prk = labeled_extract(&s, b"", b"eae_prk", dh);
    labeled_expand(&s, &prk, b"shared_secret", kem_context, 32)
}

/// KeySchedule (§5.1) in base mode: the AEAD key and base nonce.
fn key_schedule(shared: &[u8], info: &[u8]) -> (Vec<u8>, Vec<u8>) {
    let s = hpke_suite();
    let psk_id_hash = labeled_extract(&s, b"", b"psk_id_hash", b"");
    let info_hash = labeled_extract(&s, b"", b"info_hash", info);
    let mut ctx = vec![0u8];
    ctx.extend_from_slice(&psk_id_hash);
    ctx.extend_from_slice(&info_hash);
    let secret = labeled_extract(&s, shared, b"secret", b"");
    (labeled_expand(&s, &secret, b"key", &ctx, 16), labeled_expand(&s, &secret, b"base_nonce", &ctx, 12))
}

/// SealBase with the ephemeral key derived from `ikm_e`: (enc, ciphertext).
pub fn seal_base(pk_r: &[u8; 32], info: &[u8], aad: &[u8], pt: &[u8], ikm_e: &[u8]) -> Result<(Vec<u8>, Vec<u8>), HpkeError> {
    let (sk_e, pk_e) = derive_key_pair(ikm_e);
    let dh = x25519(&sk_e, pk_r)?;
    let mut kem_context = pk_e.to_vec();
    kem_context.extend_from_slice(pk_r);
    let shared = extract_and_expand(&dh, &kem_context);
    let (key, nonce) = key_schedule(&shared, info);
    let c = Aes128Gcm::new_from_slice(&key).map_err(|_| HpkeError)?;
    let ct = c.encrypt(Nonce::from_slice(&nonce), Payload { msg: pt, aad }).map_err(|_| HpkeError)?;
    Ok((pk_e.to_vec(), ct))
}

/// OpenBase.
pub fn open_base(enc: &[u8], sk_r: &[u8; 32], info: &[u8], aad: &[u8], ct: &[u8]) -> Result<Vec<u8>, HpkeError> {
    let pk_e: [u8; 32] = enc.try_into().map_err(|_| HpkeError)?;
    let dh = x25519(sk_r, &pk_e)?;
    let mut kem_context = pk_e.to_vec();
    kem_context.extend_from_slice(&public_key(sk_r));
    let shared = extract_and_expand(&dh, &kem_context);
    let (key, nonce) = key_schedule(&shared, info);
    let c = Aes128Gcm::new_from_slice(&key).map_err(|_| HpkeError)?;
    c.decrypt(Nonce::from_slice(&nonce), Payload { msg: ct, aad }).map_err(|_| HpkeError)
}

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

    fn h(s: &str) -> Vec<u8> {
        (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap()).collect()
    }

    /// RFC 9180 Appendix A.1.1 and the first encryption of A.1.1.1.
    #[test]
    fn rfc9180_a_1_1() {
        let (sk_e, pk_e) = derive_key_pair(&h("7268600d403fce431561aef583ee1613527cff655c1343f29812e66706df3234"));
        assert_eq!(sk_e.to_vec(), h("52c4a758a802cd8b936eceea314432798d5baf2d7e9235dc084ab1b9cfa2f736"));
        assert_eq!(pk_e.to_vec(), h("37fda3567bdbd628e88668c3c8d7e97d1d1253b6d4ea6d44c150f741f1bf4431"));
        let (sk_r, pk_r) = derive_key_pair(&h("6db9df30aa07dd42ee5e8181afdb977e538f5e1fec8a06223f33f7013e525037"));
        assert_eq!(sk_r.to_vec(), h("4612c550263fc8ad58375df3f557aac531d26850903e55a9f23f21d8534e8ac8"));
        assert_eq!(pk_r.to_vec(), h("3948cfe0ad1ddb695d780e59077195da6c56506b027329794ab02bca80815c4d"));
        let info = h("4f6465206f6e2061204772656369616e2055726e");
        let pt = h("4265617574792069732074727574682c20747275746820626561757479");
        let aad = h("436f756e742d30");
        let (enc, ct) = seal_base(&pk_r, &info, &aad, &pt, &h("7268600d403fce431561aef583ee1613527cff655c1343f29812e66706df3234")).unwrap();
        assert_eq!(enc, pk_e.to_vec());
        assert_eq!(ct, h("f938558b5d72f1a23810b4be2ab4f84331acc02fc97babc53a52ae8218a355a96d8770ac83d07bea87e13c512a"));
        assert_eq!(open_base(&enc, &sk_r, &info, &aad, &ct).unwrap(), pt);
        let mut bad = ct.clone();
        bad[3] ^= 1;
        assert!(open_base(&enc, &sk_r, &info, &aad, &bad).is_err());
        assert!(open_base(&enc, &sk_r, &info, b"other", &ct).is_err());
    }
}
