foxhunt/services/trading_service/src/jwt_revocation.rs

//! JWT Token Revocation System with Redis-backed Blacklist
//!
//! This module implements comprehensive JWT session revocation to address the critical
//! vulnerability where compromised tokens remain valid until expiration (CVSS 8.8).
//!
//! ## Features
//! - Redis-backed token blacklist with JTI tracking
//! - Immediate revocation capability for compromised tokens
//! - Refresh token mechanism for session continuity
//! - Admin endpoints for forced revocation
//! - Integration with existing JWT validation flow
//!
//! ## Architecture
//! ```text
//! JWT Validation Flow with Revocation:
//! 1. Extract JWT from request
//! 2. Decode and validate JWT structure
//! 3. Check JTI against Redis blacklist (CRITICAL)
//! 4. Validate expiration and claims
//! 5. Allow/Deny request
//!
//! Revocation Flow:
//! 1. Admin/User requests revocation
//! 2. Add JTI to Redis with TTL = remaining token lifetime
//! 3. Token immediately invalid on next validation
//! 4. Redis automatically cleans up expired blacklist entries
//! ```

use anyhow::{Context, Result};
use redis::{aio::ConnectionManager, AsyncCommands};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::{debug, error, info, warn};
use uuid::Uuid;

/// JWT Token ID (JTI) for unique token identification
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct Jti(pub String);

impl Jti {
    /// Generate a new random JTI
    pub fn new() -> Self {
        Self(Uuid::new_v4().to_string())
    }

    /// Parse JTI from string
    pub fn from_string(s: String) -> Self {
        Self(s)
    }

    /// Get JTI as string reference
    pub fn as_str(&self) -> &str {
        &self.0
    }

    /// Get Redis key for this JTI
    fn redis_key(&self) -> String {
        format!("jwt:blacklist:{}", self.0)
    }
}

impl Default for Jti {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Display for Jti {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Enhanced JWT claims with JTI and refresh token support
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct EnhancedJwtClaims {
    /// JWT ID for revocation tracking
    pub jti: String,
    /// Subject (user ID)
    pub sub: String,
    /// Issued at timestamp
    pub iat: u64,
    /// Expiration timestamp
    pub exp: u64,
    /// Not before timestamp
    pub nbf: u64,
    /// Issuer
    pub iss: String,
    /// Audience
    pub aud: String,
    /// User roles
    pub roles: Vec<String>,
    /// Additional permissions
    pub permissions: Vec<String>,
    /// Token type: "access" or "refresh"
    pub token_type: String,
    /// Session ID for tracking related tokens
    pub session_id: String,
}

impl EnhancedJwtClaims {
    /// Create new access token claims
    pub fn new_access_token(
        user_id: String,
        roles: Vec<String>,
        permissions: Vec<String>,
        issuer: String,
        audience: String,
        ttl_seconds: u64,
    ) -> Result<Self> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .context("System time error")?
            .as_secs();

        Ok(Self {
            jti: Jti::new().to_string(),
            sub: user_id,
            iat: now,
            exp: now + ttl_seconds,
            nbf: now,
            iss: issuer,
            aud: audience,
            roles,
            permissions,
            token_type: "access".to_string(),
            session_id: Uuid::new_v4().to_string(),
        })
    }

    /// Create new refresh token claims (longer TTL, limited permissions)
    pub fn new_refresh_token(
        user_id: String,
        issuer: String,
        audience: String,
        session_id: String,
        ttl_seconds: u64,
    ) -> Result<Self> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .context("System time error")?
            .as_secs();

        Ok(Self {
            jti: Jti::new().to_string(),
            sub: user_id,
            iat: now,
            exp: now + ttl_seconds,
            nbf: now,
            iss: issuer,
            aud: audience,
            roles: vec![], // Refresh tokens have no roles
            permissions: vec!["refresh_token".to_string()], // Only refresh permission
            token_type: "refresh".to_string(),
            session_id,
        })
    }

    /// Get remaining TTL in seconds
    pub fn remaining_ttl(&self) -> Result<u64> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .context("System time error")?
            .as_secs();

        if self.exp > now {
            Ok(self.exp - now)
        } else {
            Ok(0)
        }
    }
}

/// Token pair containing access and refresh tokens
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenPair {
    /// Access token (short-lived, full permissions)
    pub access_token: String,
    /// Refresh token (long-lived, limited to refresh permission)
    pub refresh_token: String,
    /// Access token expiration timestamp
    pub access_token_expires_at: u64,
    /// Refresh token expiration timestamp
    pub refresh_token_expires_at: u64,
    /// Token type (always "Bearer")
    pub token_type: String,
}

