eric / rustwarden

use axum::{
    Extension, Json,
    extract::{Path, Query, State},
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Redirect, Response},
};
use sea_orm::DatabaseConnection;
use serde::Deserialize;
use std::sync::Arc;
use tower_sessions::Session;

use crate::auth::jwt::JwtService;
use crate::config::auth::AuthConfig;
use crate::config::oauth::OAuthConfig;
use crate::error::AppError;
use crate::hooks::RustwardenHooks;
use crate::services::oauth::OAuthService;
use crate::services::oauth::providers::ProviderRegistry;
use crate::services::refresh_token_service::RefreshTokenService;
use crate::services::user_preferences_service::UserPreferencesService;
use crate::services::user_service::UserService;

const OAUTH_STATE_KEY: &str = "oauth_state";
const OAUTH_NONCE_KEY: &str = "oauth_nonce";
const OAUTH_PKCE_KEY: &str = "oauth_pkce_verifier";
const OAUTH_PROVIDER_KEY: &str = "oauth_provider";

/// Query parameters received on the OAuth callback URL.
#[derive(Deserialize)]
pub struct OAuthCallbackQuery {
    /// Authorization code returned by the provider on success.
    pub code: Option<String>,
    /// CSRF state value to verify against the session.
    pub state: Option<String>,
    /// OAuth error code returned by the provider on failure.
    pub error: Option<String>,
    /// Human-readable description of the OAuth error.
    pub error_description: Option<String>,
}

/// GET /api/v1/auth/oauth/:provider/redirect
///
/// Initiates the OAuth flow by redirecting the user to the provider's auth page.
/// Stores CSRF state, nonce, and PKCE verifier in the session.
pub async fn oauth_redirect(
    State(db): State<DatabaseConnection>,
    session: Session,
    Extension(registry): Extension<Arc<ProviderRegistry>>,
    Path(provider_name): Path<String>,
) -> Result<Response, AppError> {
    let oauth_config = OAuthConfig::load(&db).await;
    if !oauth_config.enabled {
        return Err(AppError::Validation("OAuth is not enabled".to_string()));
    }

    let provider = registry.get(&provider_name).ok_or_else(|| {
        AppError::NotFound(format!("OAuth provider '{}' not found", provider_name))
    })?;

    let redirect_uri = oauth_config.callback_url(&provider_name);
    let auth_url = provider.authorization_url(&redirect_uri)?;

    // Store security params in session
    session
        .insert(OAUTH_STATE_KEY, &auth_url.csrf_state)
        .await
        .map_err(|e| AppError::Internal(format!("Session error: {}", e)))?;
    session
        .insert(OAUTH_PROVIDER_KEY, &provider_name)
        .await
        .map_err(|e| AppError::Internal(format!("Session error: {}", e)))?;

    if let Some(ref nonce) = auth_url.nonce {
        session
            .insert(OAUTH_NONCE_KEY, nonce)
            .await
            .map_err(|e| AppError::Internal(format!("Session error: {}", e)))?;
    }
    if let Some(ref verifier) = auth_url.pkce_verifier {
        session
            .insert(OAUTH_PKCE_KEY, verifier)
            .await
            .map_err(|e| AppError::Internal(format!("Session error: {}", e)))?;
    }

    tracing::info!(
        "OAuth redirect to {} for provider {}",
        &auth_url.url[..50.min(auth_url.url.len())],
        provider_name
    );
    Ok(Redirect::temporary(&auth_url.url).into_response())
}

