oxirs-did 0.4.0

//! Key agreement protocols for DID-based communication.
//!
//! This module provides:
//!
//! * [`EcdhKeyAgreement`] — **real** X25519 (Curve25519) Elliptic-Curve
//!   Diffie–Hellman backed by `curve25519-dalek`. Both parties provably derive
//!   the same 32-byte shared secret, which is the intended mechanism for
//!   DID-based communication.
//! * [`KeyAgreementProtocol`] — a small-prime textbook Diffie–Hellman over
//!   `u64` modular arithmetic. This is a teaching/testing toy and is **NOT**
//!   cryptographically secure; use [`EcdhKeyAgreement`] for anything real.
//! * `hkdf_extract` / `hkdf_expand` — a lightweight XOR-chain key-derivation
//!   helper (**not** HMAC-SHA-256 HKDF; not for production key stretching).
//!
//! # Example
//!
//! ```rust
//! use oxirs_did::key_agreement::{DhParams, KeyAgreementProtocol};
//!
//! let protocol = KeyAgreementProtocol::new(DhParams::standard());
//!
//! // Alice and Bob choose private keys
//! let alice_kp = protocol.generate_keypair(6);
//! let bob_kp   = protocol.generate_keypair(15);
//!
//! // They exchange public keys and compute the shared secret
//! let alice_secret = protocol.compute_shared_secret(alice_kp.private_key, bob_kp.public_key);
//! let bob_secret   = protocol.compute_shared_secret(bob_kp.private_key, alice_kp.public_key);
//!
//! assert_eq!(alice_secret.value, bob_secret.value);
//! ```

use crate::{DidError, DidResult};

/// Diffie–Hellman parameter set (prime and generator)
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DhParams {
    /// A prime modulus
    pub prime: u64,
    /// A primitive root (generator) modulo `prime`
    pub generator: u64,
}

impl DhParams {
    /// Small-prime parameters for unit tests (prime=23, generator=5).
    /// **Not secure for production use.**
    pub fn standard() -> Self {
        Self {
            prime: 23,
            generator: 5,
        }
    }

    /// Larger parameters using the Mersenne prime 2^31 − 1 = 2 147 483 647
    /// with generator 7.  Still **not production-secure** but provides more
    /// room for key values in integration tests.
    pub fn large() -> Self {
        Self {
            prime: 2_147_483_647,
            generator: 7,
        }
    }
}

/// A Diffie–Hellman key pair
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KeyPair {
    /// The private scalar chosen by the local party
    pub private_key: u64,
    /// The public value `generator^private_key mod prime`
    pub public_key: u64,
}

/// A shared DH secret derived by both parties
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SharedSecret {
    /// The agreed value `other_public^private_key mod prime`
    pub value: u64,
}

/// A real X25519 (Curve25519) key pair.
///
/// `private_bytes` is a 32-byte secret scalar seed (clamped on use, per
/// RFC 7748); `public_bytes` is the corresponding Montgomery-u public key,
/// `X25519_BASEPOINT * clamp(private_bytes)`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct X25519KeyPair {
    /// 32-byte X25519 private key seed.
    pub private_bytes: [u8; 32],
    /// 32-byte X25519 public key (Montgomery-u coordinate).
    pub public_bytes: [u8; 32],
}

/// DH-based key agreement protocol engine
#[derive(Debug, Clone)]
pub struct KeyAgreementProtocol {
    /// The DH parameters this protocol instance uses
    pub params: DhParams,
}

impl KeyAgreementProtocol {
    /// Create a new protocol instance with the given parameters
    pub fn new(params: DhParams) -> Self {
        Self { params }
    }

    /// Derive a key pair from the given private key.
    ///
    /// `public_key = generator^private_key mod prime`
    pub fn generate_keypair(&self, private_key: u64) -> KeyPair {
        let public_key = Self::modpow(self.params.generator, private_key, self.params.prime);
        KeyPair {
            private_key,
            public_key,
        }
    }

    /// Compute the shared secret: `other_public^private_key mod prime`
    pub fn compute_shared_secret(&self, private_key: u64, other_public: u64) -> SharedSecret {
        SharedSecret {
            value: Self::modpow(other_public, private_key, self.params.prime),
        }
    }

    /// Fast binary modular exponentiation.
    ///
    /// Computes `base^exp mod modulus` in O(log exp) multiplications.
    /// Returns 1 if `modulus` is 0 to avoid division by zero.
    pub fn modpow(base: u64, exp: u64, modulus: u64) -> u64 {
        if modulus == 0 {
            return 1;
        }
        if modulus == 1 {
            return 0;
        }
        let mut result: u128 = 1;
        let mut b = (base % modulus) as u128;
        let mut e = exp;
        let m = modulus as u128;
        while e > 0 {
            if e & 1 == 1 {
                result = result * b % m;
            }
            b = b * b % m;
            e >>= 1;
        }
        result as u64
    }

