wme_client/

auth.rs

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

use parking_lot::Mutex;
use serde::Deserialize;

use crate::{ClientError, HttpTransport, Result};

/// Authentication tokens returned from login.
#[derive(Debug, Clone)]
pub struct Tokens {
    /// ID token (JWT with user claims)
    pub id_token: String,
    /// Access token for API requests
    pub access_token: String,
    /// Refresh token for obtaining new access tokens
    pub refresh_token: String,
    /// Token expiration time
    pub expires_at: Instant,
}

/// Thread-safe token manager with lazy refresh.
#[derive(Clone)]
pub struct TokenManager {
    state: Arc<Mutex<TokenState>>,
    transport: Arc<dyn HttpTransport>,
    config: AuthConfig,
}

#[derive(Debug, Clone)]
struct TokenState {
    tokens: Option<Tokens>,
}

/// Authentication configuration.
#[derive(Debug, Clone)]
pub struct AuthConfig {
    /// Username for authentication
    pub username: String,
    /// Password for authentication (only used for initial login)
    pub password: String,
    /// Authentication endpoint URL
    pub auth_url: String,
}

impl TokenManager {
    /// Create a new token manager.
    pub fn new(transport: Arc<dyn HttpTransport>, config: AuthConfig) -> Self {
        Self {
            state: Arc::new(Mutex::new(TokenState { tokens: None })),
            transport,
            config,
        }
    }

    /// Get a valid access token, refreshing if necessary.
    pub async fn get_access_token(&self) -> Result<String> {
        // Check if we need to refresh (without holding lock across await)
        let needs_refresh = {
            let state = self.state.lock();
            match &state.tokens {
                Some(tokens) => Instant::now() + Duration::from_secs(300) > tokens.expires_at,
                None => true,
            }
        };

        if needs_refresh {
            self.refresh().await?;
        }

        let state = self.state.lock();
        state
            .tokens
            .as_ref()
            .map(|t| t.access_token.clone())
            .ok_or(ClientError::Auth("No tokens available".to_string()))
    }

    /// Get the ID token (contains user claims).
    pub fn id_token(&self) -> Option<String> {
        let state = self.state.lock();
        state.tokens.as_ref().map(|t| t.id_token.clone())
    }

    /// Get all tokens if available.
    ///
    /// Returns the tokens including access_token, refresh_token, and expiration.
    /// This is useful for CLI applications that need to store tokens after login.
    pub fn get_tokens(&self) -> Option<Tokens> {
        let state = self.state.lock();
        state.tokens.clone()
    }

    /// Set tokens from external source (e.g., config file).
    ///
    /// This allows the CLI to restore tokens from a previous session.
    /// The tokens will be used for authentication and refreshed automatically when needed.
    pub fn set_tokens(&self, tokens: Tokens) {
        let mut state = self.state.lock();
        state.tokens = Some(tokens);
    }

    /// Authenticate and obtain tokens.
    ///
    /// This method performs authentication if needed:
    /// - If tokens exist and are valid, does nothing
    /// - If tokens are expired, attempts to refresh
    /// - If no tokens exist, performs fresh login
    ///
    /// After calling this, you can use [`TokenManager::get_tokens`] to retrieve the tokens for storage.
    pub async fn authenticate(&self) -> Result<()> {
        // Simply call get_access_token which handles all the auth logic
        let _ = self.get_access_token().await?;
        Ok(())
    }

    /// Refresh the access token using the refresh token endpoint.
    ///
    /// This method attempts to refresh the access token using the stored refresh token.
    /// If successful, the new access token and ID token are stored, while the refresh
    /// token remains unchanged.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - No refresh token is available
    /// - The refresh request fails
    /// - The response cannot be parsed
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use wme_client::auth::TokenManager;
    /// # async fn example(token_manager: &TokenManager) -> Result<(), Box<dyn std::error::Error>> {
    /// token_manager.refresh().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn refresh(&self) -> Result<()> {
        // Check if we have a refresh token (without holding lock across await)
        let has_refresh_token = {
            let state = self.state.lock();
            state
                .tokens
                .as_ref()
                .map(|t| !t.refresh_token.is_empty())
                .unwrap_or(false)
        };

        if has_refresh_token {
            // Try to use refresh token endpoint
            return self.refresh_with_token().await;
        }

        // No refresh token available, do fresh login
        let tokens = self.login().await?;

        // Store the new tokens
        let mut state = self.state.lock();
        state.tokens = Some(tokens);