/// Revocation reason for audit logging
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RevocationReason {
    /// User-initiated logout
    UserLogout,
    /// Admin-forced revocation
    AdminRevocation,
    /// Suspicious activity detected
    SuspiciousActivity,
    /// Password change
    PasswordChange,
    /// Account locked
    AccountLocked,
    /// Token compromised
    TokenCompromised,
    /// Session timeout
    SessionTimeout,
    /// Other reason
    Other(String),
}

impl std::fmt::Display for RevocationReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::UserLogout => write!(f, "user_logout"),
            Self::AdminRevocation => write!(f, "admin_revocation"),
            Self::SuspiciousActivity => write!(f, "suspicious_activity"),
            Self::PasswordChange => write!(f, "password_change"),
            Self::AccountLocked => write!(f, "account_locked"),
            Self::TokenCompromised => write!(f, "token_compromised"),
            Self::SessionTimeout => write!(f, "session_timeout"),
            Self::Other(reason) => write!(f, "other:{}", reason),
        }
    }
}

/// Revocation metadata stored in Redis
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
struct RevocationMetadata {
    /// User ID who owned the token
    user_id: String,
    /// Reason for revocation
    reason: String,
    /// Who initiated the revocation (user_id or "admin" or "system")
    revoked_by: String,
    /// Timestamp of revocation
    revoked_at: u64,
    /// Client IP if available
    client_ip: Option<String>,
}

impl RevocationMetadata {
    /// Public accessor for user_id (for audit logging)
    pub fn user_id(&self) -> &str {
        &self.user_id
    }
    
    /// Public accessor for reason (for audit logging)
    pub fn reason(&self) -> &str {
        &self.reason
    }
    
    /// Public accessor for revoked_by (for audit logging)
    pub fn revoked_by(&self) -> &str {
        &self.revoked_by
    }
}

impl std::fmt::Debug for JwtRevocationService {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("JwtRevocationService")
            .field("config", &self.config)
            .finish()
    }
}

/// Redis-backed JWT revocation service
#[derive(Clone)]
pub struct JwtRevocationService {
    /// Redis connection manager
    redis: ConnectionManager,
    /// Configuration
    config: Arc<RevocationConfig>,
}

/// Revocation service configuration
#[derive(Debug, Clone, PartialEq)]
pub struct RevocationConfig {
    /// Redis key prefix for blacklist entries
    pub redis_prefix: String,
    /// Redis key prefix for user session tracking
    pub session_prefix: String,
    /// Enable audit logging
    pub enable_audit_logging: bool,
    /// Maximum tokens per user to track for bulk revocation
    pub max_tokens_per_user: usize,
}

impl Default for RevocationConfig {
    fn default() -> Self {
        Self {
            redis_prefix: "jwt:blacklist:".to_string(),
            session_prefix: "jwt:user_sessions:".to_string(),
            enable_audit_logging: true,
            max_tokens_per_user: 100,
        }
    }
}

impl JwtRevocationService {
    /// Create new revocation service
    pub async fn new(redis_url: &str, config: RevocationConfig) -> Result<Self> {
        let client = redis::Client::open(redis_url)
            .context("Failed to create Redis client for JWT revocation")?;

        let redis = ConnectionManager::new(client)
            .await
            .context("Failed to connect to Redis for JWT revocation")?;

        info!("JWT revocation service connected to Redis: {}", redis_url);

        Ok(Self {
            redis,
            config: Arc::new(config),
        })
    }

    /// Check if a token is revoked (blacklisted)
    pub async fn is_revoked(&self, jti: &Jti) -> Result<bool> {
        let key = jti.redis_key();
        let mut conn = self.redis.clone();

        let exists: bool = conn
            .exists(&key)
            .await
            .context("Failed to check token revocation status in Redis")?;

        if exists {
            debug!("Token {} is blacklisted (revoked)", jti);
        }

        Ok(exists)
    }

