authkestra-engine 0.3.3

use crate::auth::{
    error::AuthError, state::Identity, state::OAuth2State, state::OAuthToken, ErasedOAuthFlow,
    OAuthProvider, UserMapper,
};
use crate::flow::{Flow, FlowContext, FlowResult};
use async_trait::async_trait;

/// Orchestrates the standard OAuth2 Authorization Code flow.
pub struct OAuth2Flow<P: OAuthProvider, M: UserMapper = ()> {
    provider: P,
    mapper: Option<M>,
    scopes: Vec<String>,
    use_pkce: bool,
}

#[async_trait]
impl<P: OAuthProvider + 'static, M: UserMapper + 'static> Flow for OAuth2Flow<P, M> {
    fn id(&self) -> &str {
        self.provider.provider_id()
    }

    async fn execute(&self, ctx: FlowContext) -> Result<FlowResult, AuthError> {
        if let Some(_code) = ctx.params.get("code") {
            let _received_state = ctx.params.get("state").ok_or(AuthError::CsrfMismatch)?;

            // In the new model, expected_state must be provided via some context.
            // For now, if it's missing from ctx, we might need to adjust FlowContext.
            // But ErasedOAuthFlow is what the adapters use.
            Err(AuthError::Token(
                "Direct Flow execution not updated for encrypted state".to_string(),
            ))
        } else {
            // Assume initiation if no code is present
            let scopes_str = ctx.params.get("scopes").map(|s| s.as_str()).unwrap_or("");
            let scopes_vec: Vec<&str> = if scopes_str.is_empty() {
                Vec::new()
            } else {
                scopes_str.split(',').collect()
            };

            let pkce_challenge = ctx.params.get("pkce_challenge").map(|s| s.as_str());
            let (url, _state) = self.initiate_login(&scopes_vec, pkce_challenge);
            Ok(FlowResult::Redirect(url))
        }
    }
}

#[async_trait]
impl<P: OAuthProvider + 'static, M: UserMapper + 'static> ErasedOAuthFlow for OAuth2Flow<P, M> {
    fn provider_id(&self) -> String {
        self.provider.provider_id().to_string()
    }

    fn initiate_login(
        &self,
        scopes: &[&str],
        pkce_challenge: Option<&str>,
    ) -> (String, OAuth2State) {
        let effective_scopes = if !scopes.is_empty() {
            scopes
        } else {
            &self
                .scopes
                .iter()
                .map(|s| s.as_str())
                .collect::<Vec<&str>>()
        };

        self.initiate_login(effective_scopes, pkce_challenge)
    }

    async fn finalize_login(
        &self,
        code: &str,
        received_state: &str,
        expected_state: &OAuth2State,
    ) -> Result<(Identity, OAuthToken), AuthError> {
        let (identity, token, _) = self
            .finalize_login(code, received_state, expected_state)
            .await?;
        Ok((identity, token))
    }
}

impl<P: OAuthProvider> OAuth2Flow<P, ()> {
    /// Create a new `OAuth2Flow` with the given provider.
    pub fn new(provider: P) -> Self {
        Self {
            provider,
            mapper: None,
            scopes: Vec::new(),
            use_pkce: true,
        }
    }
}

impl<P: OAuthProvider, M: UserMapper> OAuth2Flow<P, M> {
    /// Create a new `OAuth2Flow` with the given provider and user mapper.
    pub fn with_mapper(provider: P, mapper: M) -> Self {
        Self {
            provider,
            mapper: Some(mapper),
            scopes: Vec::new(),
            use_pkce: true,
        }
    }

    /// Set the scopes for the OAuth2 flow.
    pub fn with_scopes(mut self, scopes: Vec<impl Into<String>>) -> Self {
        self.scopes = scopes.into_iter().map(|s| s.into()).collect();
        self
    }

