qssh 0.0.2-alpha

//! Rate limiting for authentication attempts
//!
//! Prevents brute force attacks by limiting authentication attempts
//! per IP address with exponential backoff

use crate::{Result, QsshError};
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use serde::{Serialize, Deserialize};

/// Configuration for rate limiting
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RateLimitConfig {
    /// Maximum attempts before temporary ban
    pub max_attempts: u32,
    /// Time window for counting attempts (seconds)
    pub window_seconds: u64,
    /// Base ban duration (seconds) - doubles with each violation
    pub base_ban_seconds: u64,
    /// Maximum ban duration (seconds)
    pub max_ban_seconds: u64,
    /// Whether to use exponential backoff
    pub exponential_backoff: bool,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            max_attempts: 5,           // 5 attempts
            window_seconds: 60,         // per minute
            base_ban_seconds: 60,       // 1 minute base ban
            max_ban_seconds: 3600,      // 1 hour max ban
            exponential_backoff: true,  // Use exponential backoff
        }
    }
}

/// Track authentication attempts for an IP
#[derive(Debug, Clone)]
struct AttemptRecord {
    /// Number of failed attempts
    failed_attempts: u32,
    /// Time of first attempt in current window
    window_start: Instant,
    /// Time when ban expires (if banned)
    ban_until: Option<Instant>,
    /// Number of times this IP has been banned
    ban_count: u32,
    /// Last successful authentication time
    last_success: Option<Instant>,
}

impl AttemptRecord {
    fn new() -> Self {
        Self {
            failed_attempts: 0,
            window_start: Instant::now(),
            ban_until: None,
            ban_count: 0,
            last_success: None,
        }
    }

    /// Check if currently banned
    fn is_banned(&self) -> bool {
        if let Some(ban_until) = self.ban_until {
            ban_until > Instant::now()
        } else {
            false
        }
    }

    /// Get remaining ban time
    fn ban_remaining(&self) -> Option<Duration> {
        if let Some(ban_until) = self.ban_until {
            let now = Instant::now();
            if ban_until > now {
                Some(ban_until - now)
            } else {
                None
            }
        } else {
            None
        }
    }

    /// Reset the window if expired
    fn reset_if_expired(&mut self, window: Duration) {
        let now = Instant::now();
        if now.duration_since(self.window_start) > window {
            self.failed_attempts = 0;
            self.window_start = now;
        }
    }
}

/// Rate limiter for authentication attempts
pub struct RateLimiter {
    config: RateLimitConfig,
    attempts: Arc<RwLock<HashMap<IpAddr, AttemptRecord>>>,
}

impl RateLimiter {
    /// Create new rate limiter with default config
    pub fn new() -> Self {
        Self::with_config(RateLimitConfig::default())
    }

