eric / rustwarden

//! Service for managing refresh tokens
//!
//! This service handles the lifecycle of refresh tokens:
//! - Generation of secure random tokens
//! - Storage and retrieval from database
//! - Token rotation on use
//! - Reuse detection (security)
//! - Token revocation

use crate::error::AppError;
use crate::models::refresh_token;
use chrono::{DateTime, Duration, Utc};
use rand::{Rng, distr::Alphanumeric};
use sea_orm::{
    ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, PaginatorTrait, QueryFilter,
    QueryOrder, Set,
};
use sha2::{Digest, Sha256};
use uuid::Uuid;

/// Active session information for display
#[derive(Debug, Clone, serde::Serialize)]
pub struct ActiveSession {
    /// Unique identifier of the refresh token / session.
    pub id: Uuid,
    /// Optional device or user-agent description.
    pub device_info: Option<String>,
    /// IP address from which the session was created.
    pub ip_address: Option<String>,
    /// Timestamp of the last activity associated with this session.
    pub last_activity: DateTime<Utc>,
    /// Timestamp when this session was originally created.
    pub created_at: DateTime<Utc>,
    /// Whether this session corresponds to the currently authenticated request.
    #[serde(default)]
    pub is_current: bool,
}

/// Result of creating a refresh token: the DB model (with hash) + plain token for the client.
pub struct CreatedRefreshToken {
    /// The persisted refresh token model (stores the SHA-256 hash, not the plain token).
    pub model: refresh_token::Model,
    /// The plaintext token to be sent to the client (never stored directly).
    pub plain_token: String,
}

/// Service for refresh token operations
pub struct RefreshTokenService;

impl RefreshTokenService {
    /// SEC-L1: Hash a token with SHA-256 before storage/lookup.
    fn hash_token(token: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(token.as_bytes());
        hex::encode(hasher.finalize())
    }

    /// Generate a secure random refresh token
    ///
    /// Generates a cryptographically secure random token using `rand::thread_rng()`.
    /// The token is 64 characters long, composed of alphanumeric characters.
    ///
    /// # Returns
    ///
    /// A secure random token string
    ///
    /// # Security
    ///
    /// - Uses `rand::thread_rng()` (cryptographically secure PRNG)
    /// - 64 characters = ~380 bits of entropy (very secure)
    /// - Alphanumeric only (URL-safe)
    ///
    /// # Examples
    ///
    /// ```
    /// use rustwarden_core::services::refresh_token_service::RefreshTokenService;
    ///
    /// let token = RefreshTokenService::generate_token();
    /// assert_eq!(token.len(), 64);
    /// ```
    pub fn generate_token() -> String {
        rand::rng()
            .sample_iter(&Alphanumeric)
            .take(64)
            .map(char::from)
            .collect()
    }

    /// Create a new refresh token in the database
    ///
    /// Generates a secure token, stores it in the database with expiration.
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    /// * `user_id` - User ID for whom to create the token
    /// * `expiration_days` - Number of days until token expires (typically 7)
    ///
    /// # Returns
    ///
    /// The created refresh token model
    ///
    /// # Errors
    ///
    /// Returns `AppError::Database` if database operation fails
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use rustwarden_core::services::refresh_token_service::RefreshTokenService;
    /// use uuid::Uuid;
    /// use sea_orm::DatabaseConnection;
    ///
    /// async fn create_token(db: DatabaseConnection, user_id: Uuid) {
    ///     let token = RefreshTokenService::create_refresh_token(&db, user_id, 7)
    ///         .await
    ///         .unwrap();
    ///     println!("Token: {}", token.plain_token);
    /// }
    /// ```
    /// SEC-M4: Maximum active (non-revoked, non-expired) refresh tokens per user.
    const MAX_ACTIVE_TOKENS_PER_USER: u64 = 10;

