use super::types::{ApiPermission, ApiToken};
use crate::{ZoeyError, Result};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, warn};
pub struct ApiAuthManager {
tokens: Arc<RwLock<HashMap<String, ApiToken>>>,
}
impl ApiAuthManager {
pub fn new(tokens: Vec<ApiToken>) -> Self {
let token_map: HashMap<String, ApiToken> =
tokens.into_iter().map(|t| (t.token.clone(), t)).collect();
Self {
tokens: Arc::new(RwLock::new(token_map)),
}
}
pub fn disabled() -> Self {
Self {
tokens: Arc::new(RwLock::new(HashMap::new())),
}
}
pub fn hash_token(token: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(token.as_bytes());
format!("{:x}", hasher.finalize())
}
pub async fn validate_token(&self, token: &str) -> Result<Vec<ApiPermission>> {
let token_hash = Self::hash_token(token);
let tokens = self.tokens.read().await;
if tokens.is_empty() {
debug!("Authentication disabled, allowing access");
return Ok(vec![
ApiPermission::Read,
ApiPermission::Write,
ApiPermission::Execute,
]);
}
let api_token = tokens.get(&token_hash).ok_or_else(|| {
warn!("Authentication failed: invalid token");
ZoeyError::Config("Invalid authentication token".to_string())
})?;
if let Some(expires_at) = api_token.expires_at {
let now = chrono::Utc::now().timestamp();
if now > expires_at {
warn!("Authentication failed: token expired");
return Err(ZoeyError::Config(
"Authentication token has expired".to_string(),
));
}
}
debug!(
"Token validated successfully: {} (permissions: {:?})",
api_token.name, api_token.permissions
);
Ok(api_token.permissions.clone())
}
pub async fn has_permission(&self, token: &str, permission: ApiPermission) -> Result<bool> {
let permissions = self.validate_token(token).await?;
if permissions.contains(&ApiPermission::Admin) {
return Ok(true);
}
Ok(permissions.contains(&permission))
}
pub async fn add_token(&self, token: ApiToken) {
let mut tokens = self.tokens.write().await;
tokens.insert(token.token.clone(), token);
}
pub async fn remove_token(&self, token_hash: &str) -> bool {
let mut tokens = self.tokens.write().await;
tokens.remove(token_hash).is_some()
}
pub async fn list_tokens(&self) -> Vec<String> {
let tokens = self.tokens.read().await;
tokens.values().map(|t| t.name.clone()).collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use uuid::Uuid;
#[tokio::test]
async fn test_authentication() {
let token = ApiToken {
token: ApiAuthManager::hash_token("test-token"),
name: "Test Token".to_string(),
permissions: vec![ApiPermission::Read, ApiPermission::Write],
expires_at: None,
agent_id: Some(Uuid::new_v4()),
};
let auth = ApiAuthManager::new(vec![token]);
let perms = auth.validate_token("test-token").await.unwrap();
assert_eq!(perms.len(), 2);
assert!(perms.contains(&ApiPermission::Read));
assert!(auth.validate_token("invalid").await.is_err());
}
#[tokio::test]
async fn test_disabled_auth() {
let auth = ApiAuthManager::disabled();
let perms = auth.validate_token("any-token").await.unwrap();
assert!(perms.contains(&ApiPermission::Read));
}
#[tokio::test]
async fn test_has_permission() {
let admin_token = ApiToken {
token: ApiAuthManager::hash_token("admin-token"),
name: "Admin Token".to_string(),
permissions: vec![ApiPermission::Admin],
expires_at: None,
agent_id: None,
};
let auth = ApiAuthManager::new(vec![admin_token]);
assert!(auth
.has_permission("admin-token", ApiPermission::Execute)
.await
.unwrap());
}
#[test]
fn test_token_hashing() {
let hash1 = ApiAuthManager::hash_token("test");
let hash2 = ApiAuthManager::hash_token("test");
let hash3 = ApiAuthManager::hash_token("different");
assert_eq!(hash1, hash2);
assert_ne!(hash1, hash3);
}
}