use std::collections::HashMap;
#[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)
}
pub fn fingerprint(&self) -> String {
self.public_key
.iter()
.take(8)
.map(|b| format!("{:02x}", b))
.collect::<Vec<_>>()
.join(":")
}
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)
}
}
}
#[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,
}
#[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,
}
}
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]));
}
}