    /// SEC-L1: Stores the SHA-256 hash of the token in the database.
    /// SEC-M4: Revokes oldest tokens if the per-user limit is exceeded.
    /// Returns the plain token (for the client) along with the DB model.
    pub async fn create_refresh_token(
        db: &DatabaseConnection,
        user_id: Uuid,
        expiration_days: i64,
    ) -> Result<CreatedRefreshToken, AppError> {
        // SEC-M4: Enforce per-user token limit by revoking the oldest excess tokens
        let now_chrono = Utc::now();
        let active_tokens: Vec<refresh_token::Model> = refresh_token::Entity::find()
            .filter(refresh_token::Column::UserId.eq(user_id))
            .filter(refresh_token::Column::Revoked.eq(false))
            .filter(refresh_token::Column::ExpiresAt.gt(now_chrono))
            .order_by_asc(refresh_token::Column::CreatedAt)
            .all(db)
            .await?;

        if active_tokens.len() as u64 >= Self::MAX_ACTIVE_TOKENS_PER_USER {
            let excess = active_tokens.len() as u64 - Self::MAX_ACTIVE_TOKENS_PER_USER + 1;
            for old_token in active_tokens.iter().take(excess as usize) {
                let mut active: refresh_token::ActiveModel = old_token.clone().into();
                active.revoked = Set(true);
                active.revoked_at = Set(Some(now_chrono));
                active.update(db).await?;
            }
            tracing::info!(
                "SEC-M4: Revoked {} excess refresh tokens for user {}",
                excess,
                user_id
            );
        }

        let plain_token = Self::generate_token();
        let token_hash = Self::hash_token(&plain_token);
        let expires_at = now_chrono + Duration::days(expiration_days);

        let refresh_token = refresh_token::ActiveModel {
            id: Set(Uuid::new_v4()),
            user_id: Set(user_id),
            token: Set(token_hash),
            expires_at: Set(expires_at),
            created_at: Set(now_chrono),
            revoked: Set(false),
            revoked_at: Set(None),
            replaced_by: Set(None),
        };

        let model = refresh_token.insert(db).await?;

        tracing::debug!(
            "Created refresh token for user {} (expires in {} days)",
            user_id,
            expiration_days
        );

        Ok(CreatedRefreshToken { model, plain_token })
    }

    /// Validate and use a refresh token
    ///
    /// Validates the token (exists, not expired, not revoked) and returns it.
    /// The token should then be marked as used/replaced by the caller.
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    /// * `token` - Refresh token string to validate
    ///
    /// # Returns
    ///
    /// The validated refresh token model
    ///
    /// # Errors
    ///
    /// - `AppError::InvalidToken` - Token not found or malformed
    /// - `AppError::TokenExpired` - Token has expired
    /// - `AppError::InvalidToken` - Token has been revoked
    /// - `AppError::TokenReused` - Token has already been used (rotation)
    ///
    /// # Security
    ///
    /// This method performs all necessary validations but does NOT automatically
    /// mark the token as used. The caller must call `mark_as_replaced()` after
    /// successfully generating new tokens to complete the rotation.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use rustwarden_core::services::refresh_token_service::RefreshTokenService;
    /// use sea_orm::DatabaseConnection;
    ///
    /// async fn validate(db: DatabaseConnection, token_str: &str) {
    ///     match RefreshTokenService::validate_refresh_token(&db, token_str).await {
    ///         Ok(token) => println!("Valid token for user {}", token.user_id),
    ///         Err(e) => println!("Invalid token: {}", e),
    ///     }
    /// }
    /// ```
    /// SEC-L1: Hashes the input token with SHA-256 before DB lookup.
    pub async fn validate_refresh_token(
        db: &DatabaseConnection,
        token: &str,
    ) -> Result<refresh_token::Model, AppError> {
        let token_hash = Self::hash_token(token);

        let refresh_token = refresh_token::Entity::find()
            .filter(refresh_token::Column::Token.eq(token_hash))
            .one(db)
            .await?
            .ok_or_else(|| {
                tracing::warn!("Refresh token not found in database");
                AppError::InvalidToken
            })?;

        // Check if token has been revoked
        if refresh_token.revoked {
            tracing::warn!(
                "Attempt to use revoked refresh token for user {}",
                refresh_token.user_id
            );
            return Err(AppError::InvalidToken);
        }

        // Check if token has expired
        if refresh_token.is_expired() {
            tracing::debug!("Refresh token expired");
            return Err(AppError::TokenExpired);
        }

        // Check if token has already been used (reuse detection)
        if refresh_token.has_been_used() {
            tracing::error!(
                "SECURITY ALERT: Refresh token reuse detected for user {}",
                refresh_token.user_id
            );
            return Err(AppError::TokenReused);
        }

        tracing::debug!("Refresh token validated for user {}", refresh_token.user_id);

        Ok(refresh_token)
    }