    /// Enable or disable PKCE for the OAuth2 flow.
    pub fn with_pkce(mut self, use_pkce: bool) -> Self {
        self.use_pkce = use_pkce;
        self
    }

    /// Generates the redirect URL and CSRF state.
    #[tracing::instrument(skip(self), fields(provider_id = %self.provider.provider_id()))]
    pub fn initiate_login(
        &self,
        scopes: &[&str],
        pkce_challenge: Option<&str>,
    ) -> (String, OAuth2State) {
        let state = uuid::Uuid::new_v4().to_string();
        let nonce = Some(uuid::Uuid::new_v4().to_string());

        let effective_scopes = if !scopes.is_empty() {
            scopes
        } else {
            &self
                .scopes
                .iter()
                .map(|s| s.as_str())
                .collect::<Vec<&str>>()
        };

        tracing::debug!(scopes = ?effective_scopes, "generating authorization URL");

        let url = self.provider.get_authorization_url(
            &state,
            effective_scopes,
            pkce_challenge,
            nonce.as_deref(),
        );

        let auth_state = OAuth2State {
            state: state.clone(),
            nonce,
            code_verifier: None, // Will be set by the caller if needed before encryption
            success_url: None,
            provider_id: self.provider.provider_id().to_string(),
            expires_at: chrono::Utc::now().timestamp() + 600,
        };

        tracing::info!("authorization login initiated successfully");
        (url, auth_state)
    }

    /// Completes the flow by exchanging the code.
    /// If a mapper is provided, it will also map the identity to a local user.
    #[tracing::instrument(skip(self, code, expected_state), fields(provider_id = %self.provider.provider_id()))]
    pub async fn finalize_login(
        &self,
        code: &str,
        received_state: &str,
        expected_state: &OAuth2State,
    ) -> Result<(Identity, OAuthToken, Option<M::LocalUser>), AuthError> {
        if received_state != expected_state.state {
            tracing::error!("CSRF mismatch: received state does not match expected state");
            return Err(AuthError::CsrfMismatch);
        }

        tracing::debug!("exchanging code for identity");
        let (identity, token) = self
            .provider
            .exchange_code_for_identity(
                code,
                expected_state.code_verifier.as_deref(),
                expected_state.nonce.as_deref(),
            )
            .await
            .map_err(|e| {
                tracing::error!(error = %e, "failed to exchange code for identity");
                e
            })?;

        tracing::info!(user_id = %identity.external_id, "successfully retrieved identity from provider");

        if let Some(expected_nonce) = expected_state.nonce.as_deref() {
            if identity.attributes.get("nonce").map(|s| s.as_str()) != Some(expected_nonce) {
                tracing::error!("nonce mismatch or missing in identity attributes");
                return Err(AuthError::Token("Nonce mismatch".to_string()));
            }
        }

        let local_user = if let Some(mapper) = &self.mapper {
            tracing::debug!("mapping user identity");
            Some(mapper.map_user(&identity).await.map_err(|e| {
                tracing::error!(error = %e, "failed to map user");
                e
            })?)
        } else {
            None
        };

        Ok((identity, token, local_user))
    }

    /// Refresh an access token using a refresh token.
    pub async fn refresh_access_token(&self, refresh_token: &str) -> Result<OAuthToken, AuthError> {
        self.provider.refresh_token(refresh_token).await
    }

