<?php
if ( ! defined( 'ABSPATH' ) ) { exit; }

class Vulnity_Brute_Force_Alert extends Vulnity_Alert_Base {

    const WINDOW_INDEX_OPTION = 'vulnity_bf_window_index';

    private $aggregation_window = 60; // 60 seconds window
    private $cooldown_period = 300;   // 5 minutes cooldown after alert
    private $initial_threshold = 3;   // Minimum attempts to start tracking
    private $windows_processed = false;
    
    // Updated severity thresholds
    private $severity_thresholds = array(
        'critical' => 41,  // 41+ attempts
        'high' => 21,      // 21-40 attempts
        'medium' => 6,     // 6-20 attempts
        'low' => 3         // 3-5 attempts
    );
    
    public function __construct() {
        $this->alert_type = 'brute_force';
        parent::__construct();
    }
    
    protected function register_hooks() {
        add_action('wp_login_failed', array($this, 'handle_login_failed'), 10, 2);
        add_action('wp_login', array($this, 'handle_successful_login'), 10, 2);

        // Fallback processing when WP-Cron is disabled
        add_action('init', array($this, 'maybe_process_expired_windows'));
        
        // Add cron to process pending windows
        add_action('vulnity_process_brute_force_windows', array($this, 'process_aggregation_windows'));
        
        // Schedule cron if not already scheduled
        if (!wp_next_scheduled('vulnity_process_brute_force_windows')) {
            wp_schedule_event(time(), 'vulnity_every_minute', 'vulnity_process_brute_force_windows');
        }
    }
    
    public function handle_login_failed($username, $error = null) {
        $ip = $this->get_client_ip();
        $current_time = time();
        
        // Check if IP is in cooldown
        if ($this->is_ip_in_cooldown($ip)) {
            // Still track attempts but don't create new alerts
            $this->add_attempt_to_window($ip, $username, $current_time);
            return;
        }
        
        // Get or create aggregation window for this IP
        $window_key = 'vulnity_bf_window_' . md5($ip);
        $window = get_transient($window_key);
        
        if ($window === false) {
            // No active window, check if we should start one
            $recent_attempts = $this->get_recent_attempts($ip);
            
            if (count($recent_attempts) >= ($this->initial_threshold - 1)) {
                // Start new aggregation window
                $window = array(
                    'ip' => $ip,
                    'started_at' => $current_time,
                    'expires_at' => $current_time + $this->aggregation_window,
                    'attempts' => array(
                        array(
                            'username' => $username,
                            'timestamp' => $current_time
                        )
                    ),
                    'usernames' => array($username),
                    'total_attempts' => 1
                );
                
                // Add previous recent attempts to window
                foreach ($recent_attempts as $attempt) {
                    $window['attempts'][] = $attempt;
                    $window['total_attempts']++;
                    if (!in_array($attempt['username'], $window['usernames'])) {
                        $window['usernames'][] = $attempt['username'];
                    }
                }
                
                set_transient($window_key, $window, $this->aggregation_window + 10);
                self::remember_window_key($window_key);
                vulnity_log('[Vulnity] Started brute force aggregation window for IP: ' . $ip); // Surface brute force lifecycle events for operators.
            } else {
                // Just track the attempt for future reference
                $this->track_attempt($ip, $username, $current_time);
            }
        } else {
            // Active window exists, add attempt to it
            $window['attempts'][] = array(
                'username' => $username,
                'timestamp' => $current_time
            );
            
            if (!in_array($username, $window['usernames'])) {
                $window['usernames'][] = $username;
            }
            
            $window['total_attempts']++;
            
            set_transient($window_key, $window, $this->aggregation_window + 10);
            self::remember_window_key($window_key);
        }
    }
    
    public function handle_successful_login($user_login, $user) {
        $ip = $this->get_client_ip();
        $window_key = 'vulnity_bf_window_' . md5($ip);
        $window = get_transient($window_key);
        
        if ($window !== false && $window['total_attempts'] >= $this->initial_threshold) {
            // Successful login after brute force attempts - Critical alert!
            $this->create_breach_alert($window, $user_login, $ip);
            
            // Clear the window
            delete_transient($window_key);
            self::forget_window_key($window_key);
            
            // Set cooldown
            $this->set_ip_cooldown($ip);
        }
        
        // Clear tracking for this IP
        $this->clear_tracking($ip);
    }
    