    /// Mark a refresh token as replaced by a new token
    ///
    /// Updates the token's `replaced_by` field to indicate it has been used
    /// and replaced during rotation. This is used for reuse detection.
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    /// * `old_token_id` - ID of the token being replaced
    /// * `new_token_id` - ID of the new token
    ///
    /// # Returns
    ///
    /// The updated token model
    ///
    /// # Errors
    ///
    /// Returns `AppError::Database` if update fails
    pub async fn mark_as_replaced(
        db: &DatabaseConnection,
        old_token_id: Uuid,
        new_token_id: Uuid,
    ) -> Result<refresh_token::Model, AppError> {
        // Find the old token
        let old_token = refresh_token::Entity::find_by_id(old_token_id)
            .one(db)
            .await?
            .ok_or_else(|| AppError::NotFound("Token not found".to_string()))?;

        // Update it with the new token ID
        let mut active_model: refresh_token::ActiveModel = old_token.into();
        active_model.replaced_by = Set(Some(new_token_id));

        let updated = active_model.update(db).await?;

        tracing::debug!(
            "Marked token {} as replaced by {}",
            old_token_id,
            new_token_id
        );

        Ok(updated)
    }

    /// Revoke all refresh tokens for a user
    ///
    /// Marks all active tokens for a user as revoked. This is used:
    /// - On logout (revoke all sessions)
    /// - On password change (security measure)
    /// - On suspicious activity (e.g., token reuse detected)
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    /// * `user_id` - User ID whose tokens should be revoked
    ///
    /// # Returns
    ///
    /// Number of tokens revoked
    ///
    /// # Errors
    ///
    /// Returns `AppError::Database` if operation fails
    ///
    /// # Security
    ///
    /// This is a critical security operation that should be called when:
    /// - Token reuse is detected (potential theft)
    /// - User explicitly logs out from all devices
    /// - Password is changed
    /// - Account is compromised
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use rustwarden_core::services::refresh_token_service::RefreshTokenService;
    /// use uuid::Uuid;
    /// use sea_orm::DatabaseConnection;
    ///
    /// async fn revoke_all(db: DatabaseConnection, user_id: Uuid) {
    ///     let count = RefreshTokenService::revoke_all_tokens_for_user(&db, user_id)
    ///         .await
    ///         .unwrap();
    ///     println!("Revoked {} tokens", count);
    /// }
    /// ```
    pub async fn revoke_all_tokens_for_user(
        db: &DatabaseConnection,
        user_id: Uuid,
    ) -> Result<u64, AppError> {
        let now = Utc::now();

        // Find all non-revoked tokens for this user
        let tokens = refresh_token::Entity::find()
            .filter(refresh_token::Column::UserId.eq(user_id))
            .filter(refresh_token::Column::Revoked.eq(false))
            .all(db)
            .await?;

        let count = tokens.len() as u64;

        // Revoke each token
        for token in tokens {
            let mut active_model: refresh_token::ActiveModel = token.into();
            active_model.revoked = Set(true);
            active_model.revoked_at = Set(Some(now));
            active_model.update(db).await?;
        }

        tracing::info!("Revoked {} refresh tokens for user {}", count, user_id);

        Ok(count)
    }