    /// Revoke a single token by JTI
    pub async fn revoke_token(
        &self,
        jti: &Jti,
        user_id: &str,
        ttl_seconds: u64,
        reason: RevocationReason,
        revoked_by: &str,
        client_ip: Option<String>,
    ) -> Result<()> {
        if ttl_seconds == 0 {
            warn!("Token {} already expired, skipping revocation", jti);
            return Ok(());
        }

        let key = jti.redis_key();
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .context("System time error")?
            .as_secs();

        let metadata = RevocationMetadata {
            user_id: user_id.to_string(),
            reason: reason.to_string(),
            revoked_by: revoked_by.to_string(),
            revoked_at: now,
            client_ip,
        };

        let metadata_json = serde_json::to_string(&metadata)
            .context("Failed to serialize revocation metadata")?;

        let mut conn = self.redis.clone();

        // Set the blacklist entry with TTL = remaining token lifetime
        // Redis will automatically clean up when token would have expired anyway
        let _: () = conn
            .set_ex(&key, metadata_json, ttl_seconds)
            .await
            .context("Failed to add token to Redis blacklist")?;

        if self.config.enable_audit_logging {
            info!(
                "Token revoked: jti={} user={} reason={} revoked_by={} ttl={}s",
                jti, user_id, reason, revoked_by, ttl_seconds
            );
        }

        // Track this token in user's session list for bulk revocation
        self.track_user_token(user_id, jti).await?;

        Ok(())
    }

    /// Track a token under a user's session list
    async fn track_user_token(&self, user_id: &str, jti: &Jti) -> Result<()> {
        let key = format!("{}{}", self.config.session_prefix, user_id);
        let mut conn = self.redis.clone();

        // Add JTI to user's token set
        let _: () = conn
            .sadd(&key, jti.as_str())
            .await
            .context("Failed to add token to user session tracking")?;

        // Limit the size of the set to prevent memory issues
        let set_size: usize = conn
            .scard(&key)
            .await
            .context("Failed to get user session set size")?;

        if set_size > self.config.max_tokens_per_user {
            warn!(
                "User {} has {} tracked tokens, trimming oldest",
                user_id, set_size
            );
            // In production, you might want to implement FIFO cleanup
        }

        Ok(())
    }

    /// Revoke all tokens for a specific user
    pub async fn revoke_all_user_tokens(
        &self,
        user_id: &str,
        reason: RevocationReason,
        revoked_by: &str,
    ) -> Result<usize> {
        let key = format!("{}{}", self.config.session_prefix, user_id);
        let mut conn = self.redis.clone();

        // Get all JTIs for this user
        let jtis: Vec<String> = conn
            .smembers(&key)
            .await
            .context("Failed to get user's token list from Redis")?;

        if jtis.is_empty() {
            info!("No tokens found for user {}, nothing to revoke", user_id);
            return Ok(0);
        }

        let mut revoked_count = 0;

        for jti_str in &jtis {
            let jti = Jti::from_string(jti_str.clone());
            let blacklist_key = jti.redis_key();

            // Check if already revoked
            let exists: bool = conn
                .exists(&blacklist_key)
                .await
                .context("Failed to check if token is already revoked")?;

            if exists {
                continue; // Already revoked
            }

            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .context("System time error")?
                .as_secs();

            let metadata = RevocationMetadata {
                user_id: user_id.to_string(),
                reason: reason.to_string(),
                revoked_by: revoked_by.to_string(),
                revoked_at: now,
                client_ip: None,
            };

            let metadata_json = serde_json::to_string(&metadata)
                .context("Failed to serialize revocation metadata")?;

            // Revoke with a generous TTL (24 hours) since we don't know the exact token expiration
            let ttl = 86400u64; // 24 hours

            let _: () = conn
                .set_ex(&blacklist_key, metadata_json, ttl)
                .await
                .context("Failed to add token to Redis blacklist")?;

            revoked_count += 1;
        }

        // Clean up the user's token tracking set
        let _: () = conn
            .del(&key)
            .await
            .context("Failed to delete user token tracking set")?;

        if self.config.enable_audit_logging {
            info!(
                "Revoked {} tokens for user {} (reason: {}, revoked_by: {})",
                revoked_count, user_id, reason, revoked_by
            );
        }

        Ok(revoked_count)
    }

    /// Get revocation metadata for a token
    pub async fn get_revocation_metadata(&self, jti: &Jti) -> Result<Option<RevocationMetadata>> {
        let key = jti.redis_key();
        let mut conn = self.redis.clone();

        let metadata_json: Option<String> = conn
            .get(&key)
            .await
            .context("Failed to get revocation metadata from Redis")?;

        match metadata_json {
            Some(ref json) => {
                let metadata: RevocationMetadata = serde_json::from_str(&json)
                    .context("Failed to deserialize revocation metadata")?;
                Ok(Some(metadata))
            }
            None => Ok(None),
        }
    }