    /**
     * Process aggregation windows (called by cron every minute)
     */
    public function process_aggregation_windows() {
        $current_time = time();
        $window_keys = self::get_window_keys();

        if (empty($window_keys)) {
            return;
        }

        foreach ($window_keys as $window_key) {
            $window = get_transient($window_key);

            if ($window === false) {
                self::forget_window_key($window_key);
                continue;
            }

            if (!is_array($window) || !isset($window['expires_at'])) {
                continue;
            }

            if ($current_time < $window['expires_at']) {
                continue;
            }

            if ($window['total_attempts'] >= $this->initial_threshold) {
                $this->create_aggregated_alert($window);
                $this->set_ip_cooldown($window['ip']);
            }

            delete_transient($window_key);
            self::forget_window_key($window_key);
            $this->clear_tracking($window['ip']);

            vulnity_log('[Vulnity] Processed brute force window for IP: ' . $window['ip'] . ' with ' . $window['total_attempts'] . ' attempts');
        }
    }

    /**
     * Process aggregation windows during normal requests to handle sites without WP-Cron.
     */
    public function maybe_process_expired_windows() {
        if ($this->windows_processed) {
            return;
        }

        $window_keys = self::get_window_keys();

        if (empty($window_keys)) {
            $this->windows_processed = true;
            return;
        }

        $this->windows_processed = true;
        $this->process_aggregation_windows();
    }
    
    /**
     * Create aggregated alert after window expires
     */
    private function create_aggregated_alert($window) {
        $duration = $window['expires_at'] - $window['started_at'];
        $attempts_per_minute = round(($window['total_attempts'] / $duration) * 60, 2);
        
        // Determine severity based on total attempts
        $severity = $this->calculate_severity($window['total_attempts']);
        
        $this->create_alert(array(
            'severity' => $severity,
            'title' => 'Brute Force Attack Detected',
            'message' => sprintf(
                '%d failed login attempts from IP %s targeting %d username%s in %d seconds',
                $window['total_attempts'],
                $window['ip'],
                count($window['usernames']),
                count($window['usernames']) > 1 ? 's' : '',
                $duration
            ),
            'details' => array(
                'ip' => $window['ip'],
                'total_attempts' => $window['total_attempts'],
                'unique_usernames' => count($window['usernames']),
                'usernames' => array_slice($window['usernames'], 0, 10),
                'attack_started' => wp_date('Y-m-d H:i:s', $window['started_at']),
                'attack_ended' => wp_date('Y-m-d H:i:s', $window['expires_at']),
                'duration_seconds' => $duration,
                'attempts_per_minute' => $attempts_per_minute,
                'attack_pattern' => $this->analyze_attack_pattern($window),
                'peak_intensity' => $this->calculate_peak_intensity($window['attempts'])
            )
        ));
    }
    
    /**
     * Create critical alert for successful login after brute force
     */
    private function create_breach_alert($window, $user_login, $ip) {
        $this->create_alert(array(
            'severity' => 'critical',
            'title' => 'Successful Login After Brute Force Attack',
            'message' => sprintf(
                'BREACH: User "%s" successfully logged in after %d failed attempts from IP %s',
                $user_login,
                $window['total_attempts'],
                $ip
            ),
            'details' => array(
                'ip' => $ip,
                'successful_user' => $user_login,
                'total_attempts' => $window['total_attempts'],
                'unique_usernames_tried' => count($window['usernames']),
                'usernames_tried' => $window['usernames'],
                'attack_duration' => time() - $window['started_at'],
                'breach' => true,
                'risk_level' => 'critical',
                'recommended_action' => 'Immediately review account access and consider password reset'
            )
        ));
    }
    
    /**
     * Calculate severity based on attempt count
     */
    private function calculate_severity($total_attempts) {
        if ($total_attempts >= $this->severity_thresholds['critical']) {
            return 'critical';
        } elseif ($total_attempts >= $this->severity_thresholds['high']) {
            return 'high';
        } elseif ($total_attempts >= $this->severity_thresholds['medium']) {
            return 'medium';
        } else {
            return 'low';
        }
    }
    
    /**
     * Analyze attack pattern
     */
    private function analyze_attack_pattern($window) {
        $usernames_count = count($window['usernames']);
        $attempts_count = $window['total_attempts'];
        
        if ($usernames_count == 1) {
            return 'targeted_single_user';
        } elseif ($usernames_count < 5) {
            return 'targeted_multiple_users';
        } elseif (in_array('admin', $window['usernames']) || in_array('administrator', $window['usernames'])) {
            return 'dictionary_attack_common_usernames';
        } else {
            return 'broad_dictionary_attack';
        }
    }
    