    /// Get all active (non-revoked, non-expired) tokens for a user
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    /// * `user_id` - User ID
    ///
    /// # Returns
    ///
    /// Vector of active refresh tokens
    ///
    /// # Errors
    ///
    /// Returns `AppError::Database` if query fails
    #[allow(dead_code)] // Admin utility for monitoring user sessions
    pub async fn get_active_tokens_for_user(
        db: &DatabaseConnection,
        user_id: Uuid,
    ) -> Result<Vec<refresh_token::Model>, AppError> {
        let now = Utc::now();

        let tokens = refresh_token::Entity::find()
            .filter(refresh_token::Column::UserId.eq(user_id))
            .filter(refresh_token::Column::Revoked.eq(false))
            .filter(refresh_token::Column::ExpiresAt.gt(now))
            .all(db)
            .await?;

        Ok(tokens)
    }

    /// Clean up expired tokens from the database
    ///
    /// Deletes all tokens that have expired. This is a maintenance operation
    /// that should be run periodically (e.g., daily cron job).
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    ///
    /// # Returns
    ///
    /// Number of tokens deleted
    ///
    /// # Errors
    ///
    /// Returns `AppError::Database` if operation fails
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use rustwarden_core::services::refresh_token_service::RefreshTokenService;
    /// use sea_orm::DatabaseConnection;
    ///
    /// async fn cleanup(db: DatabaseConnection) {
    ///     let count = RefreshTokenService::cleanup_expired_tokens(&db)
    ///         .await
    ///         .unwrap();
    ///     println!("Cleaned up {} expired tokens", count);
    /// }
    /// ```
    #[allow(dead_code)] // Maintenance task, should be run as cron job
    pub async fn cleanup_expired_tokens(db: &DatabaseConnection) -> Result<u64, AppError> {
        let now = Utc::now();

        // Find all expired tokens
        let expired_tokens = refresh_token::Entity::find()
            .filter(refresh_token::Column::ExpiresAt.lt(now))
            .all(db)
            .await?;

        let count = expired_tokens.len() as u64;

        // Delete them
        for token in expired_tokens {
            let active_model: refresh_token::ActiveModel = token.into();
            active_model.delete(db).await?;
        }

        if count > 0 {
            tracing::info!("Cleaned up {} expired refresh tokens", count);
        }

        Ok(count)
    }

    /// Count total refresh tokens for a user
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    /// * `user_id` - User ID
    ///
    /// # Returns
    ///
    /// Number of tokens (active + expired + revoked)
    ///
    /// # Errors
    ///
    /// Returns `AppError::Database` if query fails
    #[allow(dead_code)] // Admin utility for monitoring
    pub async fn count_tokens_for_user(
        db: &DatabaseConnection,
        user_id: Uuid,
    ) -> Result<u64, AppError> {
        let count = refresh_token::Entity::find()
            .filter(refresh_token::Column::UserId.eq(user_id))
            .count(db)
            .await?;

        Ok(count)
    }

    /// Count active sessions (non-revoked, non-expired refresh tokens)
    ///
    /// This method counts all active sessions across all users.
    /// Used for admin dashboard statistics.
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    ///
    /// # Returns
    ///
    /// Number of active sessions
    ///
    /// # Errors
    ///
    /// Returns `AppError::Database` if query fails
    pub async fn count_active_sessions(db: &DatabaseConnection) -> Result<u64, AppError> {
        let now = Utc::now();

        let count = refresh_token::Entity::find()
            .filter(refresh_token::Column::Revoked.eq(false))
            .filter(refresh_token::Column::ExpiresAt.gt(now))
            .count(db)
            .await?;

        Ok(count)
    }

