oxirs-did 0.2.4

//! Key agreement protocols for DID-based communication.
//!
//! This module implements a Diffie–Hellman key agreement protocol using
//! modular arithmetic, along with a simulated X25519-style ECDH, HKDF
//! key derivation, and a helper to format public keys as `did:key` URIs.
//!
//! **Security Notice**: The implementations here use simplified arithmetic
//! (small prime DH, XOR-based ECDH simulation, XOR-chain HKDF) and are
//! **NOT cryptographically secure**.  They are provided for demonstration,
//! testing, and protocol prototyping purposes only.  For production use,
//! replace these primitives with a proper cryptographic library such as
//! `ring`, `dalek-cryptography`, or `rustls`.
//!
//! # 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);
//! ```

/// 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 simulated X25519 key pair (32-byte arrays)
///
/// The bytes here are generated by a deterministic XorShift and are
/// **not** real Curve25519 scalars or points.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct X25519KeyPair {
    /// 32-byte private key (simulated)
    pub private_bytes: [u8; 32],
    /// 32-byte public key (simulated)
    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)
    }
}

/// Simulated ECDH key agreement for named elliptic curves ("X25519" or "P-256")
#[derive(Debug, Clone)]
pub struct EcdhKeyAgreement {
    /// Curve name ("X25519" or "P-256")
    pub curve: String,
    /// Internal counter for deterministic key generation
    counter: u64,
}

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

    /// Generate a simulated key pair using a deterministic XorShift PRNG
    /// seeded from 42 XOR-ed with the internal counter.
    pub fn generate_keypair(&mut self) -> X25519KeyPair {
        let seed = 42u64 ^ self.counter;
        self.counter = self.counter.wrapping_add(1);
        let private_bytes = xorshift_bytes(seed);
        let public_bytes = xorshift_bytes(seed ^ 0xDEAD_BEEF_CAFE_1234);
        X25519KeyPair {
            private_bytes,
            public_bytes,
        }
    }

    /// Derive a simulated shared secret by XOR-ing the local private key
    /// with the remote public key.
    ///
    /// **Not cryptographically secure** — for protocol simulation only.
    pub fn derive_shared_secret(
        &self,
        local: &X25519KeyPair,
        remote_public: &[u8; 32],
    ) -> [u8; 32] {
        let mut shared = [0u8; 32];
        for i in 0..32 {
            shared[i] = local.private_bytes[i] ^ remote_public[i];
        }
        shared
    }
}

/// 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
// ---------------------------------------------------------------------------

/// Generate 32 pseudo-random bytes from a 64-bit XorShift PRNG seed.
fn xorshift_bytes(seed: u64) -> [u8; 32] {
    let mut state = if seed == 0 {
        0x123456789ABCDEF0u64
    } else {
        seed
    };
    let mut bytes = [0u8; 32];
    for chunk in bytes.chunks_mut(8) {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        for (i, b) in state.to_le_bytes().iter().enumerate() {
            if i < chunk.len() {
                chunk[i] = *b;
            }
        }
    }
    bytes
}

/// 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();
        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();
        let kp2 = ecdh.generate_keypair();
        // 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();
        let bob = ecdh.generate_keypair();
        let shared = ecdh.derive_shared_secret(&alice, &bob.public_bytes);
        assert_eq!(shared.len(), 32);
    }

    #[test]
    fn test_ecdh_derive_shared_secret_symmetry() {
        // XOR-based: XOR(alice_priv, bob_pub) and XOR(bob_priv, alice_pub)
        // are NOT the same (XOR ECDH is not symmetric in this simplified form)
        // — we only verify it doesn't panic and produces 32 bytes.
        let mut ecdh = EcdhKeyAgreement::new("P-256");
        let alice = ecdh.generate_keypair();
        let bob = ecdh.generate_keypair();
        let s1 = ecdh.derive_shared_secret(&alice, &bob.public_bytes);
        let s2 = ecdh.derive_shared_secret(&bob, &alice.public_bytes);
        assert_eq!(s1.len(), 32);
        assert_eq!(s2.len(), 32);
    }

    #[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);
        }
    }
}