/// GET /api/v1/auth/oauth/:provider/callback
///
/// Handles the OAuth callback from the provider.
/// Validates state, exchanges code for tokens, finds or creates user,
/// generates JWT + refresh token, sets cookies, redirects to /home.
// Axum handler: each argument is a distinct extractor; factoring them into a
// combined extractor would obscure the contract for marginal benefit.
#[allow(clippy::too_many_arguments)]
pub async fn oauth_callback(
    State(db): State<DatabaseConnection>,
    headers: HeaderMap,
    session: Session,
    Extension(registry): Extension<Arc<ProviderRegistry>>,
    Extension(auth_config): Extension<Arc<AuthConfig>>,
    hooks: Option<Extension<Arc<dyn RustwardenHooks>>>,
    Path(provider_name): Path<String>,
    Query(params): Query<OAuthCallbackQuery>,
) -> Result<Response, AppError> {
    // SEC-C2: Rate limiting — surface the proper 429 + Retry-After response
    // instead of masking it as a generic 400.
    if let Err(resp) = crate::middleware::rate_limit::check_auth_rate_limit(&headers) {
        return Ok(resp);
    }

    // Check for error from provider
    if let Some(ref error) = params.error {
        tracing::warn!(
            "OAuth error from {}: {} - {}",
            provider_name,
            error,
            params.error_description.as_deref().unwrap_or("no details")
        );
        return Ok(Redirect::to("/login?error=oauth_denied").into_response());
    }

    let code = params
        .code
        .as_deref()
        .ok_or_else(|| AppError::Validation("Missing authorization code".to_string()))?;

    // Validate CSRF state
    let session_state: Option<String> = session
        .get(OAUTH_STATE_KEY)
        .await
        .map_err(|e| AppError::Internal(format!("Session error: {}", e)))?;
    let session_provider: Option<String> = session
        .get(OAUTH_PROVIDER_KEY)
        .await
        .map_err(|e| AppError::Internal(format!("Session error: {}", e)))?;

    let expected_state = session_state
        .ok_or_else(|| AppError::Validation("Missing OAuth state in session".to_string()))?;
    let received_state = params
        .state
        .as_deref()
        .ok_or_else(|| AppError::Validation("Missing state parameter".to_string()))?;

    // SEC: Constant-time comparison for state
    use subtle::ConstantTimeEq;
    let state_match: bool = expected_state
        .as_bytes()
        .ct_eq(received_state.as_bytes())
        .into();
    if !state_match {
        tracing::warn!("OAuth state mismatch for provider {}", provider_name);
        return Err(AppError::Validation(
            "OAuth state mismatch (possible CSRF attack)".to_string(),
        ));
    }

    if session_provider.as_deref() != Some(&provider_name) {
        return Err(AppError::Validation("Provider mismatch".to_string()));
    }

    // Retrieve nonce and PKCE verifier from session
    let nonce: Option<String> = session.get(OAUTH_NONCE_KEY).await.unwrap_or(None);
    let pkce_verifier: Option<String> = session.get(OAUTH_PKCE_KEY).await.unwrap_or(None);

    // Clean up session oauth data
    session.remove::<String>(OAUTH_STATE_KEY).await.ok();
    session.remove::<String>(OAUTH_NONCE_KEY).await.ok();
    session.remove::<String>(OAUTH_PKCE_KEY).await.ok();
    session.remove::<String>(OAUTH_PROVIDER_KEY).await.ok();

    let oauth_config = OAuthConfig::load(&db).await;
    let redirect_uri = oauth_config.callback_url(&provider_name);

    let provider = registry
        .get(&provider_name)
        .ok_or_else(|| AppError::NotFound(format!("Provider '{}' not found", provider_name)))?;

    // Exchange code for user info
    let user_info = provider
        .exchange_code(
            code,
            &redirect_uri,
            pkce_verifier.as_deref(),
            nonce.as_deref(),
        )
        .await?;

    // Find or create user
    let (user, is_new_user) =
        OAuthService::find_or_create_user(&db, &provider_name, &user_info).await?;

    // Generate JWT
    let roles = UserService::get_user_roles(&db, user.id).await?;
    let permissions =
        crate::services::permission_service::PermissionService::get_permissions_for_user(
            &db, user.id,
        )
        .await?;

    let access_token = JwtService::generate_token(
        &auth_config,
        user.id,
        &user.username,
        roles.clone(),
        permissions,
    )?;

    // Generate refresh token
    let _created_rt = RefreshTokenService::create_refresh_token(
        &db,
        user.id,
        auth_config.refresh_token_expiration_days,
    )
    .await?;

    // SEC-M1: Regenerate session ID
    if let Err(e) = session.cycle_id().await {
        tracing::warn!("Failed to cycle session ID: {}", e);
    }

    // Build cookies
    let expires_in_seconds = auth_config.jwt_expiration_minutes * 60;
    let same_site = match auth_config.cookie_same_site.to_lowercase().as_str() {
        "strict" => axum_extra::extract::cookie::SameSite::Strict,
        "lax" => axum_extra::extract::cookie::SameSite::Lax,
        "none" => axum_extra::extract::cookie::SameSite::None,
        _ => axum_extra::extract::cookie::SameSite::Lax,
    };

    let jwt_cookie = axum_extra::extract::cookie::Cookie::build(("jwt_token", access_token))
        .http_only(true)
        .same_site(same_site)
        .secure(auth_config.cookie_secure)
        .path("/")
        .max_age(time::Duration::seconds(expires_in_seconds));

    let mut response = Redirect::to("/home").into_response();
    if let Ok(hv) = jwt_cookie.to_string().parse() {
        response
            .headers_mut()
            .insert(axum::http::header::SET_COOKIE, hv);
    }

    // Set theme/locale cookies from preferences
    if let Ok(Some(prefs)) = UserPreferencesService::get_preferences(&db, user.id).await {
        let theme_cookie = axum_extra::extract::cookie::Cookie::build(("theme", prefs.theme))
            .http_only(false)
            .same_site(axum_extra::extract::cookie::SameSite::Lax)
            .path("/")
            .max_age(time::Duration::seconds(31536000));
        if let Ok(hv) = theme_cookie.to_string().parse() {
            response
                .headers_mut()
                .append(axum::http::header::SET_COOKIE, hv);
        }

        let locale_cookie = axum_extra::extract::cookie::Cookie::build(("locale", prefs.language))
            .http_only(false)
            .same_site(axum_extra::extract::cookie::SameSite::Lax)
            .path("/")
            .max_age(time::Duration::seconds(31536000));
        if let Ok(hv) = locale_cookie.to_string().parse() {
            response
                .headers_mut()
                .append(axum::http::header::SET_COOKIE, hv);
        }
    }

    tracing::info!(
        "User {} logged in via OAuth ({}) [new_user={}]",
        user.username,
        provider_name,
        is_new_user
    );

    // Fire hooks
    if let Some(Extension(ref hooks)) = hooks {
        if let Err(e) = hooks.on_user_login(&db, user.id, &user.username).await {
            tracing::warn!("on_user_login hook error: {}", e);
        }
        if let Err(e) = hooks
            .on_oauth_login(&db, user.id, &provider_name, is_new_user)
            .await
        {
            tracing::warn!("on_oauth_login hook error: {}", e);
        }
    }

    Ok(response)
}

