bascule-oss/crates/bascule-auth-agent-id/src/lib.rs

//! Microsoft Entra Agent ID authentication provider for Bascule.
//!
//! Validates AI agent identity tokens issued by Microsoft Entra.
//! Agents authenticate via OAuth 2.0 client_credentials and present
//! their token as the SSH password (token-as-password pattern).
//!
//! Agent metadata extracted from token claims:
//! - Agent application/client ID
//! - Display name and agent type
//! - Blueprint ID (agent type template)
//! - Sponsor (human or org that registered the agent)
//! - Delegation chain (on-behalf-of claims)
//! - Scopes and roles

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use async_trait::async_trait;
use tokio::sync::RwLock;

use bascule_core::auth::AuthProvider;

/// Metadata extracted from an Entra Agent ID token.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct AgentIdentity {
    pub agent_id: String,
    pub display_name: String,
    pub agent_type: Option<String>,
    pub blueprint_id: Option<String>,
    pub sponsor: Option<String>,
    pub on_behalf_of: Option<String>,
    pub tenant_id: String,
    pub scopes: Vec<String>,
    pub expires_at: Option<chrono::DateTime<chrono::Utc>>,
    pub tags: HashMap<String, String>,
}

/// Microsoft Entra Agent ID authentication provider.
pub struct EntraAgentIdProvider {
    tenant_id: String,
    expected_audiences: Vec<String>,
    jwks_cache: Arc<RwLock<Option<JwksCache>>>,
    allow_multi_tenant: bool,
}

struct JwksCache {
    keys: jsonwebtoken::jwk::JwkSet,
    fetched_at: Instant,
}

impl EntraAgentIdProvider {
    /// Create a provider for a specific Entra tenant.
    pub fn new(tenant_id: &str, expected_audiences: Vec<String>) -> Self {
        Self {
            tenant_id: tenant_id.to_string(),
            expected_audiences,
            jwks_cache: Arc::new(RwLock::new(None)),
            allow_multi_tenant: false,
        }
    }

    /// Create a multi-tenant provider (accepts agents from any tenant).
    pub fn multi_tenant(expected_audiences: Vec<String>) -> Self {
        Self {
            tenant_id: "common".to_string(),
            expected_audiences,
            jwks_cache: Arc::new(RwLock::new(None)),
            allow_multi_tenant: true,
        }
    }

    async fn get_jwks(&self) -> anyhow::Result<jsonwebtoken::jwk::JwkSet> {
        {
            let cache = self.jwks_cache.read().await;
            if let Some(ref c) = *cache {
                if c.fetched_at.elapsed() < Duration::from_secs(3600) {
                    return Ok(c.keys.clone());
                }
            }
        }

        let url = format!(
            "https://login.microsoftonline.com/{}/discovery/v2.0/keys",
            self.tenant_id
        );
        let client = reqwest::Client::new();
        let jwks: jsonwebtoken::jwk::JwkSet = client.get(&url).send().await?.json().await?;

        {
            let mut cache = self.jwks_cache.write().await;
            *cache = Some(JwksCache {
                keys: jwks.clone(),
                fetched_at: Instant::now(),
            });
        }

        Ok(jwks)
    }

    async fn validate_token(&self, token: &str) -> anyhow::Result<AgentIdentity> {
        let jwks = self.get_jwks().await?;

        let header = jsonwebtoken::decode_header(token)?;
        let kid = header
            .kid
            .as_deref()
            .ok_or_else(|| anyhow::anyhow!("Token missing kid header"))?;

        let jwk = jwks
            .keys
            .iter()
            .find(|k| k.common.key_id.as_deref() == Some(kid))
            .ok_or_else(|| anyhow::anyhow!("Key ID not found in JWKS"))?;

        let decoding_key = jsonwebtoken::DecodingKey::from_jwk(jwk)?;

        let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::RS256);
        validation.set_audience(&self.expected_audiences);

        let tenant_issuer = format!("https://login.microsoftonline.com/{}/v2.0", self.tenant_id);
        if self.allow_multi_tenant {
            let issuers = [
                tenant_issuer,
                "https://login.microsoftonline.com/common/v2.0".to_string(),
            ];
            validation.set_issuer(&issuers);
        } else {
            validation.set_issuer(&[tenant_issuer]);
        }

        let token_data =
            jsonwebtoken::decode::<serde_json::Value>(token, &decoding_key, &validation)?;
        let claims = token_data.claims;

        let agent_id = claims
            .get("azp")
            .or_else(|| claims.get("appid"))
            .and_then(|v| v.as_str())
            .unwrap_or("unknown")
            .to_string();

        let display_name = claims
            .get("app_displayname")
            .or_else(|| claims.get("name"))
            .and_then(|v| v.as_str())
            .unwrap_or(&agent_id)
            .to_string();

        let tenant_id = claims
            .get("tid")
            .and_then(|v| v.as_str())
            .unwrap_or(&self.tenant_id)
            .to_string();

        let scopes = if let Some(roles) = claims.get("roles") {
            roles
                .as_array()
                .map(|arr| {
                    arr.iter()
                        .filter_map(|v| v.as_str().map(|s| s.to_string()))
                        .collect()
                })
                .unwrap_or_default()
        } else if let Some(scp) = claims.get("scp") {
            scp.as_str()
                .map(|s| s.split(' ').map(|s| s.to_string()).collect())
                .unwrap_or_default()
        } else {
            vec![]
        };

        let mut tags = HashMap::new();
        for key in ["agent_type", "blueprint_id", "sponsor"] {
            if let Some(val) = claims.get(key).and_then(|v| v.as_str()) {
                tags.insert(key.to_string(), val.to_string());
            }
        }

        let on_behalf_of = claims
            .get("oid")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string())
            .filter(|_| claims.get("azp").is_some());

        let expires_at = claims
            .get("exp")
            .and_then(|v| v.as_i64())
            .and_then(|ts| chrono::DateTime::from_timestamp(ts, 0));

        Ok(AgentIdentity {
            agent_id,
            display_name,
            agent_type: tags.get("agent_type").cloned(),
            blueprint_id: tags.get("blueprint_id").cloned(),
            sponsor: tags.get("sponsor").cloned(),
            on_behalf_of,
            tenant_id,
            scopes,
            expires_at,
            tags,
        })
    }
}

#[async_trait]
impl AuthProvider for EntraAgentIdProvider {
    async fn check_password(&self, user: &str, password: &str) -> bool {
        match self.validate_token(password).await {
            Ok(identity) => {
                tracing::info!(
                    agent_id = %identity.agent_id,
                    display_name = %identity.display_name,
                    agent_type = ?identity.agent_type,
                    tenant = %identity.tenant_id,
                    "Agent authenticated via Entra Agent ID"
                );
                true
            }
            Err(e) => {
                tracing::debug!(user = %user, error = %e, "Agent ID token validation failed");
                false
            }
        }
    }

    fn principal_for_user(&self, user: &str) -> String {
        format!("agent:{}", user)
    }
}