//! Authentication Manager //! //! Provides unified authentication management for LLM providers with support for: //! - API key authentication //! - OAuth device code flow //! - Automatic token refresh //! - Credential persistence //! //! # Usage //! //! ```rust,ignore //! use auth_manager::AuthManager; //! use llm_core::ProviderType; //! //! let mut manager = AuthManager::new()?; //! //! // Login with OAuth device flow //! manager.login(ProviderType::Anthropic).await?; //! //! // Get auth for making API calls //! let auth = manager.get_auth(ProviderType::Anthropic)?; //! //! // Check status //! for status in manager.status() { //! println!("{}: authenticated={}", status.provider, status.authenticated); //! } //! ``` mod login; mod refresh; pub use login::LoginFlow; pub use refresh::{RefreshConfig, TokenRefresher}; use credentials::{CredentialError, CredentialManager, HelperManager}; use llm_core::{AuthMethod, ProviderStatus, ProviderType, StoredCredentials}; use std::collections::HashMap; use std::sync::{Arc, RwLock}; use std::time::{SystemTime, UNIX_EPOCH}; use thiserror::Error; // ============================================================================ // Error Types // ============================================================================ /// Errors that can occur during authentication operations #[derive(Error, Debug)] pub enum AuthError { #[error("Credential error: {0}")] Credential(#[from] CredentialError), #[error("OAuth error: {0}")] OAuth(String), #[error("Provider error: {0}")] Provider(#[from] llm_core::LlmError), #[error("Provider not supported for OAuth: {0}")] NotSupported(String), #[error("Not authenticated for provider: {0}")] NotAuthenticated(String), #[error("Token expired and no refresh token available")] TokenExpired, #[error("Login cancelled by user")] Cancelled, } pub type Result = std::result::Result; // ============================================================================ // Authentication Manager // ============================================================================ /// Manages authentication for all LLM providers pub struct AuthManager { /// Credential storage credentials: CredentialManager, /// Cached auth methods per provider (for performance) cache: RwLock>, /// Environment variable overrides (checked first) env_overrides: HashMap, /// API key helpers (1Password, pass, etc.) helpers: RwLock, } impl AuthManager { /// Create a new authentication manager pub fn new() -> Result { let credentials = CredentialManager::new()?; let env_overrides = Self::load_env_overrides(); Ok(Self { credentials, cache: RwLock::new(HashMap::new()), env_overrides, helpers: RwLock::new(HelperManager::new()), }) } /// Register an API key helper for a provider /// /// Helpers are checked after environment variables but before stored credentials. /// This allows using password managers like 1Password, Bitwarden, or pass. /// /// # Example /// /// ```rust,ignore /// use credentials::one_password_helper; /// /// let manager = AuthManager::new()?; /// manager.register_helper( /// ProviderType::Anthropic, /// Box::new(one_password_helper("Private", "Anthropic", "api_key")) /// ); /// ``` pub fn register_helper( &self, provider: ProviderType, helper: Box, ) -> Result<()> { let mut helpers = self.helpers.write().map_err(|_| { AuthError::OAuth("Failed to acquire helpers write lock".to_string()) })?; helpers.register(provider.as_str(), helper); Ok(()) } /// Register a command-based API key helper /// /// Convenience method for registering a shell command that outputs the API key. /// /// # Example /// /// ```rust,ignore /// manager.register_command_helper( /// ProviderType::Anthropic, /// "op read 'op://Private/Anthropic/api_key'" /// )?; /// ``` pub fn register_command_helper( &self, provider: ProviderType, command: impl Into, ) -> Result<()> { let mut helpers = self.helpers.write().map_err(|_| { AuthError::OAuth("Failed to acquire helpers write lock".to_string()) })?; helpers.register_command(provider.as_str(), command); Ok(()) } /// Load API key overrides from environment variables fn load_env_overrides() -> HashMap { let mut overrides = HashMap::new(); // Check standard environment variables if let Ok(key) = std::env::var("ANTHROPIC_API_KEY") { overrides.insert(ProviderType::Anthropic, key); } if let Ok(key) = std::env::var("OPENAI_API_KEY") { overrides.insert(ProviderType::OpenAI, key); } if let Ok(key) = std::env::var("OLLAMA_API_KEY") { overrides.insert(ProviderType::Ollama, key); } // Also check OWLEN_ prefixed variants if let Ok(key) = std::env::var("OWLEN_ANTHROPIC_API_KEY") { overrides.insert(ProviderType::Anthropic, key); } if let Ok(key) = std::env::var("OWLEN_OPENAI_API_KEY") { overrides.insert(ProviderType::OpenAI, key); } if let Ok(key) = std::env::var("OWLEN_API_KEY") { // Generic API key - use for Ollama if set if !overrides.contains_key(&ProviderType::Ollama) { overrides.insert(ProviderType::Ollama, key); } } overrides } /// Get authentication for a provider /// /// Priority order: /// 1. Environment variables (ANTHROPIC_API_KEY, etc.) /// 2. API key helpers (1Password, pass, Bitwarden, etc.) /// 3. In-memory cache /// 4. Stored credentials /// 5. None (for providers that don't require auth) pub fn get_auth(&self, provider: ProviderType) -> Result { // 1. Check environment variable override first if let Some(key) = self.env_overrides.get(&provider) { return Ok(AuthMethod::ApiKey(key.clone())); } // 2. Check API key helpers (1Password, pass, etc.) { let helpers = self.helpers.read().map_err(|_| { AuthError::OAuth("Failed to acquire helpers read lock".to_string()) })?; if let Some(result) = helpers.get_key(provider.as_str()) { match result { Ok(key) => return Ok(AuthMethod::ApiKey(key)), Err(e) => { // Log helper error but continue to other sources tracing::warn!( provider = provider.as_str(), error = %e, "API key helper failed, trying other sources" ); } } } } // 3. Check cache { let cache = self.cache.read().map_err(|_| { AuthError::OAuth("Failed to acquire cache read lock".to_string()) })?; if let Some(auth) = cache.get(&provider) { // Check if OAuth token needs refresh if !auth.needs_refresh() { return Ok(auth.clone()); } } } // 4. Load from stored credentials let provider_name = provider.as_str(); if let Some(creds) = self.credentials.retrieve(provider_name)? { let auth = Self::credentials_to_auth(&creds); // Cache it let mut cache = self.cache.write().map_err(|_| { AuthError::OAuth("Failed to acquire cache write lock".to_string()) })?; cache.insert(provider, auth.clone()); return Ok(auth); } // 5. No credentials found // For Ollama, return None (local doesn't need auth) if provider == ProviderType::Ollama { return Ok(AuthMethod::None); } Err(AuthError::NotAuthenticated(provider_name.to_string())) } /// Convert stored credentials to AuthMethod fn credentials_to_auth(creds: &StoredCredentials) -> AuthMethod { if creds.refresh_token.is_some() || creds.expires_at.is_some() { // This is OAuth AuthMethod::OAuth { access_token: creds.access_token.clone(), refresh_token: creds.refresh_token.clone(), expires_at: creds.expires_at, } } else { // This is an API key AuthMethod::ApiKey(creds.access_token.clone()) } } /// Store an API key for a provider pub fn store_api_key(&self, provider: ProviderType, api_key: &str) -> Result<()> { let creds = StoredCredentials { provider: provider.as_str().to_string(), access_token: api_key.to_string(), refresh_token: None, expires_at: None, }; self.credentials.store(provider.as_str(), creds)?; // Update cache let mut cache = self.cache.write().map_err(|_| { AuthError::OAuth("Failed to acquire cache write lock".to_string()) })?; cache.insert(provider, AuthMethod::ApiKey(api_key.to_string())); Ok(()) } /// Store OAuth credentials for a provider pub fn store_oauth( &self, provider: ProviderType, access_token: &str, refresh_token: Option<&str>, expires_in: Option, ) -> Result<()> { let expires_at = expires_in.map(|secs| { SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_secs() + secs) .unwrap_or(0) }); let creds = StoredCredentials { provider: provider.as_str().to_string(), access_token: access_token.to_string(), refresh_token: refresh_token.map(|s| s.to_string()), expires_at, }; self.credentials.store(provider.as_str(), creds)?; // Update cache let auth = AuthMethod::OAuth { access_token: access_token.to_string(), refresh_token: refresh_token.map(|s| s.to_string()), expires_at, }; let mut cache = self.cache.write().map_err(|_| { AuthError::OAuth("Failed to acquire cache write lock".to_string()) })?; cache.insert(provider, auth); Ok(()) } /// Perform OAuth login for a provider pub async fn login(&self, provider: ProviderType) -> Result<()> { let flow = LoginFlow::new(provider)?; let auth = flow.execute().await?; // Store the credentials match &auth { AuthMethod::OAuth { access_token, refresh_token, expires_at, } => { let expires_in = expires_at.map(|exp| { let now = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_secs()) .unwrap_or(0); exp.saturating_sub(now) }); self.store_oauth( provider, access_token, refresh_token.as_deref(), expires_in, )?; } AuthMethod::ApiKey(key) => { self.store_api_key(provider, key)?; } AuthMethod::None => { // Nothing to store } } Ok(()) } /// Logout from a provider (delete stored credentials) pub fn logout(&self, provider: ProviderType) -> Result<()> { // Remove from credential store self.credentials.delete(provider.as_str())?; // Remove from cache let mut cache = self.cache.write().map_err(|_| { AuthError::OAuth("Failed to acquire cache write lock".to_string()) })?; cache.remove(&provider); Ok(()) } /// Get authentication status for all providers pub fn status(&self) -> Vec { let providers = [ ProviderType::Ollama, ProviderType::Anthropic, ProviderType::OpenAI, ]; // Check which providers have helpers configured let helper_providers: Vec = self.helpers.read() .map(|h| h.providers().into_iter().map(|s| s.to_string()).collect()) .unwrap_or_default(); providers .iter() .map(|provider| { let auth_result = self.get_auth(*provider); let authenticated = matches!( &auth_result, Ok(AuthMethod::ApiKey(_)) | Ok(AuthMethod::OAuth { .. }) ); let has_helper = helper_providers.contains(&provider.as_str().to_string()); let (source, message) = if self.env_overrides.contains_key(provider) { ("env", Some("API key from environment variable".to_string())) } else if has_helper && authenticated { ("helper", Some("API key from helper command".to_string())) } else if authenticated { ("stored", None) } else { ("none", auth_result.err().map(|e| e.to_string())) }; ProviderStatus { provider: provider.as_str().to_string(), authenticated, account: None, // Would need to call provider API to get this model: provider.default_model().to_string(), endpoint: match provider { ProviderType::Ollama => "http://localhost:11434".to_string(), ProviderType::Anthropic => "https://api.anthropic.com".to_string(), ProviderType::OpenAI => "https://api.openai.com".to_string(), }, reachable: true, // Would need to check connectivity message: message.or_else(|| Some(format!("Source: {}", source))), } }) .collect() } /// Check if a helper is registered for a provider pub fn has_helper(&self, provider: ProviderType) -> bool { self.helpers.read() .map(|h| h.has_helper(provider.as_str())) .unwrap_or(false) } /// Check if a provider has stored credentials pub fn is_authenticated(&self, provider: ProviderType) -> bool { matches!( self.get_auth(provider), Ok(AuthMethod::ApiKey(_)) | Ok(AuthMethod::OAuth { .. }) ) } /// Get the credential storage backend name pub fn storage_name(&self) -> &'static str { self.credentials.storage_name() } /// Clear all cached credentials (does not delete stored credentials) pub fn clear_cache(&self) { if let Ok(mut cache) = self.cache.write() { cache.clear(); } } /// Start background token refresh with default configuration /// /// Spawns a background task that periodically checks for expiring tokens /// and refreshes them before they expire. /// /// Returns a `TokenRefresher` handle that can be used to stop the refresh task. pub fn start_background_refresh(self: Arc) -> TokenRefresher { self.start_background_refresh_with_config(RefreshConfig::default()) } /// Start background token refresh with custom configuration pub fn start_background_refresh_with_config( self: Arc, config: RefreshConfig, ) -> TokenRefresher { TokenRefresher::start(self, config) } // ========================================================================= // MCP Server OAuth Methods // ========================================================================= /// Get the credential key for an MCP server fn mcp_credential_key(server_name: &str) -> String { format!("mcp:{}", server_name) } /// Store OAuth token for an MCP server /// /// MCP tokens are stored with the key pattern `mcp:{server_name}`. pub fn store_mcp_token( &self, server_name: &str, access_token: &str, refresh_token: Option<&str>, expires_in: Option, ) -> Result<()> { let key = Self::mcp_credential_key(server_name); let expires_at = expires_in.map(|secs| { SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_secs() + secs) .unwrap_or(0) }); let creds = StoredCredentials { provider: key.clone(), access_token: access_token.to_string(), refresh_token: refresh_token.map(|s| s.to_string()), expires_at, }; self.credentials.store(&key, creds)?; Ok(()) } /// Get OAuth token for an MCP server /// /// Returns the access token if available, or None if not authenticated. pub fn get_mcp_token(&self, server_name: &str) -> Result> { let key = Self::mcp_credential_key(server_name); if let Some(creds) = self.credentials.retrieve(&key)? { // Check if token is expired if let Some(expires_at) = creds.expires_at { let now = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_secs()) .unwrap_or(0); if expires_at <= now { // Token expired // TODO: Auto-refresh if refresh token is available return Ok(None); } } Ok(Some(creds.access_token)) } else { Ok(None) } } /// Check if an MCP server is authenticated pub fn is_mcp_authenticated(&self, server_name: &str) -> bool { self.get_mcp_token(server_name) .ok() .flatten() .is_some() } /// Remove stored token for an MCP server pub fn logout_mcp(&self, server_name: &str) -> Result<()> { let key = Self::mcp_credential_key(server_name); self.credentials.delete(&key)?; Ok(()) } /// List all authenticated MCP servers pub fn list_mcp_servers(&self) -> Result> { let providers = self.credentials.list_providers()?; Ok(providers .into_iter() .filter_map(|p| p.strip_prefix("mcp:").map(|s| s.to_string())) .collect()) } } impl Default for AuthManager { fn default() -> Self { Self::new().expect("Failed to create auth manager") } } // ============================================================================ // Tests // ============================================================================ #[cfg(test)] mod tests { use super::*; use credentials::HelperManager; #[test] fn test_env_override_loading() { // Create a temp dir for config to avoid loading system credentials let temp_dir = tempfile::tempdir().unwrap(); let temp_path = temp_dir.path().to_str().unwrap(); // Clear any env vars that might interfere and set XDG_CONFIG_HOME // Also disable keyring to ensure we only use the temp config dir // SAFETY: Single-threaded test context unsafe { std::env::set_var("XDG_CONFIG_HOME", temp_path); std::env::set_var("OWLEN_KEYRING_DISABLE", "1"); std::env::remove_var("ANTHROPIC_API_KEY"); std::env::remove_var("OWLEN_ANTHROPIC_API_KEY"); } // Set env var (unsafe in Rust 2024 due to potential thread safety issues) // SAFETY: This is a single-threaded test, no concurrent access unsafe { std::env::set_var("ANTHROPIC_API_KEY", "test-key-123"); } let manager = AuthManager::new().unwrap(); let auth = manager.get_auth(ProviderType::Anthropic).unwrap(); match auth { AuthMethod::ApiKey(key) => assert_eq!(key, "test-key-123"), _ => panic!("Expected API key auth"), } // Clean up // SAFETY: Single-threaded test unsafe { std::env::remove_var("ANTHROPIC_API_KEY"); std::env::remove_var("XDG_CONFIG_HOME"); std::env::remove_var("OWLEN_KEYRING_DISABLE"); } } #[test] fn test_ollama_no_auth_required() { // Create a temp dir for config let temp_dir = tempfile::tempdir().unwrap(); let temp_path = temp_dir.path().to_str().unwrap(); // Clear any env vars that might interfere and set XDG_CONFIG_HOME // Also disable keyring // SAFETY: Single-threaded test context unsafe { std::env::set_var("XDG_CONFIG_HOME", temp_path); std::env::set_var("OWLEN_KEYRING_DISABLE", "1"); std::env::remove_var("OLLAMA_API_KEY"); std::env::remove_var("OWLEN_API_KEY"); } let manager = AuthManager::new().unwrap(); let auth = manager.get_auth(ProviderType::Ollama).unwrap(); assert!(matches!(auth, AuthMethod::None)); // Clean up unsafe { std::env::remove_var("XDG_CONFIG_HOME"); std::env::remove_var("OWLEN_KEYRING_DISABLE"); } } #[test] fn test_status() { let manager = AuthManager::new().unwrap(); let status = manager.status(); assert_eq!(status.len(), 3); let provider_names: Vec<&str> = status.iter().map(|s| s.provider.as_str()).collect(); assert!(provider_names.contains(&"ollama")); assert!(provider_names.contains(&"anthropic")); assert!(provider_names.contains(&"openai")); } #[test] fn test_helper_registration() { // Clear any env vars that might interfere (from other tests) // SAFETY: Single-threaded test context unsafe { std::env::remove_var("ANTHROPIC_API_KEY"); std::env::remove_var("OWLEN_ANTHROPIC_API_KEY"); } // Create a fresh manager after clearing env vars let manager = AuthManager::new().unwrap(); // No helper initially assert!(!manager.has_helper(ProviderType::Anthropic)); // Register a command helper manager .register_command_helper(ProviderType::Anthropic, "echo 'helper-test-key'") .unwrap(); // Now has helper assert!(manager.has_helper(ProviderType::Anthropic)); // Should get auth from helper let auth = manager.get_auth(ProviderType::Anthropic).unwrap(); match auth { AuthMethod::ApiKey(key) => assert_eq!(key, "helper-test-key"), _ => panic!("Expected API key from helper"), } } #[test] fn test_helper_in_status() { let manager = AuthManager::new().unwrap(); // Register a helper manager .register_command_helper(ProviderType::OpenAI, "echo 'status-test-key'") .unwrap(); let status = manager.status(); let openai_status = status.iter().find(|s| s.provider == "openai").unwrap(); assert!(openai_status.authenticated); assert!(openai_status .message .as_ref() .map(|m| m.contains("helper")) .unwrap_or(false)); } #[test] fn test_mcp_token_storage() { // Use file-only credential manager to avoid keyring issues in test environment let credentials = CredentialManager::file_only().unwrap(); let manager = AuthManager { credentials, cache: RwLock::new(HashMap::new()), env_overrides: HashMap::new(), helpers: RwLock::new(HelperManager::new()), }; // Use a unique server name to avoid conflicts with other tests let server_name = format!("test-mcp-{}", std::process::id()); // Initially not authenticated assert!(!manager.is_mcp_authenticated(&server_name)); // Store a token manager .store_mcp_token(&server_name, "test-token-123", Some("refresh-456"), Some(3600)) .unwrap(); // Now authenticated assert!( manager.is_mcp_authenticated(&server_name), "Expected to be authenticated after storing token" ); // Get token let token = manager.get_mcp_token(&server_name).unwrap(); assert_eq!(token, Some("test-token-123".to_string())); // Cleanup manager.logout_mcp(&server_name).unwrap(); assert!(!manager.is_mcp_authenticated(&server_name)); } }