    /// Encode the public key as a `did:key` URI with base58btc multibase prefix.
    ///
    /// Format: `did:key:z<base58(big-endian bytes of public)>`
    pub fn key_to_did_key_format(public: u64) -> String {
        let bytes = public.to_be_bytes();
        // Strip leading zero bytes for a compact representation
        let significant: Vec<u8> = bytes.iter().copied().skip_while(|&b| b == 0).collect();
        let encoded = if significant.is_empty() {
            "1".to_string() // base58 encoding of zero
        } else {
            base58_encode(&significant)
        };
        format!("did:key:z{}", encoded)
    }
}

/// Real X25519 (Curve25519) Elliptic-Curve Diffie–Hellman key agreement.
///
/// Backed by `curve25519-dalek`. Keys are generated from OS entropy and the
/// shared secret is the standard X25519 function, so Alice and Bob both compute
/// the identical secret from their own private key and the peer's public key.
///
/// The `curve` label is retained for API compatibility; the cryptographic
/// operation is always X25519 regardless of its value.
#[derive(Debug, Clone)]
pub struct EcdhKeyAgreement {
    /// Curve label (informational; the operation is always X25519).
    pub curve: String,
}

impl EcdhKeyAgreement {
    /// Create a new ECDH agreement for the named curve.
    pub fn new(curve: &str) -> Self {
        Self {
            curve: curve.to_string(),
        }
    }

    /// Generate a fresh X25519 key pair from a cryptographically secure RNG.
    ///
    /// Takes `&mut self` for backward-compatible call sites, but no longer keeps
    /// any deterministic internal state — each call draws fresh OS entropy.
    ///
    /// Fails closed if the OS entropy source is unavailable.
    pub fn generate_keypair(&mut self) -> DidResult<X25519KeyPair> {
        Self::generate_keypair_static()
    }

    /// Generate a fresh X25519 key pair (does not require `&mut self`).
    ///
    /// Fails closed if the OS entropy source is unavailable.
    pub fn generate_keypair_static() -> DidResult<X25519KeyPair> {
        use curve25519_dalek::constants::X25519_BASEPOINT;

        // 32-byte X25519 private seed from the OS CSPRNG (oxicrypto-rand →
        // getrandom). Any 32-byte value is a valid X25519 scalar (clamped on use).
        let private_bytes = oxicrypto_rand::random_nonce::<32>()
            .map_err(|e| DidError::InternalError(format!("OS entropy source failed: {e}")))?;
        // Public key = clamp(private) * basepoint (mul_clamped applies the
        // RFC 7748 clamping internally).
        let public_bytes = X25519_BASEPOINT.mul_clamped(private_bytes).to_bytes();
        Ok(X25519KeyPair {
            private_bytes,
            public_bytes,
        })
    }

    /// Derive the X25519 shared secret from the local private key and the
    /// remote public key.
    ///
    /// This is real ECDH: `derive_shared_secret(alice_priv, bob_pub)` equals
    /// `derive_shared_secret(bob_priv, alice_pub)`.
    pub fn derive_shared_secret(
        &self,
        local: &X25519KeyPair,
        remote_public: &[u8; 32],
    ) -> [u8; 32] {
        use curve25519_dalek::montgomery::MontgomeryPoint;
        MontgomeryPoint(*remote_public)
            .mul_clamped(local.private_bytes)
            .to_bytes()
    }
}

/// HKDF Extract phase — produces a pseudo-random key (PRK) from `salt` and `ikm`.
///
/// Uses a simplified XOR-chain construction (not HMAC-SHA-256).
pub fn hkdf_extract(salt: &[u8], ikm: &[u8]) -> Vec<u8> {
    // Combine salt and IKM into a 32-byte PRK using XOR-chain mixing
    let mut state = [0u8; 32];

    // Absorb salt
    for (i, &b) in salt.iter().enumerate() {
        state[i % 32] ^= b;
        // Diffuse by rotation
        state[(i + 1) % 32] = state[(i + 1) % 32].wrapping_add(b.rotate_left(3));
    }

    // Absorb IKM
    for (i, &b) in ikm.iter().enumerate() {
        state[i % 32] ^= b;
        state[(i + 3) % 32] = state[(i + 3) % 32].wrapping_add(b.rotate_right(5));
    }

    // Final mixing pass
    for i in 0..32 {
        state[i] = state[i].wrapping_add(state[(i + 7) % 32]).rotate_left(1);
    }

    state.to_vec()
}