        Ok(())
    }

    /// Refresh using the /v1/token-refresh endpoint.
    async fn refresh_with_token(&self) -> Result<()> {
        let refresh_token = {
            let state = self.state.lock();
            state
                .tokens
                .as_ref()
                .map(|t| t.refresh_token.clone())
                .ok_or(ClientError::Auth("No refresh token available".to_string()))?
        };

        let url = format!("{}/v1/token-refresh", self.config.auth_url);

        let body = serde_json::json!({
            "username": self.config.username,
            "refresh_token": refresh_token,
        });

        let mut headers = std::collections::HashMap::new();
        headers.insert("Content-Type".to_string(), "application/json".to_string());

        let response = self
            .transport
            .request(
                reqwest::Method::POST,
                &url,
                Some(headers),
                Some(body.to_string()),
            )
            .await?;

        if !response.status().is_success() {
            // If refresh fails, try fresh login
            if response.status().as_u16() == 401 {
                return self.login().await.map(|tokens| {
                    let mut state = self.state.lock();
                    state.tokens = Some(tokens);
                });
            }

            return Err(ClientError::Auth(format!(
                "Token refresh failed: {}",
                response.status()
            )));
        }

        let refresh_response: TokenRefreshResponse =
            response.json().await.map_err(ClientError::from)?;

        // Update tokens (keep the same refresh token)
        let mut state = self.state.lock();
        if let Some(existing) = &mut state.tokens {
            existing.id_token = refresh_response.id_token;
            existing.access_token = refresh_response.access_token;
            existing.expires_at = Instant::now() + Duration::from_secs(refresh_response.expires_in);
        }

        Ok(())
    }

    /// Perform initial login using /v1/login endpoint.
    async fn login(&self) -> Result<Tokens> {
        let url = format!("{}/v1/login", self.config.auth_url);

        let body = serde_json::json!({
            "username": self.config.username,
            "password": self.config.password,
        });

        let mut headers = std::collections::HashMap::new();
        headers.insert("Content-Type".to_string(), "application/json".to_string());

        tracing::debug!(
            "Login request to {} with username: {}",
            url,
            self.config.username
        );

        let response = self
            .transport
            .request(
                reqwest::Method::POST,
                &url,
                Some(headers),
                Some(body.to_string()),
            )
            .await?;

        if !response.status().is_success() {
            let status = response.status();
            let body_text = response.text().await.unwrap_or_default();
            // Log full response body at debug level only (may contain sensitive info)
            tracing::debug!("Login failed with status {}: {}", status, body_text);
            return Err(ClientError::Auth(format!("Login failed: {}", status)));
        }

        let body_text = response.text().await.map_err(ClientError::from)?;
        tracing::debug!("Login response: {}", body_text);

        let login_response: LoginResponse = serde_json::from_str(&body_text).map_err(|e| {
            tracing::debug!("Failed to parse login response: {}", e);
            ClientError::from(e)
        })?;

        Ok(Tokens {
            id_token: login_response.id_token,
            access_token: login_response.access_token,
            refresh_token: login_response.refresh_token,
            expires_at: Instant::now() + Duration::from_secs(login_response.expires_in),
        })
    }

    /// Revoke the current token using /v1/token-revoke endpoint.
    ///
    /// This invalidates the refresh token on the server, preventing
    /// any future token refreshes. The local token state is cleared
    /// regardless of whether the server request succeeds.
    ///
    /// # Errors
    ///
    /// Returns an error if no tokens are available to revoke or if
    /// the HTTP request fails.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use wme_client::auth::TokenManager;
    /// # async fn example(token_manager: &TokenManager) -> Result<(), Box<dyn std::error::Error>> {
    /// token_manager.revoke_token().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn revoke_token(&self) -> Result<()> {
        let refresh_token = {
            let state = self.state.lock();
            state
                .tokens
                .as_ref()
                .map(|t| t.refresh_token.clone())
                .ok_or(ClientError::Auth("No tokens to revoke".to_string()))?
        };

        let url = format!("{}/v1/token-revoke", self.config.auth_url);

        let body = serde_json::json!({
            "refresh_token": refresh_token,
        });

        let mut headers = std::collections::HashMap::new();
        headers.insert("Content-Type".to_string(), "application/json".to_string());

        let _response = self
            .transport
            .request(
                reqwest::Method::POST,
                &url,
                Some(headers),
                Some(body.to_string()),
            )
            .await?;

        // Clear local state regardless of server response
        self.clear_state();

        Ok(())
    }

    /// Clear internal state (force refresh on next use).
    pub fn clear_state(&self) {
        let mut state = self.state.lock();
        state.tokens = None;
    }

    /// Trigger forgot password email.
    ///
    /// Sends a password reset email to the user with a confirmation code.
    /// The code can be used with [`TokenManager::forgot_password_confirm`] to set a new password.
    ///
    /// # Arguments
    ///
    /// * `username` - The username to send the reset email to
    ///
    /// # Errors
    ///
    /// Returns an error if the request fails or the user is not found.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use wme_client::auth::TokenManager;
    /// # async fn example(token_manager: &TokenManager) -> Result<(), Box<dyn std::error::Error>> {
    /// token_manager.forgot_password("user@example.com").await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn forgot_password(&self, username: &str) -> Result<()> {
        let url = format!("{}/v1/forgot-password", self.config.auth_url);

        let body = serde_json::json!({
            "username": username,
        });

        let mut headers = std::collections::HashMap::new();
        headers.insert("Content-Type".to_string(), "application/json".to_string());

        let response = self
            .transport
            .request(
                reqwest::Method::POST,
                &url,
                Some(headers),
                Some(body.to_string()),
            )
            .await?;

        if !response.status().is_success() {
            return Err(ClientError::Auth(format!(
                "Forgot password request failed: {}",
                response.status()
            )));
        }

        Ok(())
    }

    /// Confirm forgot password with code.
    ///
    /// Completes the password reset flow by providing the confirmation code
    /// received via email and the new password.
    ///
    /// # Arguments
    ///
    /// * `username` - The username being reset
    /// * `password` - The new password to set
    /// * `code` - The confirmation code from the email
    ///
    /// # Errors
    ///
    /// Returns an error if the code is invalid, expired, or the request fails.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use wme_client::auth::TokenManager;
    /// # async fn example(token_manager: &TokenManager) -> Result<(), Box<dyn std::error::Error>> {
    /// token_manager.forgot_password_confirm("user@example.com", "new_password", "123456").await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn forgot_password_confirm(
        &self,
        username: &str,
        password: &str,
        code: &str,
    ) -> Result<()> {
        let url = format!("{}/v1/forgot-password-confirm", self.config.auth_url);

        let body = serde_json::json!({
            "username": username,
            "password": password,
            "code": code,
        });

        let mut headers = std::collections::HashMap::new();
        headers.insert("Content-Type".to_string(), "application/json".to_string());

        let response = self
            .transport
            .request(
                reqwest::Method::POST,
                &url,
                Some(headers),
                Some(body.to_string()),
            )
            .await?;

        if !response.status().is_success() {
            return Err(ClientError::Auth(format!(
                "Password reset confirmation failed: {}",
                response.status()
            )));
        }

        Ok(())
    }

    /// Change password (requires authentication).
    ///
    /// Changes the password for the authenticated user. The user must
    /// provide their current password for verification.
    ///
    /// # Arguments
    ///
    /// * `previous_password` - The current password
    /// * `proposed_password` - The new password to set
    ///
    /// # Errors
    ///
    /// Returns an error if the current password is incorrect or the request fails.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use wme_client::auth::TokenManager;
    /// # async fn example(token_manager: &TokenManager) -> Result<(), Box<dyn std::error::Error>> {
    /// token_manager.change_password("old_password", "new_password").await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn change_password(
        &self,
        previous_password: &str,
        proposed_password: &str,
    ) -> Result<()> {
        let url = format!("{}/v1/change-password", self.config.auth_url);
        let headers = self.auth_headers().await?;

        let body = serde_json::json!({
            "previous_password": previous_password,
            "proposed_password": proposed_password,
        });

        let response = self
            .transport
            .request(reqwest::Method::POST, &url, headers, Some(body.to_string()))
            .await?;

        if !response.status().is_success() {
            return Err(ClientError::Auth(format!(
                "Password change failed: {}",
                response.status()
            )));
        }

        Ok(())
    }

    /// Set new password for NEW_PASSWORD_REQUIRED challenge.
    ///
    /// Completes the authentication flow when a user is required to set
    /// a new password on first login or after a password reset.
    ///
    /// # Arguments
    ///
    /// * `username` - The username
    /// * `session` - The session identifier from the authentication challenge
    /// * `new_password` - The new password to set
    ///
    /// # Errors
    ///
    /// Returns an error if the session is invalid, password doesn't meet
    /// requirements, or the request fails.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use wme_client::auth::TokenManager;
    /// # async fn example(token_manager: &TokenManager) -> Result<(), Box<dyn std::error::Error>> {
    /// token_manager.set_new_password("user@example.com", "session_id", "new_secure_password").await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn set_new_password(
        &self,
        username: &str,
        session: &str,
        new_password: &str,
    ) -> Result<()> {
        let url = format!("{}/v1/set-new-password", self.config.auth_url);

        let body = serde_json::json!({
            "username": username,
            "session": session,
            "new_password": new_password,
        });

        let mut headers = std::collections::HashMap::new();
        headers.insert("Content-Type".to_string(), "application/json".to_string());

        let response = self
            .transport
            .request(
                reqwest::Method::POST,
                &url,
                Some(headers),
                Some(body.to_string()),
            )
            .await?;

        if !response.status().is_success() {
            return Err(ClientError::Auth(format!(
                "Set new password failed: {}",
                response.status()
            )));
        }

        Ok(())
    }

    /// Get authentication headers.
    async fn auth_headers(&self) -> Result<Option<HashMap<String, String>>> {
        let token = self.get_access_token().await?;
        let mut headers = HashMap::new();
        headers.insert("Authorization".to_string(), format!("Bearer {}", token));
        Ok(Some(headers))
    }
}

/// Login response from /v1/login endpoint.
#[derive(Debug, Deserialize)]
struct LoginResponse {
    id_token: String,
    access_token: String,
    refresh_token: String,
    expires_in: u64,
}

/// Token refresh response from /v1/token-refresh endpoint.
#[derive(Debug, Deserialize)]
struct TokenRefreshResponse {
    id_token: String,
    access_token: String,
    expires_in: u64,
}