    /// Create rate limiter with custom config
    pub fn with_config(config: RateLimitConfig) -> Self {
        Self {
            config,
            attempts: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Check if an IP is allowed to attempt authentication
    pub async fn check_allowed(&self, ip: IpAddr) -> Result<()> {
        let mut attempts = self.attempts.write().await;
        let record = attempts.entry(ip).or_insert_with(AttemptRecord::new);

        // Check if currently banned
        if record.is_banned() {
            let remaining = record.ban_remaining()
                .map(|d| d.as_secs())
                .unwrap_or(0);

            return Err(QsshError::RateLimited(format!(
                "Too many failed attempts. Try again in {} seconds",
                remaining
            )));
        }

        // Reset window if expired
        let window = Duration::from_secs(self.config.window_seconds);
        record.reset_if_expired(window);

        // Check if over limit
        if record.failed_attempts >= self.config.max_attempts {
            // Apply ban
            let ban_duration = if self.config.exponential_backoff {
                // Exponential backoff: base * 2^ban_count
                let multiplier = 2_u64.saturating_pow(record.ban_count);
                let duration = self.config.base_ban_seconds.saturating_mul(multiplier);
                duration.min(self.config.max_ban_seconds)
            } else {
                self.config.base_ban_seconds
            };

            record.ban_until = Some(Instant::now() + Duration::from_secs(ban_duration));
            record.ban_count = record.ban_count.saturating_add(1);

            return Err(QsshError::RateLimited(format!(
                "Too many failed attempts. Banned for {} seconds",
                ban_duration
            )));
        }

        Ok(())
    }

    /// Record a failed authentication attempt
    pub async fn record_failure(&self, ip: IpAddr) {
        let mut attempts = self.attempts.write().await;
        let record = attempts.entry(ip).or_insert_with(AttemptRecord::new);

        // Reset window if expired
        let window = Duration::from_secs(self.config.window_seconds);
        record.reset_if_expired(window);

        // Increment failure count
        record.failed_attempts = record.failed_attempts.saturating_add(1);

        log::warn!(
            "Failed auth attempt from {} ({}/{} in window)",
            ip,
            record.failed_attempts,
            self.config.max_attempts
        );
    }

    /// Record a successful authentication
    pub async fn record_success(&self, ip: IpAddr) {
        let mut attempts = self.attempts.write().await;
        let record = attempts.entry(ip).or_insert_with(AttemptRecord::new);

        // Reset on success
        record.failed_attempts = 0;
        record.last_success = Some(Instant::now());

        // Don't reset ban_count immediately - decay it over time
        // This prevents attackers from resetting their penalty with stolen creds

        log::info!("Successful auth from {}", ip);
    }

    /// Clean up old records (housekeeping)
    pub async fn cleanup(&self) {
        let mut attempts = self.attempts.write().await;
        let now = Instant::now();
        let window = Duration::from_secs(self.config.window_seconds * 10); // Keep 10x window

        // Remove old records that haven't been used recently
        attempts.retain(|ip, record| {
            // Keep if:
            // - Currently banned
            // - Had activity in last 10 windows
            // - Has recent success

            if record.is_banned() {
                return true;
            }

            if now.duration_since(record.window_start) < window {
                return true;
            }

            if let Some(last_success) = record.last_success {
                if now.duration_since(last_success) < Duration::from_secs(86400) {
                    return true; // Keep successful logins for 24 hours
                }
            }

            log::debug!("Cleaning up rate limit record for {}", ip);
            false
        });
    }

    /// Get current statistics
    pub async fn get_stats(&self) -> RateLimitStats {
        let attempts = self.attempts.read().await;

        let total_tracked = attempts.len();
        let currently_banned = attempts.values()
            .filter(|r| r.is_banned())
            .count();

        let high_risk = attempts.iter()
            .filter(|(_, r)| r.ban_count > 2)
            .map(|(ip, _)| *ip)
            .collect();

        RateLimitStats {
            total_tracked,
            currently_banned,
            high_risk_ips: high_risk,
        }
    }

    /// Manually unban an IP (admin function)
    pub async fn unban(&self, ip: IpAddr) -> Result<()> {
        let mut attempts = self.attempts.write().await;

        if let Some(record) = attempts.get_mut(&ip) {
            record.ban_until = None;
            record.failed_attempts = 0;
            log::info!("Manually unbanned {}", ip);
            Ok(())
        } else {
            Err(QsshError::NotFound(format!("No record for IP {}", ip)))
        }
    }

    /// Check if IP is suspicious (for monitoring)
    pub async fn is_suspicious(&self, ip: IpAddr) -> bool {
        let attempts = self.attempts.read().await;

        if let Some(record) = attempts.get(&ip) {
            // Suspicious if:
            // - Has been banned multiple times
            // - Currently has high failure rate
            // - Is currently banned

            record.ban_count > 1 ||
            record.failed_attempts > self.config.max_attempts / 2 ||
            record.is_banned()
        } else {
            false
        }
    }
}

/// Statistics from the rate limiter
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RateLimitStats {
    /// Total IPs being tracked
    pub total_tracked: usize,
    /// IPs currently banned
    pub currently_banned: usize,
    /// High risk IPs (banned multiple times)
    pub high_risk_ips: Vec<IpAddr>,
}

/// Background task to periodically clean up old records
pub async fn cleanup_task(limiter: Arc<RateLimiter>) {
    let mut interval = tokio::time::interval(Duration::from_secs(300)); // Every 5 minutes

    loop {
        interval.tick().await;
        limiter.cleanup().await;
    }
}

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

    #[tokio::test]
    async fn test_rate_limiting() {
        let config = RateLimitConfig {
            max_attempts: 3,
            window_seconds: 60,
            base_ban_seconds: 10,
            max_ban_seconds: 100,
            exponential_backoff: true,
        };

        let limiter = RateLimiter::with_config(config);
        let ip = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 100));

        // First attempts should succeed
        assert!(limiter.check_allowed(ip).await.is_ok());
        limiter.record_failure(ip).await;

        assert!(limiter.check_allowed(ip).await.is_ok());
        limiter.record_failure(ip).await;

        assert!(limiter.check_allowed(ip).await.is_ok());
        limiter.record_failure(ip).await;

        // Fourth attempt should be rate limited
        assert!(limiter.check_allowed(ip).await.is_err());
    }

    #[tokio::test]
    async fn test_exponential_backoff() {
        let config = RateLimitConfig {
            max_attempts: 2,
            window_seconds: 60,
            base_ban_seconds: 1,
            max_ban_seconds: 100,
            exponential_backoff: true,
        };

        let limiter = RateLimiter::with_config(config);
        let ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1));

        // Get banned once
        limiter.record_failure(ip).await;
        limiter.record_failure(ip).await;
        assert!(limiter.check_allowed(ip).await.is_err());

        // Wait for ban to expire
        tokio::time::sleep(Duration::from_secs(2)).await;

        // Get banned again - should be longer
        assert!(limiter.check_allowed(ip).await.is_ok());
        limiter.record_failure(ip).await;
        limiter.record_failure(ip).await;

        let err = limiter.check_allowed(ip).await.unwrap_err();
        assert!(err.to_string().contains("Banned for 2 seconds"));
    }

    #[tokio::test]
    async fn test_success_resets() {
        let limiter = RateLimiter::new();
        let ip = IpAddr::V4(Ipv4Addr::new(172, 16, 0, 1));

        // Some failures
        limiter.record_failure(ip).await;
        limiter.record_failure(ip).await;

        // Success should reset
        limiter.record_success(ip).await;

        // Should be able to try again
        assert!(limiter.check_allowed(ip).await.is_ok());
    }

    #[tokio::test]
    async fn test_cleanup() {
        let mut config = RateLimitConfig::default();
        config.window_seconds = 1; // Short window for testing

        let limiter = RateLimiter::with_config(config);
        let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));

        // Add a record
        limiter.record_failure(ip).await;

        // Should have one record
        let stats = limiter.get_stats().await;
        assert_eq!(stats.total_tracked, 1);

        // Wait for window to expire significantly
        tokio::time::sleep(Duration::from_secs(15)).await;

        // Cleanup should remove it
        limiter.cleanup().await;
        let stats = limiter.get_stats().await;
        assert_eq!(stats.total_tracked, 0);
    }
}