    /**
     * Calculate peak intensity (max attempts in any 10-second period)
     */
    private function calculate_peak_intensity($attempts) {
        if (empty($attempts)) {
            return 0;
        }
        
        $timestamps = array_column($attempts, 'timestamp');
        sort($timestamps);
        
        $max_in_10_seconds = 0;
        $window_size = 10; // 10 seconds
        
        for ($i = 0; $i < count($timestamps); $i++) {
            $count = 1;
            $window_start = $timestamps[$i];
            
            for ($j = $i + 1; $j < count($timestamps); $j++) {
                if ($timestamps[$j] - $window_start <= $window_size) {
                    $count++;
                } else {
                    break;
                }
            }
            
            $max_in_10_seconds = max($max_in_10_seconds, $count);
        }
        
        return $max_in_10_seconds;
    }
    
    /**
     * Track individual attempt (before window starts)
     */
    private function track_attempt($ip, $username, $timestamp) {
        $tracking_key = 'vulnity_bf_track_' . md5($ip);
        $attempts = get_transient($tracking_key);
        
        if ($attempts === false) {
            $attempts = array();
        }
        
        $attempts[] = array(
            'username' => $username,
            'timestamp' => $timestamp
        );
        
        // Keep only last 10 attempts
        $attempts = array_slice($attempts, -10);
        
        set_transient($tracking_key, $attempts, 300); // Keep for 5 minutes
    }
    
    /**
     * Get recent attempts for IP
     */
    private function get_recent_attempts($ip) {
        $tracking_key = 'vulnity_bf_track_' . md5($ip);
        $attempts = get_transient($tracking_key);
        
        if ($attempts === false) {
            return array();
        }
        
        // Only return attempts from last 30 seconds
        $cutoff = time() - 30;
        return array_filter($attempts, function($attempt) use ($cutoff) {
            return $attempt['timestamp'] >= $cutoff;
        });
    }
    
    /**
     * Clear tracking for IP
     */
    private function clear_tracking($ip) {
        $tracking_key = 'vulnity_bf_track_' . md5($ip);
        delete_transient($tracking_key);
    }
    
    /**
     * Add attempt to existing window
     */
    private function add_attempt_to_window($ip, $username, $timestamp) {
        $window_key = 'vulnity_bf_window_' . md5($ip);
        $window = get_transient($window_key);
        
        if ($window !== false) {
            $window['attempts'][] = array(
                'username' => $username,
                'timestamp' => $timestamp
            );
            
            if (!in_array($username, $window['usernames'])) {
                $window['usernames'][] = $username;
            }
            
            $window['total_attempts']++;
            
            set_transient($window_key, $window, $this->aggregation_window + 10);
            self::remember_window_key($window_key);
        }
    }
    
    /**
     * Check if IP is in cooldown period
     */
    private function is_ip_in_cooldown($ip) {
        $cooldown_key = 'vulnity_bf_cooldown_' . md5($ip);
        return get_transient($cooldown_key) !== false;
    }
    
    /**
     * Set cooldown for IP
     */
    private function set_ip_cooldown($ip) {
        $cooldown_key = 'vulnity_bf_cooldown_' . md5($ip);
        set_transient($cooldown_key, true, $this->cooldown_period);
        vulnity_log('[Vulnity] Set cooldown for IP ' . $ip . ' for ' . $this->cooldown_period . ' seconds');
    }

    private static function get_window_keys() {
        $keys = get_option(self::WINDOW_INDEX_OPTION, array());

        if (!is_array($keys)) {
            $keys = array();
        }

        return $keys;
    }

    private static function remember_window_key($key) {
        $key = sanitize_key($key);
        $keys = self::get_window_keys();

        if (!in_array($key, $keys, true)) {
            $keys[] = $key;
            update_option(self::WINDOW_INDEX_OPTION, $keys, false);
        }
    }

    private static function forget_window_key($key) {
        $key = sanitize_key($key);
        $keys = self::get_window_keys();
        $new_keys = array_values(array_diff($keys, array($key)));

        if ($new_keys !== $keys) {
            update_option(self::WINDOW_INDEX_OPTION, $new_keys, false);
        }
    }
    
    protected function evaluate($data) {
        // This method is required by parent but not used in new implementation
        // All logic is handled in handle_login_failed and process_aggregation_windows
    }
}