    /// List all active sessions for a user (Feature 009 - Phase 9)
    ///
    /// Returns all non-revoked, non-expired refresh tokens for the specified user.
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    /// * `user_id` - UUID of the user
    ///
    /// # Returns
    ///
    /// Vector of active session information
    ///
    /// # Errors
    ///
    /// Returns `AppError::Database` if query fails
    pub async fn list_user_sessions(
        db: &DatabaseConnection,
        user_id: Uuid,
    ) -> Result<Vec<ActiveSession>, AppError> {
        let now = Utc::now();

        let tokens = refresh_token::Entity::find()
            .filter(refresh_token::Column::UserId.eq(user_id))
            .filter(refresh_token::Column::Revoked.eq(false))
            .filter(refresh_token::Column::ExpiresAt.gt(now))
            .all(db)
            .await?;

        let sessions: Vec<ActiveSession> = tokens
            .into_iter()
            .map(|token| ActiveSession {
                id: token.id,
                device_info: None, // TODO: Store user_agent in refresh_token table
                ip_address: None,  // TODO: Store IP in refresh_token table
                last_activity: token.created_at,
                created_at: token.created_at,
                is_current: false, // Will be set by handler
            })
            .collect();

        Ok(sessions)
    }

    /// Revoke a specific session (Feature 009 - Phase 9)
    ///
    /// Marks a refresh token as revoked. Optionally blacklists associated JWT.
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    /// * `token_id` - UUID of the refresh token to revoke
    ///
    /// # Returns
    ///
    /// Ok(()) if revocation successful
    ///
    /// # Errors
    ///
    /// Returns `AppError::NotFound` if token doesn't exist
    /// Returns `AppError::Database` if database operation fails
    pub async fn revoke_session(db: &DatabaseConnection, token_id: Uuid) -> Result<(), AppError> {
        // Find the token
        let token = refresh_token::Entity::find_by_id(token_id)
            .one(db)
            .await?
            .ok_or_else(|| AppError::NotFound("Session not found".to_string()))?;

        // Mark as revoked
        let mut active_token: refresh_token::ActiveModel = token.into();
        active_token.revoked = Set(true);
        active_token.update(db).await?;

        tracing::info!("Session {} revoked", token_id);

        // TODO: Blacklist associated JWT if applicable

        Ok(())
    }

    /// Delete a specific refresh token (Feature 009 - Phase 9)
    ///
    /// Permanently deletes a refresh token from the database.
    ///
    /// # Arguments
    ///
    /// * `db` - Database connection
    /// * `token_id` - UUID of the refresh token to delete
    ///
    /// # Returns
    ///
    /// Ok(()) if deletion successful
    ///
    /// # Errors
    ///
    /// Returns `AppError::Database` if database operation fails
    pub async fn delete_session(db: &DatabaseConnection, token_id: Uuid) -> Result<(), AppError> {
        refresh_token::Entity::delete_by_id(token_id)
            .exec(db)
            .await?;

        tracing::info!("Session {} deleted", token_id);
        Ok(())
    }
}

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

    #[test]
    fn test_generate_token() {
        let token1 = RefreshTokenService::generate_token();
        let token2 = RefreshTokenService::generate_token();

        // Verify length
        assert_eq!(token1.len(), 64);
        assert_eq!(token2.len(), 64);

        // Verify uniqueness (extremely unlikely to be the same)
        assert_ne!(token1, token2);

        // Verify alphanumeric only
        assert!(token1.chars().all(|c| c.is_alphanumeric()));
    }

    #[test]
    fn test_token_entropy() {
        // Generate multiple tokens and verify they're all different
        let tokens: Vec<String> = (0..100)
            .map(|_| RefreshTokenService::generate_token())
            .collect();

        // All should be unique
        let unique_count = tokens
            .iter()
            .collect::<std::collections::HashSet<_>>()
            .len();
        assert_eq!(unique_count, 100);
    }

    // Note: Database operations are tested in integration tests
}