sb-mesh 0.1.1

use chrono::Utc;
use curve25519_dalek::constants::RISTRETTO_BASEPOINT_POINT;
use curve25519_dalek::ristretto::CompressedRistretto;
use curve25519_dalek::scalar::Scalar;
use rand::rngs::OsRng;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha512};
use std::collections::HashMap;

pub const ACL_DOMAIN_SEPARATOR: &[u8] = b"SBM_CAPABILITY_TOKEN_V1";

/// Decentralized Capability Token granting specific port permissions
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CapabilityToken {
    pub issuer_node_id: String,
    pub subject_node_id: String,
    pub allowed_ports: Vec<u16>,
    pub valid_until: u64,
    pub authority_pubkey: [u8; 32],
    pub signature_r: [u8; 32],
    pub signature_s: [u8; 32],
}

impl CapabilityToken {
    pub fn digest(&self) -> [u8; 64] {
        let mut hasher = Sha512::new();
        hasher.update(ACL_DOMAIN_SEPARATOR);
        hasher.update(self.issuer_node_id.as_bytes());
        hasher.update(self.subject_node_id.as_bytes());
        for &port in &self.allowed_ports {
            hasher.update(&port.to_be_bytes());
        }
        hasher.update(&self.valid_until.to_be_bytes());
        let result = hasher.finalize();
        let mut out = [0u8; 64];
        out.copy_from_slice(&result);
        out
    }

    /// Issues and signs a new CapabilityToken using the issuer's private scalar key
    pub fn issue(
        issuer_node_id: &str,
        subject_node_id: &str,
        allowed_ports: Vec<u16>,
        valid_until: u64,
        authority_secret: &Scalar,
    ) -> Self {
        let authority_pub = (authority_secret * RISTRETTO_BASEPOINT_POINT).compress().to_bytes();

        let mut dummy = Self {
            issuer_node_id: issuer_node_id.to_string(),
            subject_node_id: subject_node_id.to_string(),
            allowed_ports,
            valid_until,
            authority_pubkey: authority_pub,
            signature_r: [0u8; 32],
            signature_s: [0u8; 32],
        };

        let digest = dummy.digest();

        // Schnorr signature
        let k = Scalar::random(&mut OsRng);
        let commitment_r = (k * RISTRETTO_BASEPOINT_POINT).compress().to_bytes();

        let mut chal_hasher = Sha512::new();
        chal_hasher.update(ACL_DOMAIN_SEPARATOR);
        chal_hasher.update(&commitment_r);
        chal_hasher.update(&authority_pub);
        chal_hasher.update(&digest);
        let chal_hash: [u8; 64] = chal_hasher.finalize().into();
        let challenge = Scalar::from_bytes_mod_order_wide(&chal_hash);

        let response_s = k + challenge * authority_secret;

        dummy.signature_r = commitment_r;
        dummy.signature_s = response_s.to_bytes();
        dummy
    }

    /// Verifies the cryptographic Schnorr signature and validity epoch of the token
    pub fn verify(&self, current_epoch: u64) -> Result<(), String> {
        if current_epoch > self.valid_until {
            return Err("Capability token expired".to_string());
        }

        let compressed_x = CompressedRistretto(self.authority_pubkey);
        let x_point = compressed_x
            .decompress()
            .ok_or_else(|| "Invalid authority public key point".to_string())?;

        let compressed_r = CompressedRistretto(self.signature_r);
        let r_point = compressed_r
            .decompress()
            .ok_or_else(|| "Invalid commitment R point".to_string())?;

        let s_scalar: Option<Scalar> = Scalar::from_canonical_bytes(self.signature_s).into();
        let s_scalar = match s_scalar {
            Some(s) => s,
            None => return Err("Invalid scalar S".to_string()),
        };

        let digest = self.digest();

        let mut chal_hasher = Sha512::new();
        chal_hasher.update(ACL_DOMAIN_SEPARATOR);
        chal_hasher.update(&self.signature_r);
        chal_hasher.update(&self.authority_pubkey);
        chal_hasher.update(&digest);
        let chal_hash: [u8; 64] = chal_hasher.finalize().into();
        let challenge = Scalar::from_bytes_mod_order_wide(&chal_hash);

        let lhs = s_scalar * RISTRETTO_BASEPOINT_POINT;
        let rhs = r_point + challenge * x_point;

        if lhs == rhs {
            Ok(())
        } else {
            Err("Signature verification failed".to_string())
        }
    }
}

/// Zero-Trust Capability-based Access Control Engine
#[derive(Debug, Clone)]
pub struct AclEngine {
    // Subject Node ID -> List of CapabilityTokens
    tokens: HashMap<String, Vec<CapabilityToken>>,
    pub default_allow: bool,
}

impl AclEngine {
    pub fn new(default_allow: bool) -> Self {
        Self {
            tokens: HashMap::new(),
            default_allow,
        }
    }

    /// Grants and stores a verified capability token
    pub fn grant(&mut self, token: CapabilityToken) -> Result<(), String> {
        let now = Utc::now().timestamp() as u64;
        token.verify(now)?;
        self.tokens
            .entry(token.subject_node_id.clone())
            .or_default()
            .push(token);
        Ok(())
    }

    /// Checks whether subject node has permission to access the target port
    pub fn is_port_allowed(&self, subject_node_id: &str, port: u16, current_epoch: u64) -> bool {
        if let Some(token_list) = self.tokens.get(subject_node_id) {
            for token in token_list {
                if token.verify(current_epoch).is_ok() {
                    if token.allowed_ports.is_empty() || token.allowed_ports.contains(&port) {
                        return true;
                    }
                }
            }
            false
        } else {
            self.default_allow
        }
    }
}

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

    #[test]
    fn test_capability_token_signing_and_verification() {
        let authority_secret = Scalar::random(&mut OsRng);
        let now = Utc::now().timestamp() as u64;

        let token = CapabilityToken::issue(
            "sbm-0xissuer00",
            "sbm-0xsubject1",
            vec![22, 8080],
            now + 3600,
            &authority_secret,
        );

        assert!(token.verify(now).is_ok());

        // Tampering test: altered port
        let mut tampered = token.clone();
        tampered.allowed_ports.push(5432);
        assert!(tampered.verify(now).is_err());

        // Expiration test
        assert!(token.verify(now + 4000).is_err());
    }

    #[test]
    fn test_acl_engine_port_enforcement() {
        let authority_secret = Scalar::random(&mut OsRng);
        let now = Utc::now().timestamp() as u64;

        let token = CapabilityToken::issue(
            "sbm-0xvault",
            "sbm-0xclient",
            vec![22, 8080], // Only SSH and HTTP allowed
            now + 3600,
            &authority_secret,
        );

        let mut engine = AclEngine::new(false); // Default deny
        engine.grant(token).expect("Grant failed");

        // Allowed ports
        assert!(engine.is_port_allowed("sbm-0xclient", 22, now));
        assert!(engine.is_port_allowed("sbm-0xclient", 8080, now));

        // Blocked database port
        assert!(!engine.is_port_allowed("sbm-0xclient", 5432, now));
        assert!(!engine.is_port_allowed("sbm-0xclient", 3306, now));

        // Unknown client blocked
        assert!(!engine.is_port_allowed("sbm-0xunknown", 22, now));
    }
}