use chrono::{DateTime, Utc};
use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{debug, info};
use uuid::Uuid;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum Permission {
CreateSession,
ExecuteCommand,
ViewSessions,
DeleteSession,
ManageAgents,
SystemAdmin,
}
impl Permission {
pub fn as_str(&self) -> &'static str {
match self {
Permission::CreateSession => "create_session",
Permission::ExecuteCommand => "execute_command",
Permission::ViewSessions => "view_sessions",
Permission::DeleteSession => "delete_session",
Permission::ManageAgents => "manage_agents",
Permission::SystemAdmin => "system_admin",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentCredentials {
pub id: String,
pub name: String,
pub api_key: String,
pub permissions: Vec<Permission>,
pub created_at: DateTime<Utc>,
pub last_used: Option<DateTime<Utc>>,
pub active: bool,
pub metadata: HashMap<String, String>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct JwtClaims {
pub sub: String, pub name: String,
pub permissions: Vec<String>,
pub exp: usize, pub iat: usize, pub jti: String, }
pub struct AuthManager {
agents: Arc<RwLock<HashMap<String, AgentCredentials>>>,
jwt_secret: String,
token_expiry: Duration,
}
impl AuthManager {
pub fn new(jwt_secret: String) -> Self {
Self {
agents: Arc::new(RwLock::new(HashMap::new())),
jwt_secret,
token_expiry: Duration::from_secs(24 * 60 * 60), }
}
pub async fn add_agent(&self, agent: AgentCredentials) -> Result<(), anyhow::Error> {
let mut agents = self.agents.write().await;
if agents.contains_key(&agent.id) {
return Err(anyhow::anyhow!("Agent with ID {} already exists", agent.id));
}
agents.insert(agent.id.clone(), agent.clone());
info!("Added new agent: {}", agent.name);
Ok(())
}
pub async fn remove_agent(&self, agent_id: &str) -> Result<(), anyhow::Error> {
let mut agents = self.agents.write().await;
if agents.remove(agent_id).is_none() {
return Err(anyhow::anyhow!("Agent with ID {} not found", agent_id));
}
info!("Removed agent: {}", agent_id);
Ok(())
}
pub async fn get_agent(&self, agent_id: &str) -> Option<AgentCredentials> {
let agents = self.agents.read().await;
agents.get(agent_id).cloned()
}
pub async fn list_agents(&self) -> Vec<AgentCredentials> {
let agents = self.agents.read().await;
agents.values().cloned().collect()
}
pub async fn validate_token(&self, token: &str) -> Result<String, anyhow::Error> {
let token_data = decode::<JwtClaims>(
token,
&DecodingKey::from_secret(self.jwt_secret.as_ref()),
&Validation::default(),
)?;
let claims = token_data.claims;
let agent_id = claims.sub.clone();
let agent = self
.get_agent(&agent_id)
.await
.ok_or_else(|| anyhow::anyhow!("Agent not found: {}", agent_id))?;
if !agent.active {
return Err(anyhow::anyhow!("Agent is inactive: {}", agent_id));
}
debug!("Successfully validated token for agent: {}", agent_id);
Ok(agent_id)
}
pub async fn authenticate_api_key(&self, api_key: &str) -> Result<String, anyhow::Error> {
let agents = self.agents.read().await;
let agent = agents
.values()
.find(|a| a.api_key == api_key && a.active)
.ok_or_else(|| anyhow::anyhow!("Invalid API key"))?;
info!("Agent authenticated via API key: {}", agent.name);
Ok(agent.id.clone())
}
pub async fn generate_token(&self, agent_id: &str) -> Result<String, anyhow::Error> {
let agent = self
.get_agent(agent_id)
.await
.ok_or_else(|| anyhow::anyhow!("Agent not found: {}", agent_id))?;
let now = Utc::now();
let exp = now
+ chrono::Duration::from_std(self.token_expiry)
.map_err(|e| anyhow::anyhow!("Invalid expiration duration: {}", e))?;
let claims = JwtClaims {
sub: agent.id.clone(),
name: agent.name.clone(),
permissions: agent
.permissions
.iter()
.map(|p| p.as_str().to_string())
.collect(),
exp: exp.timestamp() as usize,
iat: now.timestamp() as usize,
jti: Uuid::new_v4().to_string(),
};
let token = encode(
&Header::default(),
&claims,
&EncodingKey::from_secret(self.jwt_secret.as_ref()),
)?;
info!("Generated JWT token for agent: {}", agent.name);
Ok(token)
}
pub async fn check_permission(&self, agent_id: &str, permission: &Permission) -> bool {
let agents = self.agents.read().await;
if let Some(agent) = agents.get(agent_id) {
agent.permissions.contains(permission)
|| agent.permissions.contains(&Permission::SystemAdmin)
} else {
false
}
}
pub async fn update_last_used(&self, agent_id: &str) -> Result<(), anyhow::Error> {
let mut agents = self.agents.write().await;
if let Some(agent) = agents.get_mut(agent_id) {
agent.last_used = Some(Utc::now());
}
Ok(())
}
pub async fn get_agent_stats(&self) -> AgentStats {
let agents = self.agents.read().await;
let total_agents = agents.len();
let active_agents = agents.values().filter(|a| a.active).count();
let agents_with_recent_activity = agents
.values()
.filter(|a| {
if let Some(last_used) = a.last_used {
Utc::now() - last_used < chrono::Duration::hours(24)
} else {
false
}
})
.count();
AgentStats {
total_agents,
active_agents,
agents_with_recent_activity,
}
}
pub async fn init_default_agent(&self) -> Result<String, anyhow::Error> {
let admin_agent = AgentCredentials {
id: "admin".to_string(),
name: "Default Admin Agent".to_string(),
api_key: "sk-ditto-admin-2024".to_string(),
permissions: vec![
Permission::CreateSession,
Permission::ExecuteCommand,
Permission::ViewSessions,
Permission::DeleteSession,
Permission::ManageAgents,
Permission::SystemAdmin,
],
created_at: Utc::now(),
last_used: None,
active: true,
metadata: HashMap::new(),
};
let agent_id = admin_agent.id.clone();
self.add_agent(admin_agent).await?;
info!("Initialized default admin agent");
self.generate_token(&agent_id).await
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentStats {
pub total_agents: usize,
pub active_agents: usize,
pub agents_with_recent_activity: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_agent_authentication() {
let auth_manager = AuthManager::new("test_secret".to_string());
let agent = AgentCredentials {
id: "test_agent".to_string(),
name: "Test Agent".to_string(),
api_key: "test_key".to_string(),
permissions: vec![Permission::CreateSession],
created_at: Utc::now(),
last_used: None,
active: true,
metadata: HashMap::new(),
};
auth_manager.add_agent(agent).await.unwrap();
let result = auth_manager.authenticate_api_key("test_key").await;
assert!(result.is_ok());
assert_eq!(result.unwrap(), "test_agent");
let result = auth_manager.authenticate_api_key("invalid_key").await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_jwt_token_generation() {
let auth_manager = AuthManager::new("test_secret".to_string());
let agent = AgentCredentials {
id: "test_agent".to_string(),
name: "Test Agent".to_string(),
api_key: "test_key".to_string(),
permissions: vec![Permission::CreateSession],
created_at: Utc::now(),
last_used: None,
active: true,
metadata: HashMap::new(),
};
auth_manager.add_agent(agent).await.unwrap();
let token = auth_manager.generate_token("test_agent").await.unwrap();
assert!(!token.is_empty());
let agent_id = auth_manager.validate_token(&token).await.unwrap();
assert_eq!(agent_id, "test_agent");
}
#[tokio::test]
async fn test_permission_checking() {
let auth_manager = AuthManager::new("test_secret".to_string());
let agent = AgentCredentials {
id: "test_agent".to_string(),
name: "Test Agent".to_string(),
api_key: "test_key".to_string(),
permissions: vec![Permission::CreateSession],
created_at: Utc::now(),
last_used: None,
active: true,
metadata: HashMap::new(),
};
auth_manager.add_agent(agent).await.unwrap();
assert!(
auth_manager
.check_permission("test_agent", &Permission::CreateSession)
.await
);
assert!(
!auth_manager
.check_permission("test_agent", &Permission::DeleteSession)
.await
);
assert!(
!auth_manager
.check_permission("non_existent", &Permission::CreateSession)
.await
);
}
}