/// HKDF Expand phase — expands the PRK `prk` with context `info` to `length` bytes.
///
/// Uses a counter-based XOR-chain construction (not HMAC-SHA-256 counter mode).
pub fn hkdf_expand(prk: &[u8], info: &[u8], length: usize) -> Vec<u8> {
    if length == 0 || prk.is_empty() {
        return vec![0u8; length];
    }

    let mut output = Vec::with_capacity(length);
    let mut t = Vec::new(); // T(0) = empty

    let mut counter: u8 = 1;
    while output.len() < length {
        let mut block = Vec::new();
        block.extend_from_slice(&t);
        block.extend_from_slice(info);
        block.push(counter);

        // Mix PRK into block using XOR-chain
        let mut state = [0u8; 32];
        let prk_cycle: Vec<u8> = prk
            .iter()
            .copied()
            .cycle()
            .take(32.max(block.len()))
            .collect();
        for (i, &b) in block.iter().enumerate() {
            state[i % 32] ^= b ^ prk_cycle[i % prk_cycle.len()];
            state[(i + 1) % 32] = state[(i + 1) % 32].wrapping_add(b.rotate_left(2));
        }
        // Final diffusion
        for i in 0..32 {
            state[i] = state[i].wrapping_add(state[(i + 5) % 32]).rotate_right(3);
        }

        t = state.to_vec();
        output.extend_from_slice(&t);
        counter = counter.wrapping_add(1);
    }

    output.truncate(length);
    output
}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