/// GET /api/v1/auth/oauth/providers
///
/// Returns the list of enabled OAuth providers (for the login page to display buttons).
pub async fn list_providers(
    State(db): State<DatabaseConnection>,
    Extension(registry): Extension<Arc<ProviderRegistry>>,
) -> Json<Vec<serde_json::Value>> {
    let oauth_config = OAuthConfig::load(&db).await;
    if !oauth_config.enabled {
        return Json(vec![]);
    }

    let providers: Vec<serde_json::Value> = registry
        .list()
        .into_iter()
        .map(|(name, display)| {
            serde_json::json!({
                "name": name,
                "display_name": display
            })
        })
        .collect();

    Json(providers)
}

/// GET /api/v1/auth/oauth/accounts
///
/// Returns the list of OAuth accounts linked to the authenticated user.
pub async fn list_user_oauth_accounts(
    State(db): State<DatabaseConnection>,
    auth_user: crate::auth::extractors::AuthUser,
) -> Result<Json<Vec<serde_json::Value>>, AppError> {
    let accounts = OAuthService::find_user_oauth_accounts(&db, auth_user.user_id).await?;
    let result: Vec<serde_json::Value> = accounts
        .into_iter()
        .map(|a| {
            serde_json::json!({
                "id": a.id,
                "provider": a.provider,
                "email": a.email,
                "display_name": a.display_name,
                "created_at": a.created_at.to_string(),
            })
        })
        .collect();
    Ok(Json(result))
}

/// DELETE /api/v1/auth/oauth/accounts/:id
///
/// Unlinks an OAuth account from the authenticated user.
pub async fn unlink_oauth_account(
    State(db): State<DatabaseConnection>,
    auth_user: crate::auth::extractors::AuthUser,
    Path(account_id): Path<uuid::Uuid>,
) -> Result<StatusCode, AppError> {
    OAuthService::unlink_account_by_id(&db, auth_user.user_id, account_id).await?;
    Ok(StatusCode::NO_CONTENT)
}