poimen-memory/crates/mem-cli/src/auth/authentik_provider.rs

/// Authentik OIDC provider implementation.
///
/// Validates JWT tokens issued by Authentik and extracts claims.

use async_trait::async_trait;
use jsonwebtoken::{decode, decode_header, DecodingKey, Validation, Algorithm};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::Arc;
use tokio::sync::RwLock;

use super::provider::{AuthProvider, Claims, AuthError};

/// JWT token claims from Authentik.
#[derive(Debug, Deserialize, Serialize)]
pub struct TokenClaims {
    pub sub: String,
    pub iss: String,
    pub aud: String,
    pub exp: i64,
    pub iat: i64,
    pub groups: Option<Vec<String>>,
    pub attributes: Option<serde_json::Map<String, Value>>,
}

/// JWKS entry (public key).
#[derive(Debug, Deserialize)]
pub struct JwksKey {
    pub kid: String,
    pub kty: String,
    pub use_: Option<String>,
    pub n: String,
    pub e: String,
}

/// JWKS response from Authentik.
#[derive(Debug, Deserialize)]
pub struct JwkSet {
    pub keys: Vec<JwksKey>,
}

/// Authentik provider configuration.
#[derive(Clone, Debug)]
pub struct AuthentikConfig {
    pub issuer: String,           // https://authentik.riotpiao.com/application/o/memory/
    pub audience: String,         // poimen-memory
    pub jwks_uri: String,         // https://authentik.riotpiao.com/.well-known/openid-configuration
    pub cache_ttl_secs: u64,     // Default 3600
}

/// Authentik OIDC provider.
pub struct AuthentikProvider {
    config: AuthentikConfig,
    http_client: reqwest::Client,
    // TODO: Add JWKS cache
    // jwks_cache: Arc<RwLock<Option<(JwkSet, Instant)>>>,
}

impl AuthentikProvider {
    /// Create new Authentik provider.
    pub fn new(config: AuthentikConfig) -> Self {
        Self {
            config,
            http_client: reqwest::Client::new(),
        }
    }
    
    /// Fetch JWKS from Authentik (should be cached in real implementation).
    async fn fetch_jwks(&self) -> Result<JwkSet, AuthError> {
        // TODO: Implement JWKS caching (1 hour TTL)
        // For now, always fetch
        
        // First get OIDC config to find jwks_uri
        let config_url = format!("{}/.well-known/openid-configuration", self.config.issuer);
        
        let config_response = self.http_client
            .get(&config_url)
            .send()
            .await
            .map_err(|e| AuthError::ProviderUnavailable(e.to_string()))?;
        
        let config: serde_json::Value = config_response
            .json()
            .await
            .map_err(|e| AuthError::ProviderUnavailable(e.to_string()))?;
        
        let jwks_uri = config["jwks_uri"]
            .as_str()
            .ok_or(AuthError::ProviderUnavailable("No jwks_uri in config".to_string()))?;
        
        // Fetch JWKS
        let jwks_response = self.http_client
            .get(jwks_uri)
            .send()
            .await
            .map_err(|e| AuthError::ProviderUnavailable(e.to_string()))?;
        
        jwks_response
            .json::<JwkSet>()
            .await
            .map_err(|e| AuthError::ProviderUnavailable(e.to_string()))
    }
}

#[async_trait]
impl AuthProvider for AuthentikProvider {
    async fn validate_token(&self, token: &str) -> Result<Claims, AuthError> {
        // 1. Decode header to find kid
        let header = decode_header(token)
            .map_err(|_| AuthError::InvalidSignature)?;
        
        let kid = header.kid
            .ok_or(AuthError::InvalidSignature)?;
        
        // 2. Fetch JWKS to find public key
        let jwks = self.fetch_jwks().await?;
        
        let jwks_key = jwks.keys.iter()
            .find(|k| k.kid == kid)
            .ok_or(AuthError::InvalidSignature)?;
        
        // 3. Decode and verify JWT
        // TODO: Implement RSA key construction from JWKS
        // For now, this is a placeholder
        
        let claims: TokenClaims = decode::<TokenClaims>(
            token,
            &DecodingKey::from_secret(b"TODO"), // Placeholder
            &Validation::new(Algorithm::RS256),
        )
            .map_err(|_| AuthError::InvalidSignature)?
            .claims;
        
        // 4. Validate issuer and audience
        if claims.iss != self.config.issuer {
            return Err(AuthError::InvalidIssuer);
        }
        
        if claims.aud != self.config.audience {
            return Err(AuthError::InvalidAudience);
        }
        
        // 5. Check expiration
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;
        
        if claims.exp < now {
            return Err(AuthError::TokenExpired);
        }
        
        // 6. Convert to standard Claims format
        Ok(Claims {
            sub: claims.sub,
            groups: claims.groups.unwrap_or_default(),
            attributes: claims.attributes.unwrap_or_default(),
            exp: claims.exp,
            iat: claims.iat,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_authentik_config() {
        let config = AuthentikConfig {
            issuer: "https://authentik.riotpiao.com/application/o/memory/".to_string(),
            audience: "poimen-memory".to_string(),
            jwks_uri: "https://authentik.riotpiao.com/.well-known/openid-configuration".to_string(),
            cache_ttl_secs: 3600,
        };
        
        assert_eq!(config.audience, "poimen-memory");
    }
    
    #[test]
    fn test_token_claims() {
        let claims = TokenClaims {
            sub: "rock".to_string(),
            iss: "https://authentik.riotpiao.com/application/o/memory/".to_string(),
            aud: "poimen-memory".to_string(),
            exp: 1735689600,
            iat: 1735689300,
            groups: Some(vec!["memory-users".to_string()]),
            attributes: None,
        };
        
        assert_eq!(claims.sub, "rock");
    }
}