    /// Revoke an access token.
    pub async fn revoke_token(&self, token: &str) -> Result<(), AuthError> {
        self.provider.revoke_token(token).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::auth::{Provider, ProviderConfig};
    use async_trait::async_trait;
    use std::collections::HashMap;

    struct MockProvider {
        id: String,
        auth_url: String,
        expected_code: String,
        identity: Identity,
        token: OAuthToken,
    }

    #[async_trait]
    impl Provider for MockProvider {
        async fn config(&self) -> ProviderConfig {
            ProviderConfig {
                id: self.id.clone(),
                name: self.id.clone(),
                extra: HashMap::new(),
            }
        }
    }

    #[async_trait]
    impl OAuthProvider for MockProvider {
        fn provider_id(&self) -> &str {
            &self.id
        }

        fn get_authorization_url(
            &self,
            state: &str,
            _scopes: &[&str],
            _code_challenge: Option<&str>,
            _nonce: Option<&str>,
        ) -> String {
            format!("{}?state={}", self.auth_url, state)
        }

        async fn exchange_code_for_identity(
            &self,
            code: &str,
            _code_verifier: Option<&str>,
            _nonce: Option<&str>,
        ) -> Result<(Identity, OAuthToken), AuthError> {
            if code == self.expected_code {
                Ok((self.identity.clone(), self.token.clone()))
            } else {
                Err(AuthError::Token("Invalid code".to_string()))
            }
        }
    }

    #[tokio::test]
    async fn test_oauth2_flow_initiate() {
        let provider = MockProvider {
            id: "mock".to_string(),
            auth_url: "http://mock/auth".to_string(),
            expected_code: "code123".to_string(),
            identity: Identity {
                provider_id: "mock".to_string(),
                external_id: "1".to_string(),
                email: None,
                username: None,
                attributes: HashMap::new(),
            },
            token: OAuthToken {
                access_token: "acc".to_string(),
                token_type: "Bearer".to_string(),
                expires_in: None,
                refresh_token: None,
                scope: None,
                id_token: None,
            },
        };

        let flow = OAuth2Flow::new(provider).with_scopes(vec!["scope1"]);
        let (url, state) = flow.initiate_login(&["scope2"], None);
        assert!(url.contains("http://mock/auth?state="));
        assert_eq!(state.provider_id, "mock");
    }

    #[tokio::test]
    async fn test_oauth2_flow_finalize_success() {
        let provider = MockProvider {
            id: "mock".to_string(),
            auth_url: "http://mock/auth".to_string(),
            expected_code: "code123".to_string(),
            identity: Identity {
                provider_id: "mock".to_string(),
                external_id: "1".to_string(),
                email: None,
                username: None,
                attributes: HashMap::new(),
            },
            token: OAuthToken {
                access_token: "acc".to_string(),
                token_type: "Bearer".to_string(),
                expires_in: None,
                refresh_token: None,
                scope: None,
                id_token: None,
            },
        };

        let flow = OAuth2Flow::new(provider);
        let expected_state = OAuth2State {
            state: "state123".to_string(),
            nonce: None,
            code_verifier: None,
            success_url: None,
            provider_id: "mock".to_string(),
            expires_at: 0,
        };

        let (ident, tok, _) = flow
            .finalize_login("code123", "state123", &expected_state)
            .await
            .unwrap();
        assert_eq!(ident.external_id, "1");
        assert_eq!(tok.access_token, "acc");
    }

    #[tokio::test]
    async fn test_oauth2_flow_finalize_csrf_mismatch() {
        let provider = MockProvider {
            id: "mock".to_string(),
            auth_url: "http://mock/auth".to_string(),
            expected_code: "code123".to_string(),
            identity: Identity {
                provider_id: "mock".to_string(),
                external_id: "1".to_string(),
                email: None,
                username: None,
                attributes: HashMap::new(),
            },
            token: OAuthToken {
                access_token: "acc".to_string(),
                token_type: "Bearer".to_string(),
                expires_in: None,
                refresh_token: None,
                scope: None,
                id_token: None,
            },
        };

        let flow = OAuth2Flow::new(provider);
        let expected_state = OAuth2State {
            state: "state123".to_string(),
            nonce: None,
            code_verifier: None,
            success_url: None,
            provider_id: "mock".to_string(),
            expires_at: 0,
        };

        let result = flow
            .finalize_login("code123", "wrong_state", &expected_state)
            .await;
        assert!(matches!(result, Err(AuthError::CsrfMismatch)));
    }
}