    /// Get statistics about revoked tokens
    pub async fn get_statistics(&self) -> Result<RevocationStatistics> {
        let mut conn = self.redis.clone();

        // Count blacklist entries
        let pattern = format!("{}*", self.config.redis_prefix);
        let blacklist_count: usize = self.count_keys(&mut conn, &pattern).await?;

        // Count user session tracking entries
        let session_pattern = format!("{}*", self.config.session_prefix);
        let active_users: usize = self.count_keys(&mut conn, &session_pattern).await?;

        Ok(RevocationStatistics {
            revoked_tokens: blacklist_count,
            active_users,
        })
    }

    /// Helper to count keys matching a pattern
    async fn count_keys(&self, conn: &mut ConnectionManager, pattern: &str) -> Result<usize> {
        // Note: KEYS command is not recommended for production with large datasets
        // Consider using SCAN for production deployments
        let keys: Vec<String> = redis::cmd("KEYS")
            .arg(pattern)
            .query_async(conn)
            .await
            .context("Failed to count keys in Redis")?;

        Ok(keys.len())
    }

    /// Clean up expired entries (manual cleanup, Redis TTL handles most cases)
    pub async fn cleanup_expired(&self) -> Result<usize> {
        // Redis automatically deletes expired keys, but we can do manual cleanup
        // of the user session tracking sets
        let mut conn = self.redis.clone();
        let pattern = format!("{}*", self.config.session_prefix);

        let keys: Vec<String> = conn
            .keys(&pattern)
            .await
            .context("Failed to get session tracking keys")?;

        let mut cleaned = 0;

        for key in keys {
            // Get all JTIs in the set
            let jtis: Vec<String> = conn
                .smembers(&key)
                .await
                .context("Failed to get JTIs from session set")?;

            for jti_str in jtis {
                let jti = Jti::from_string(jti_str.clone());
                let blacklist_key = jti.redis_key();

                // Check if the blacklist entry still exists
                let exists: bool = conn
                    .exists(&blacklist_key)
                    .await
                    .context("Failed to check blacklist entry")?;

                // If blacklist entry doesn't exist, token has expired - remove from tracking
                if !exists {
                    let _: () = conn
                        .srem(&key, &jti_str)
                        .await
                        .context("Failed to remove expired JTI from session set")?;
                    cleaned += 1;
                }
            }

            // If the set is now empty, delete it
            let set_size: usize = conn
                .scard(&key)
                .await
                .context("Failed to get session set size")?;

            if set_size == 0 {
                let _: () = conn
                    .del(&key)
                    .await
                    .context("Failed to delete empty session set")?;
            }
        }

        debug!("Cleaned up {} expired token tracking entries", cleaned);
        Ok(cleaned)
    }
}

/// Statistics about revoked tokens
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RevocationStatistics {
    /// Number of currently revoked tokens
    pub revoked_tokens: usize,
    /// Number of active users with tracked sessions
    pub active_users: usize,
}

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

    #[test]
    fn test_jti_generation() {
        let jti1 = Jti::new();
        let jti2 = Jti::new();

        assert_ne!(jti1, jti2);
        assert!(jti1.as_str().len() > 0);
    }

    #[test]
    fn test_enhanced_jwt_claims_creation() {
        let claims = EnhancedJwtClaims::new_access_token(
            "user123".to_string(),
            vec!["trader".to_string()],
            vec!["trading.submit_order".to_string()],
            "foxhunt".to_string(),
            "trading-api".to_string(),
            3600,
        )
        .unwrap();

        assert_eq!(claims.sub, "user123");
        assert_eq!(claims.token_type, "access");
        assert!(!claims.jti.is_empty());
        assert!(claims.remaining_ttl().unwrap() > 3500); // Should be close to 3600
    }

    #[test]
    fn test_refresh_token_claims() {
        let session_id = Uuid::new_v4().to_string();
        let claims = EnhancedJwtClaims::new_refresh_token(
            "user123".to_string(),
            "foxhunt".to_string(),
            "trading-api".to_string(),
            session_id.clone(),
            86400, // 24 hours
        )
        .unwrap();

        assert_eq!(claims.sub, "user123");
        assert_eq!(claims.token_type, "refresh");
        assert_eq!(claims.session_id, session_id);
        assert_eq!(claims.permissions, vec!["refresh_token".to_string()]);
        assert!(claims.roles.is_empty());
    }
}