sb-mesh 0.1.1

use chacha20poly1305::{
    aead::{Aead, KeyInit},
    ChaCha20Poly1305, Nonce,
};
use sha2::{Digest, Sha256};
use x25519_dalek::{PublicKey, StaticSecret};

pub const PACKET_MAGIC: [u8; 4] = [b'S', b'B', b'M', b'1'];

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PacketType {
    HandshakeInit = 1,
    HandshakeResp = 2,
    Ping = 3,
    Pong = 4,
    Data = 5,
    RekeyInit = 6,
    RekeyResp = 7,
}

impl PacketType {
    pub fn from_u8(val: u8) -> Option<Self> {
        match val {
            1 => Some(PacketType::HandshakeInit),
            2 => Some(PacketType::HandshakeResp),
            3 => Some(PacketType::Ping),
            4 => Some(PacketType::Pong),
            5 => Some(PacketType::Data),
            6 => Some(PacketType::RekeyInit),
            7 => Some(PacketType::RekeyResp),
            _ => None,
        }
    }
}

/// Compute X25519 Diffie-Hellman shared key and derive 256-bit symmetric session key via SHA-256
pub fn derive_session_key(local_secret: &StaticSecret, remote_public: &PublicKey) -> [u8; 32] {
    let shared = local_secret.diffie_hellman(remote_public);
    let mut hasher = Sha256::new();
    hasher.update(shared.as_bytes());
    hasher.update(b"SBM_CHACHA20_POLY1305_SESSION_V1");
    let result = hasher.finalize();
    let mut key = [0u8; 32];
    key.copy_from_slice(&result);
    key
}

/// Build standard 12-byte AEAD nonce from a 64-bit monotonically increasing sequence number
pub fn build_nonce(counter: u64) -> Nonce {
    let mut nonce_bytes = [0u8; 12];
    nonce_bytes[4..12].copy_from_slice(&counter.to_be_bytes());
    *Nonce::from_slice(&nonce_bytes)
}

/// Encrypt payload using ChaCha20-Poly1305 with 64-bit anti-replay sequence number
pub fn encrypt_payload(
    key: &[u8; 32],
    counter: u64,
    plaintext: &[u8],
) -> Result<Vec<u8>, String> {
    let cipher = ChaCha20Poly1305::new_from_slice(key)
        .map_err(|e| format!("Invalid cipher key: {}", e))?;
    let nonce = build_nonce(counter);
    cipher
        .encrypt(&nonce, plaintext)
        .map_err(|e| format!("Encryption failure: {}", e))
}

/// Decrypt payload using ChaCha20-Poly1305 with 64-bit anti-replay sequence number
pub fn decrypt_payload(
    key: &[u8; 32],
    counter: u64,
    ciphertext: &[u8],
) -> Result<Vec<u8>, String> {
    let cipher = ChaCha20Poly1305::new_from_slice(key)
        .map_err(|e| format!("Invalid cipher key: {}", e))?;
    let nonce = build_nonce(counter);
    cipher
        .decrypt(&nonce, ciphertext)
        .map_err(|e| format!("Authentication/Decryption failure: {}", e))
}

/// Encode a complete WireGuard-style packet:
/// [0..4]: Magic (SBM1)
/// [4]: PacketType
/// [5..37]: Sender Public Key (32 bytes)
/// [37..45]: Nonce / Counter (8 bytes, Big-Endian)
/// [45..]: Encrypted payload + Poly1305 Tag (16 bytes)
pub fn pack_message(
    packet_type: PacketType,
    sender_pubkey: &[u8; 32],
    counter: u64,
    encrypted_payload: &[u8],
) -> Vec<u8> {
    let mut buf = Vec::with_capacity(45 + encrypted_payload.len());
    buf.extend_from_slice(&PACKET_MAGIC);
    buf.push(packet_type as u8);
    buf.extend_from_slice(sender_pubkey);
    buf.extend_from_slice(&counter.to_be_bytes());
    buf.extend_from_slice(encrypted_payload);
    buf
}

/// Parse and validate header of an incoming packet
pub fn unpack_header(packet: &[u8]) -> Result<(PacketType, [u8; 32], u64, &[u8]), String> {
    if packet.len() < 45 {
        return Err("Packet too short (< 45 bytes header)".to_string());
    }
    if packet[0..4] != PACKET_MAGIC {
        return Err("Invalid packet magic header".to_string());
    }
    let ptype = PacketType::from_u8(packet[4])
        .ok_or_else(|| format!("Unknown packet type: {}", packet[4]))?;

    let mut sender_pubkey = [0u8; 32];
    sender_pubkey.copy_from_slice(&packet[5..37]);

    let mut counter_bytes = [0u8; 8];
    counter_bytes.copy_from_slice(&packet[37..45]);
    let counter = u64::from_be_bytes(counter_bytes);

    let payload = &packet[45..];
    Ok((ptype, sender_pubkey, counter, payload))
}

use zeroize::{Zeroize, ZeroizeOnDrop};

/// Protected in-memory key that securely zeros its buffer upon drop
#[derive(Clone, PartialEq, Eq, Zeroize, ZeroizeOnDrop)]
pub struct ZeroizingKey(pub [u8; 32]);

impl ZeroizingKey {
    pub fn new(bytes: [u8; 32]) -> Self {
        Self(bytes)
    }

    pub fn as_bytes(&self) -> &[u8; 32] {
        &self.0
    }
}

impl std::fmt::Debug for ZeroizingKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "ZeroizingKey([REDACTED])")
    }
}

impl std::ops::Deref for ZeroizingKey {
    type Target = [u8; 32];
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Derives a protected session key wrapped in a ZeroizingKey structure
pub fn derive_session_zeroizing_key(local_secret: &StaticSecret, remote_public: &PublicKey) -> ZeroizingKey {
    ZeroizingKey::new(derive_session_key(local_secret, remote_public))
}

/// KDF Ratchet step: Derives the next epoch key from current key and ephemeral Diffie-Hellman secret.
/// Advances epoch and ensures Perfect Forward Secrecy (PFS).
pub fn ratchet_forward(current_key: &[u8; 32], ephemeral_shared: &[u8; 32], epoch: u64) -> ZeroizingKey {
    let mut hasher = Sha256::new();
    hasher.update(b"SBM_PFS_RATCHET_FORWARD_V1");
    hasher.update(current_key);
    hasher.update(ephemeral_shared);
    hasher.update(&epoch.to_be_bytes());
    let result = hasher.finalize();
    let mut new_key = [0u8; 32];
    new_key.copy_from_slice(&result);
    ZeroizingKey(new_key)
}

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

    #[test]
    fn test_zeroizing_key_clears_memory() {
        let mut key = ZeroizingKey::new([0x42; 32]);
        assert_eq!(key[0], 0x42);
        key.zeroize();
        assert_eq!(key.0, [0u8; 32]);
    }

    #[test]
    fn test_ratchet_forward_pfs() {
        let base_key = [0x11u8; 32];
        let eph_shared = [0x22u8; 32];

        let epoch1_key = ratchet_forward(&base_key, &eph_shared, 1);
        let epoch2_key = ratchet_forward(&epoch1_key, &eph_shared, 2);

        assert_ne!(&epoch1_key[..], &base_key[..]);
        assert_ne!(&epoch2_key[..], &epoch1_key[..]);

        // Deterministic ratchet with same inputs
        let verify_key = ratchet_forward(&base_key, &eph_shared, 1);
        assert_eq!(&epoch1_key[..], &verify_key[..]);
    }
}