agent_identity/

lib.rs

//! Agent identity, key management, and trust verification.

use std::collections::HashMap;

/// An agent's cryptographic identity.
#[derive(Debug, Clone)]
pub struct AgentIdentity {
    pub agent_id: String,
    pub public_key: Vec<u8>,
    pub created_at: u64,
    pub expires_at: Option<u64>,
    pub roles: Vec<String>,
    pub metadata: HashMap<String, String>,
}

impl AgentIdentity {
    pub fn new(agent_id: &str, public_key: Vec<u8>) -> Self {
        Self {
            agent_id: agent_id.to_string(),
            public_key,
            created_at: 0,
            expires_at: None,
            roles: vec![],
            metadata: HashMap::new(),
        }
    }

    pub fn with_role(mut self, role: &str) -> Self {
        self.roles.push(role.to_string());
        self
    }
    pub fn with_expiry(mut self, ts: u64) -> Self {
        self.expires_at = Some(ts);
        self
    }

    pub fn is_expired(&self, now: u64) -> bool {
        self.expires_at.is_some_and(|exp| now > exp)
    }

    pub fn has_role(&self, role: &str) -> bool {
        self.roles.iter().any(|r| r == role)
    }

    /// Simple key fingerprint (hex of first 8 bytes).
    pub fn fingerprint(&self) -> String {
        self.public_key
            .iter()
            .take(8)
            .map(|b| format!("{:02x}", b))
            .collect::<Vec<_>>()
            .join(":")
    }

    /// Verify this identity has required roles.
    pub fn verify_roles(&self, required: &[&str]) -> Result<(), Vec<String>> {
        let missing: Vec<String> = required
            .iter()
            .filter(|&&r| !self.has_role(r))
            .map(|&r| r.to_string())
            .collect();
        if missing.is_empty() {
            Ok(())
        } else {
            Err(missing)
        }
    }
}

/// A trust store for managing known agent identities.
#[derive(Debug, Clone)]
pub struct TrustStore {
    identities: HashMap<String, AgentIdentity>,
    revoked: Vec<String>,
}

impl Default for TrustStore {
    fn default() -> Self {
        Self::new()
    }
}

impl TrustStore {
    pub fn new() -> Self {
        Self {
            identities: HashMap::new(),
            revoked: vec![],
        }
    }

    pub fn register(&mut self, identity: AgentIdentity) {
        self.identities.insert(identity.agent_id.clone(), identity);
    }

    pub fn revoke(&mut self, agent_id: &str) {
        self.revoked.push(agent_id.to_string());
    }

    pub fn is_trusted(&self, agent_id: &str) -> bool {
        !self.revoked.contains(&agent_id.to_string()) && self.identities.contains_key(agent_id)
    }

    pub fn get(&self, agent_id: &str) -> Option<&AgentIdentity> {
        self.identities.get(agent_id)
    }

    pub fn verify(&self, agent_id: &str, now: u64) -> VerificationResult {
        if self.revoked.contains(&agent_id.to_string()) {
            return VerificationResult::Revoked;
        }
        match self.identities.get(agent_id) {
            None => VerificationResult::Unknown,
            Some(id) if id.is_expired(now) => VerificationResult::Expired,
            Some(_) => VerificationResult::Trusted,
        }
    }

    pub fn len(&self) -> usize {
        self.identities.len()
    }
    pub fn is_empty(&self) -> bool {
        self.identities.is_empty()
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum VerificationResult {
    Trusted,
    Unknown,
    Expired,
    Revoked,
}

/// A challenge-response token for authentication.
#[derive(Debug, Clone)]
pub struct AuthToken {
    pub agent_id: String,
    pub challenge: Vec<u8>,
    pub response: Vec<u8>,
    pub timestamp: u64,
}

impl AuthToken {
    pub fn new(agent_id: &str, challenge: Vec<u8>, response: Vec<u8>, timestamp: u64) -> Self {
        Self {
            agent_id: agent_id.to_string(),
            challenge,
            response,
            timestamp,
        }
    }

    /// Simple verification: check response matches expected.
    pub fn verify_response(&self, expected: &[u8]) -> bool {
        self.response == expected
    }
}

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

    #[test]
    fn test_create_identity() {
        let id = AgentIdentity::new("agent-1", vec![1, 2, 3, 4, 5, 6, 7, 8]).with_role("admin");
        assert_eq!(id.agent_id, "agent-1");
        assert!(id.has_role("admin"));
        assert!(!id.has_role("user"));
    }

    #[test]
    fn test_fingerprint() {
        let id = AgentIdentity::new(
            "agent-1",
            vec![0xab, 0xcd, 0xef, 0x01, 0x23, 0x45, 0x67, 0x89],
        );
        assert_eq!(id.fingerprint(), "ab:cd:ef:01:23:45:67:89");
    }

    #[test]
    fn test_expiry() {
        let id = AgentIdentity::new("agent-1", vec![1, 2, 3]).with_expiry(1000);
        assert!(id.is_expired(1001));
        assert!(!id.is_expired(999));
    }

    #[test]
    fn test_verify_roles() {
        let id = AgentIdentity::new("agent-1", vec![1, 2, 3])
            .with_role("read")
            .with_role("write");
        assert!(id.verify_roles(&["read"]).is_ok());
        assert!(id.verify_roles(&["read", "write"]).is_ok());
        assert!(id.verify_roles(&["admin"]).is_err());
    }

    #[test]
    fn test_trust_store() {
        let mut store = TrustStore::new();
        let id = AgentIdentity::new("agent-1", vec![1, 2, 3]);
        store.register(id);
        assert!(store.is_trusted("agent-1"));
        assert!(!store.is_trusted("unknown"));
    }

    #[test]
    fn test_revoke() {
        let mut store = TrustStore::new();
        store.register(AgentIdentity::new("agent-1", vec![1, 2, 3]));
        store.revoke("agent-1");
        assert!(!store.is_trusted("agent-1"));
    }

    #[test]
    fn test_verification() {
        let mut store = TrustStore::new();
        store.register(AgentIdentity::new("agent-1", vec![1, 2, 3]).with_expiry(1000));
        assert_eq!(store.verify("agent-1", 500), VerificationResult::Trusted);
        assert_eq!(store.verify("agent-1", 1001), VerificationResult::Expired);
        assert_eq!(store.verify("unknown", 500), VerificationResult::Unknown);
    }

    #[test]
    fn test_auth_token() {
        let token = AuthToken::new("agent-1", vec![1, 2, 3], vec![4, 5, 6], 0);
        assert!(token.verify_response(&[4, 5, 6]));
        assert!(!token.verify_response(&[7, 8, 9]));
    }
}