/// Minimal base58 encoder (Bitcoin alphabet).
fn base58_encode(bytes: &[u8]) -> String {
    const ALPHABET: &[u8] = b"123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";

    let mut digits: Vec<u8> = Vec::new();
    for &byte in bytes {
        let mut carry = byte as u32;
        for d in digits.iter_mut() {
            carry += (*d as u32) << 8;
            *d = (carry % 58) as u8;
            carry /= 58;
        }
        while carry > 0 {
            digits.push((carry % 58) as u8);
            carry /= 58;
        }
    }

    // Leading zero bytes become '1' characters
    let leading_ones = bytes.iter().take_while(|&&b| b == 0).count();
    let mut result = String::with_capacity(leading_ones + digits.len());
    for _ in 0..leading_ones {
        result.push('1');
    }
    for &d in digits.iter().rev() {
        result.push(ALPHABET[d as usize] as char);
    }
    result
}

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

    // --- DhParams ---

    #[test]
    fn test_standard_params() {
        let p = DhParams::standard();
        assert_eq!(p.prime, 23);
        assert_eq!(p.generator, 5);
    }

    #[test]
    fn test_large_params() {
        let p = DhParams::large();
        assert_eq!(p.prime, 2_147_483_647);
        assert_eq!(p.generator, 7);
    }

    // --- modpow ---

    #[test]
    fn test_modpow_basic() {
        // 5^6 mod 23 = 8
        assert_eq!(KeyAgreementProtocol::modpow(5, 6, 23), 8);
    }

    #[test]
    fn test_modpow_zero_exp() {
        assert_eq!(KeyAgreementProtocol::modpow(5, 0, 23), 1);
    }

    #[test]
    fn test_modpow_exp_one() {
        assert_eq!(KeyAgreementProtocol::modpow(5, 1, 23), 5);
    }

    #[test]
    fn test_modpow_modulus_one() {
        assert_eq!(KeyAgreementProtocol::modpow(100, 50, 1), 0);
    }

    #[test]
    fn test_modpow_large() {
        // Python: pow(7, 1_000_000, 2_147_483_647)
        let result = KeyAgreementProtocol::modpow(7, 1_000_000, 2_147_483_647);
        assert!(result < 2_147_483_647);
    }

    // --- generate_keypair ---

    #[test]
    fn test_generate_keypair_public_in_range() {
        let protocol = KeyAgreementProtocol::new(DhParams::standard());
        let kp = protocol.generate_keypair(6);
        assert!(kp.public_key < 23);
    }

    #[test]
    fn test_generate_keypair_private_preserved() {
        let protocol = KeyAgreementProtocol::new(DhParams::standard());
        let kp = protocol.generate_keypair(15);
        assert_eq!(kp.private_key, 15);
    }

    // --- Key agreement (shared secret symmetry) ---

    #[test]
    fn test_shared_secret_symmetry_standard() {
        let protocol = KeyAgreementProtocol::new(DhParams::standard());
        let alice = protocol.generate_keypair(6);
        let bob = protocol.generate_keypair(15);
        let alice_ss = protocol.compute_shared_secret(alice.private_key, bob.public_key);
        let bob_ss = protocol.compute_shared_secret(bob.private_key, alice.public_key);
        assert_eq!(alice_ss.value, bob_ss.value);
    }

    #[test]
    fn test_shared_secret_symmetry_large_params() {
        let protocol = KeyAgreementProtocol::new(DhParams::large());
        let alice = protocol.generate_keypair(123_456_789);
        let bob = protocol.generate_keypair(987_654_321);
        let alice_ss = protocol.compute_shared_secret(alice.private_key, bob.public_key);
        let bob_ss = protocol.compute_shared_secret(bob.private_key, alice.public_key);
        assert_eq!(alice_ss.value, bob_ss.value);
    }

    #[test]
    fn test_shared_secret_different_private_keys_may_differ() {
        let protocol = KeyAgreementProtocol::new(DhParams::standard());
        // Two completely independent pairs should generally yield different secrets
        let alice = protocol.generate_keypair(6);
        let carol = protocol.generate_keypair(3);
        let mallory = protocol.generate_keypair(11);
        let alice_ss = protocol.compute_shared_secret(alice.private_key, carol.public_key);
        let wrong_ss = protocol.compute_shared_secret(mallory.private_key, carol.public_key);
        // With these parameters Alice and Mallory have different shared secrets with Carol
        // (not always guaranteed for all combinations, but true here)
        let _ = alice_ss;
        let _ = wrong_ss;
        // Just verify no panic and values are in range
        assert!(alice_ss.value < 23);
        assert!(wrong_ss.value < 23);
    }

    // --- did:key format ---

    #[test]
    fn test_key_to_did_key_format_starts_with_prefix() {
        let uri = KeyAgreementProtocol::key_to_did_key_format(8);
        assert!(uri.starts_with("did:key:z"), "URI = {}", uri);
    }

    #[test]
    fn test_key_to_did_key_format_zero() {
        let uri = KeyAgreementProtocol::key_to_did_key_format(0);
        assert!(uri.starts_with("did:key:z"), "URI for 0 = {}", uri);
    }

    #[test]
    fn test_key_to_did_key_format_large_value() {
        let uri = KeyAgreementProtocol::key_to_did_key_format(u64::MAX);
        assert!(uri.starts_with("did:key:z"), "URI = {}", uri);
    }

    #[test]
    fn test_key_to_did_key_format_different_values_differ() {
        let u1 = KeyAgreementProtocol::key_to_did_key_format(8);
        let u2 = KeyAgreementProtocol::key_to_did_key_format(9);
        assert_ne!(u1, u2);
    }

    // --- EcdhKeyAgreement ---

    #[test]
    fn test_ecdh_generate_keypair_length() {
        let mut ecdh = EcdhKeyAgreement::new("X25519");
        let kp = ecdh.generate_keypair().expect("X25519 keygen");
        assert_eq!(kp.private_bytes.len(), 32);
        assert_eq!(kp.public_bytes.len(), 32);
    }

    #[test]
    fn test_ecdh_successive_keypairs_differ() {
        let mut ecdh = EcdhKeyAgreement::new("X25519");
        let kp1 = ecdh.generate_keypair().expect("X25519 keygen");
        let kp2 = ecdh.generate_keypair().expect("X25519 keygen");
        // Counter-based seeds → different output
        assert_ne!(kp1.public_bytes, kp2.public_bytes);
    }

    #[test]
    fn test_ecdh_derive_shared_secret_length() {
        let mut ecdh = EcdhKeyAgreement::new("X25519");
        let alice = ecdh.generate_keypair().expect("X25519 keygen");
        let bob = ecdh.generate_keypair().expect("X25519 keygen");
        let shared = ecdh.derive_shared_secret(&alice, &bob.public_bytes);
        assert_eq!(shared.len(), 32);
    }

    #[test]
    fn test_ecdh_derive_shared_secret_symmetry() {
        // Real X25519 ECDH: both parties derive the SAME shared secret.
        let mut ecdh = EcdhKeyAgreement::new("X25519");
        let alice = ecdh.generate_keypair().expect("X25519 keygen");
        let bob = ecdh.generate_keypair().expect("X25519 keygen");
        let s1 = ecdh.derive_shared_secret(&alice, &bob.public_bytes);
        let s2 = ecdh.derive_shared_secret(&bob, &alice.public_bytes);
        assert_eq!(s1, s2, "X25519 ECDH must be symmetric");
        // A real shared secret is not all-zero for independent keypairs.
        assert!(s1.iter().any(|&b| b != 0));
    }

    #[test]
    fn test_ecdh_distinct_pairs_yield_distinct_secrets() {
        // An attacker (mallory) with their own keypair cannot derive alice&bob's
        // shared secret.
        let alice = EcdhKeyAgreement::generate_keypair_static().expect("X25519 keygen");
        let bob = EcdhKeyAgreement::generate_keypair_static().expect("X25519 keygen");
        let mallory = EcdhKeyAgreement::generate_keypair_static().expect("X25519 keygen");
        let ecdh = EcdhKeyAgreement::new("X25519");

        let ab = ecdh.derive_shared_secret(&alice, &bob.public_bytes);
        let mallory_with_bob = ecdh.derive_shared_secret(&mallory, &bob.public_bytes);
        assert_ne!(
            ab, mallory_with_bob,
            "third party must not recover the Alice/Bob secret"
        );
    }

    #[test]
    fn test_ecdh_public_key_matches_basepoint_mul() {
        // The generated public key really is basepoint * clamp(private).
        use curve25519_dalek::constants::X25519_BASEPOINT;
        let kp = EcdhKeyAgreement::generate_keypair_static().expect("X25519 keygen");
        let expected = X25519_BASEPOINT.mul_clamped(kp.private_bytes).to_bytes();
        assert_eq!(kp.public_bytes, expected);
    }

    #[test]
    fn test_ecdh_curve_name_stored() {
        let ecdh = EcdhKeyAgreement::new("P-256");
        assert_eq!(ecdh.curve, "P-256");
    }

    // --- hkdf_extract ---

    #[test]
    fn test_hkdf_extract_output_length() {
        let prk = hkdf_extract(b"salt", b"input_key_material");
        assert_eq!(prk.len(), 32);
    }

    #[test]
    fn test_hkdf_extract_different_salts_differ() {
        let prk1 = hkdf_extract(b"salt1", b"ikm");
        let prk2 = hkdf_extract(b"salt2", b"ikm");
        assert_ne!(prk1, prk2);
    }

    #[test]
    fn test_hkdf_extract_different_ikm_differ() {
        let prk1 = hkdf_extract(b"salt", b"ikm1");
        let prk2 = hkdf_extract(b"salt", b"ikm2");
        assert_ne!(prk1, prk2);
    }

    #[test]
    fn test_hkdf_extract_deterministic() {
        let p1 = hkdf_extract(b"s", b"k");
        let p2 = hkdf_extract(b"s", b"k");
        assert_eq!(p1, p2);
    }

    #[test]
    fn test_hkdf_extract_empty_inputs() {
        let prk = hkdf_extract(b"", b"");
        assert_eq!(prk.len(), 32);
    }

    // --- hkdf_expand ---

    #[test]
    fn test_hkdf_expand_correct_length() {
        let prk = hkdf_extract(b"salt", b"ikm");
        let okm = hkdf_expand(&prk, b"info", 64);
        assert_eq!(okm.len(), 64);
    }

    #[test]
    fn test_hkdf_expand_zero_length() {
        let prk = hkdf_extract(b"salt", b"ikm");
        let okm = hkdf_expand(&prk, b"info", 0);
        assert!(okm.is_empty());
    }

    #[test]
    fn test_hkdf_expand_different_info_differ() {
        let prk = hkdf_extract(b"salt", b"ikm");
        let okm1 = hkdf_expand(&prk, b"info1", 32);
        let okm2 = hkdf_expand(&prk, b"info2", 32);
        assert_ne!(okm1, okm2);
    }

    #[test]
    fn test_hkdf_expand_deterministic() {
        let prk = hkdf_extract(b"salt", b"ikm");
        let o1 = hkdf_expand(&prk, b"context", 48);
        let o2 = hkdf_expand(&prk, b"context", 48);
        assert_eq!(o1, o2);
    }

    #[test]
    fn test_hkdf_full_pipeline() {
        let prk = hkdf_extract(b"shared_secret_salt", b"dh_shared_value");
        let session_key = hkdf_expand(&prk, b"session_key_v1", 32);
        assert_eq!(session_key.len(), 32);
        // Verify it's not all zeros
        assert!(session_key.iter().any(|&b| b != 0));
    }

    // --- base58 helper ---

    #[test]
    fn test_base58_encode_non_empty() {
        let encoded = base58_encode(&[1, 2, 3, 4]);
        assert!(!encoded.is_empty());
    }

    #[test]
    fn test_base58_encode_only_alphabet_chars() {
        const ALPHA: &str = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";
        let encoded = base58_encode(&[255, 128, 64]);
        for ch in encoded.chars() {
            assert!(ALPHA.contains(ch), "Invalid base58 char: {}", ch);
        }
    }
}