poimen-memory/crates/mem-ingest/src/authentik_jwt.rs

//! Authentik JWT Token Exchange
//! 
//! Uses OAuth2 client credentials flow to obtain JWT tokens from Authentik
//! These tokens are used to authenticate with LLM gateway and S3

use anyhow::{Result, anyhow};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::sync::Mutex;
use std::time::{SystemTime, Duration};

/// JWT token response from Authentik
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenResponse {
    pub access_token: String,
    pub token_type: String,
    pub expires_in: u64,
    #[serde(skip)]
    pub obtained_at: Option<SystemTime>,
}

impl TokenResponse {
    /// Check if token is still valid
    pub fn is_expired(&self) -> bool {
        match self.obtained_at {
            Some(time) => {
                let elapsed = time.elapsed().unwrap_or(Duration::from_secs(u64::MAX));
                elapsed.as_secs() >= self.expires_in - 60  // Refresh 60s before expiry
            }
            None => true,  // No timestamp = expired
        }
    }
}

/// Authentik JWT issuer client
pub struct AuthentikJwtIssuer {
    issuer_url: String,
    client_id: String,
    client_secret: String,
    cached_token: Arc<Mutex<Option<TokenResponse>>>,
}

impl AuthentikJwtIssuer {
    pub fn new(issuer_url: &str, client_id: &str, client_secret: &str) -> Self {
        Self {
            issuer_url: issuer_url.to_string(),
            client_id: client_id.to_string(),
            client_secret: client_secret.to_string(),
            cached_token: Arc::new(Mutex::new(None)),
        }
    }

    /// From environment: AUTHENTIK_ISSUER, AUTHENTIK_CLIENT_ID, AUTHENTIK_CLIENT_SECRET
    pub fn from_env() -> Result<Self> {
        let issuer = std::env::var("AUTHENTIK_ISSUER")
            .map_err(|_| anyhow!("AUTHENTIK_ISSUER not set"))?;
        let client_id = std::env::var("AUTHENTIK_CLIENT_ID")
            .map_err(|_| anyhow!("AUTHENTIK_CLIENT_ID not set"))?;
        let client_secret = std::env::var("AUTHENTIK_CLIENT_SECRET")
            .map_err(|_| anyhow!("AUTHENTIK_CLIENT_SECRET not set"))?;

        Ok(Self::new(&issuer, &client_id, &client_secret))
    }

    /// Get valid access token, using cache if available
    pub async fn get_access_token(&self) -> Result<String> {
        // Check cache
        if let Ok(lock) = self.cached_token.lock() {
            if let Some(token) = lock.as_ref() {
                if !token.is_expired() {
                    tracing::debug!("Using cached Authentik token");
                    return Ok(token.access_token.clone());
                }
            }
        }

        // Fetch new token
        let mut token = self.fetch_token().await?;
        token.obtained_at = Some(SystemTime::now());
        let access_token = token.access_token.clone();
        
        // Cache it
        if let Ok(mut lock) = self.cached_token.lock() {
            *lock = Some(token);
        }

        Ok(access_token)
    }

    /// Exchange client credentials for JWT token
    async fn fetch_token(&self) -> Result<TokenResponse> {
        let client = reqwest::Client::new();
        
        // Authentik OAuth2 token endpoint
        let token_url = format!("{}/token/", self.issuer_url.trim_end_matches('/'));
        
        let params = [
            ("grant_type", "client_credentials"),
            ("client_id", &self.client_id),
            ("client_secret", &self.client_secret),
        ];

        let response = client
            .post(&token_url)
            .form(&params)
            .timeout(Duration::from_secs(10))
            .send()
            .await?;

        if !response.status().is_success() {
            return Err(anyhow!(
                "Authentik token request failed: {} - {}",
                response.status(),
                response.text().await.unwrap_or_default()
            ));
        }

        let token_resp: TokenResponse = response.json().await?;

        tracing::info!(
            "Obtained Authentik JWT token (expires in {} seconds)",
            token_resp.expires_in
        );

        Ok(token_resp)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_token_expiry_check() {
        let mut token = TokenResponse {
            access_token: "test".to_string(),
            token_type: "Bearer".to_string(),
            expires_in: 3600,
            obtained_at: Some(SystemTime::now()),
        };

        assert!(!token.is_expired());

        // Simulate aged token
        token.obtained_at = Some(SystemTime::now() - Duration::from_secs(3600));
        assert!(token.is_expired());
    }

    #[test]
    fn test_issuer_creation() {
        let issuer = AuthentikJwtIssuer::new(
            "https://example.com",
            "client_id",
            "client_secret",
        );

        assert_eq!(issuer.issuer_url, "https://example.com");
        assert_eq!(issuer.client_id, "client_id");
    }
}