<?php
defined('ABSPATH') || exit;

class WPSEC_Firewall {
    private $whitelisted_ips = [];
    private $permanently_blocked_ips = [];
    private $is_active = false;
    private $blocked_count = 0;
    private $last_matched_rules = [];
    private static $instance = null;
    
    // Rules for logging - more generic to catch potential threats
    private $logging_rules = [
        'sql_injection' => [
            'pattern' => '/(?i)(\'|\%27|--|\%23|\#|\/\*|\*\/|\||\&|\;|union|select|insert|update|delete|drop|truncate|alter|create|rename|replace|database|table|column|into|exec|declare|cast|varchar|concat|ascii|substring|benchmark|sleep|null|like|script|having|group\s+by|order\s+by|where\s+|from\s+|limit\s+)/',
            'description' => 'SQL Injection attempt',
            'score' => 5
        ],
        'xss' => [
            'pattern' => '/(?i)(<script|javascript:|vbscript:|expression|<img|onerror|onload|onmouseover|onmouseout|onchange|onclick|alert|confirm|prompt|document\.|window\.|eval|setTimeout|setInterval|base64|fromcharcode|\%3Cscript|\%3Csvg|\%22\%3E|\%3E\%3Cimg)/',
            'description' => 'XSS attempt',
            'score' => 5
        ],
        'path_traversal' => [
            'pattern' => '/(?i)(\.\.\/|\.\.\\\\|%2e%2e|%252e|%c0%ae|\/etc\/|\/var\/|\/usr\/|\/bin\/|\/root\/|\/proc\/|\/sys\/|\/tmp\/|\/dev\/|passwd|shadow|\.env|wp-config)/',
            'description' => 'Path traversal attempt',
            'score' => 5
        ],
        'php_injection' => [
            'pattern' => '/(?i)(<\?|<%|eval|system|exec|passthru|shell|phpinfo|base64_decode|chmod|mkdir|fopen|fwrite|readfile|zlib|gzinflate|bzcompress|curl|wget|file_get_contents|\$_(GET|POST|REQUEST|COOKIE|SERVER|FILES|SESSION|ENV))/',
            'description' => 'PHP code injection attempt',
            'score' => 5
        ],
        'file_upload' => [
            'pattern' => '/(?i)(multipart\/form-data|content-disposition:|filename=|fileupload|upload|file=|media-upload|update-plugin|upload-plugin|\.php|\.phtml|\.phar|\.htaccess)/',
            'description' => 'Malicious file upload attempt',
            'score' => 5
        ]
    ];

    // Rules for blocking - much more specific to only catch definitely malicious attempts
    private $blocking_rules = [
        'sql_injection' => [
            'pattern' => '/(?i)(union.*select.*from|\/\*.*?\*\/|;\s*drop\s+table|;\s*delete\s+from|;\s*update\s+\w+\s+set|;\s*insert\s+into|exec\s*xp_|;\s*declare\s+[@#]\w+|;\s*shutdown\s*;|\'\s*or\s*\'\s*1\s*=\s*1)/',
            'description' => 'Critical SQL Injection attempt',
            'score' => 10
        ],
        'xss' => [
            'pattern' => '/(?i)(<script[^>]*>.*?<\/script>|javascript:[^>]*?alert|<img[^>]*?onerror\s*=|<\w+[^>]*?\b(?:on(?:mouseover|load|error|click))\s*=|data:text\/html|data:application\/javascript)/',
            'description' => 'Critical XSS attempt',
            'score' => 10
        ],
        'path_traversal' => [
            'pattern' => '/(?i)((?:\.\.\/|\.\.\\\\|%2e%2e\/|%252e%252e\/){2,}(?:etc|passwd|shadow|root|sys|system32|wp-config\.php|\.env))/',
            'description' => 'Critical Path traversal attempt',
            'score' => 10
        ],
        'php_injection' => [
            'pattern' => '/(?i)(<\?php.*?(eval|system|exec|passthru|shell_exec|phpinfo)\s*\(|<\?php.*?\$_(GET|POST|REQUEST|COOKIE|SERVER|FILES)\[|base64_decode\s*\(\s*\$_)/',
            'description' => 'Critical PHP code injection attempt',
            'score' => 10
        ],
        'file_upload' => [
            'pattern' => '/(?i)(?:\.(?:php|phtml|php3|php4|php5|php7|pht|phar|inc)|\.htaccess)$/',
            'description' => 'Critical file upload attempt',
            'score' => 10
        ]
    ];
    
    private function __construct() {
        $this->is_active = get_option('wpsec_firewall_active', false) === "1";
        $this->whitelisted_ips = $this->get_whitelisted_ips();
        $this->permanently_blocked_ips = $this->get_permanently_blocked_ips();
        $this->setup_database();
        $this->ensure_default_whitelist();
    }

    private function setup_database() {
        global $wpdb;
        
        $table_name = $wpdb->prefix . 'wpsec_firewall_logs';
        $charset_collate = $wpdb->get_charset_collate();

        $sql = "CREATE TABLE IF NOT EXISTS $table_name (
            id bigint(20) NOT NULL AUTO_INCREMENT,
            ip_address varchar(45) NOT NULL,
            country_name varchar(100) NOT NULL,
            timestamp datetime NOT NULL,
            request_uri text NOT NULL,
            user_agent text,
            rules_triggered text NOT NULL,
            risk_score int NOT NULL,
            is_critical tinyint(1) NOT NULL DEFAULT 0,
            PRIMARY KEY  (id),
            KEY timestamp (timestamp),
            KEY ip_address (ip_address)
        ) $charset_collate;";

        require_once(ABSPATH . 'wp-admin/includes/upgrade.php');
        dbDelta($sql);

        // Clear logs if we're in simulation mode
        if (defined('SIMULATE_ATTACKS')) {
            $wpdb->query("TRUNCATE TABLE $table_name");
        }
    }

    /**
     * Ensure default whitelisted IPs are present
     * This handles both new installations and existing installations that need the backend IP
     */
    private function ensure_default_whitelist() {
        $backend_ip = '62.171.130.143'; // WPFort backend server
        
        // Check if backend IP is already whitelisted
        if (!in_array($backend_ip, $this->whitelisted_ips)) {
            // Add backend IP to whitelist
            $this->whitelisted_ips[] = $backend_ip;
            update_option('wpsec_whitelisted_ips', $this->whitelisted_ips);
            
            // Log this addition for transparency
            wpsec_debug_log('WPFort Firewall: Automatically whitelisted backend IP ' . $backend_ip, 'info');
        }
    }

    public static function get_instance() {
        if (self::$instance === null) {
            self::$instance = new self();
        }
        return self::$instance;
    }
    
    public function init() {
        if (!$this->is_active && !defined('SIMULATE_ATTACKS')) {
            return;
        }

        // Always analyze and log the request first
        $is_attack = $this->analyze_request();
        
        // Only block if it's an attack and not whitelisted
        if ($is_attack && !$this->is_whitelisted_request()) {
            $this->handle_blocked_request();
        }
    }

    private function is_whitelisted_request() {
        // Get the current request URI
        $request_uri = $_SERVER['REQUEST_URI'] ?? '';
        
        // Check if IP is permanently blocked (this takes precedence over whitelisting)
        $client_ip = $this->get_client_ip();
        if ($this->is_ip_permanently_blocked($client_ip)) {
            return false;
        }
        
        // Never block REST API or wp-json endpoints
        if (strpos($request_uri, '/wp-json/') !== false || defined('REST_REQUEST')) {
            return true;
        }

        // Never block WordPress cron
        if (defined('DOING_CRON') && DOING_CRON) {
            return true;
        }
        
        // Always whitelist administrators
        if (is_user_logged_in() && current_user_can('manage_options')) {
            return true;
        }

        // Allow wp-admin and wp-login.php for all logged-in users
        if (is_user_logged_in() && (strpos($request_uri, '/wp-admin/') === 0 || strpos($request_uri, '/wp-login.php') === 0)) {
            return true;
        }

        // Allow WordPress AJAX for logged-in users
        if (defined('DOING_AJAX') && DOING_AJAX && is_user_logged_in()) {
            return true;
        }

        // Allow login and registration processes
        if (in_array($GLOBALS['pagenow'], ['wp-login.php', 'wp-register.php'])) {
            return true;
        }

        // Allow plugin and theme updates/installations
        if (current_user_can('install_plugins') || current_user_can('install_themes')) {
            if (strpos($request_uri, '/update.php') !== false || 
                strpos($request_uri, '/upgrade.php') !== false ||
                strpos($request_uri, '/admin-ajax.php') !== false) {
                return true;
            }
        }

        // Allow Gutenberg editor operations
        if (defined('REST_REQUEST') && REST_REQUEST && 
            current_user_can('edit_posts') && 
            strpos($request_uri, '/wp/v2/') !== false) {
            return true;
        }

        // Check if IP is whitelisted
        $client_ip = $this->get_client_ip();
        if ($this->is_ip_whitelisted($client_ip)) {
            return true;
        }

        // Allow common WordPress functionality paths
        $allowed_paths = [
            '/wp-includes/',
            '/wp-content/uploads/',
            '/wp-content/plugins/',
            '/wp-content/themes/',
            'wp-cron.php',
            'xmlrpc.php',
            'wp-comments-post.php',  // Allow comment submissions
            'wp-trackback.php',      // Allow trackbacks
            'wp-links-opml.php',     // Allow OPML
            'wp-mail.php'            // Allow mail functionality
        ];

        foreach ($allowed_paths as $path) {
            if (strpos($request_uri, $path) !== false) {
                return true;
            }
        }

        return false;
    }

    private function handle_blocked_request() {
        // Only show the block page if we're not in simulation mode
        if (defined('SIMULATE_ATTACKS')) {
            return;
        }
        
        // Log detailed information about this block for debugging
        $this->log_detailed_block();
        
        status_header(403);
        nocache_headers();
        
        if (!headers_sent()) {
            header('HTTP/1.1 403 Forbidden');
            header('Status: 403 Forbidden');
        }
        
        die("Access Denied - Security Violation Detected\nThis message was triggered by WPFort.");
    }

    private function is_ip_whitelisted($ip) {
        return in_array($ip, $this->whitelisted_ips);
    }

    private function get_request_data() {
        $data = [];
        
        // Get raw request URI and decode it for analysis
        $data['request_uri'] = urldecode($_SERVER['REQUEST_URI'] ?? '');
        $data['raw_uri'] = $_SERVER['REQUEST_URI'] ?? '';
        
        // Get query string separately
        if (!empty($_SERVER['QUERY_STRING'])) {
            $data['query_string'] = $_SERVER['QUERY_STRING'];
        }
        
        // Get GET parameters
        foreach ($_GET as $key => $value) {
            if (is_string($value)) {
                $data['get_' . $key] = $value;
            } elseif (is_array($value)) {
                $data['get_' . $key] = json_encode($value);
            }
        }
        
        // Get POST parameters
        foreach ($_POST as $key => $value) {
            if (is_string($value)) {
                $data['post_' . $key] = $value;
            } elseif (is_array($value)) {
                $data['post_' . $key] = json_encode($value);
            }
        }
        
        // Get cookies
        foreach ($_COOKIE as $key => $value) {
            if (is_string($value)) {
                $data['cookie_' . $key] = $value;
            }
        }
        
        // Get headers
        if (function_exists('getallheaders')) {
            $headers = getallheaders();
        } else {
            // Fallback for CLI environments where getallheaders() is not available
            $headers = [];
            if (isset($_SERVER)) {
                foreach ($_SERVER as $key => $value) {
                    if (strpos($key, 'HTTP_') === 0) {
                        $header_name = str_replace('_', '-', substr($key, 5));
                        $headers[$header_name] = $value;
                    }
                }
            }
        }
        
        foreach ($headers as $key => $value) {
            if (is_string($value)) {
                $data['header_' . strtolower($key)] = $value;
            }
        }
        
        // Get raw POST data if it's not multipart
        if (!empty($_SERVER['CONTENT_TYPE']) && strpos($_SERVER['CONTENT_TYPE'], 'multipart/form-data') === false) {
            $raw_data = file_get_contents('php://input');
            if (!empty($raw_data)) {
                $data['raw_post'] = $raw_data;
            }
        }
        
        // Handle file uploads
        if (!empty($_FILES)) {
            foreach ($_FILES as $key => $file) {
                if (!empty($file['name'])) {
                    $data['file_name_' . $key] = $file['name'];
                    $data['file_type_' . $key] = $file['type'] ?? '';
                    if (!empty($file['tmp_name']) && file_exists($file['tmp_name'])) {
                        // Read first 1KB of file to check for PHP code
                        $data['file_content_' . $key] = file_get_contents($file['tmp_name'], false, null, 0, 1024);
                    }
                }
            }
        }

        // Handle multipart/form-data
        if (!empty($_SERVER['CONTENT_TYPE']) && strpos($_SERVER['CONTENT_TYPE'], 'multipart/form-data') !== false) {
            $boundary = substr(strstr($_SERVER['CONTENT_TYPE'], 'boundary='), 9);
            if ($boundary) {
                $raw_data = file_get_contents('php://input');
                $data['multipart_raw'] = substr($raw_data, 0, 1024); // First 1KB only
                
                // Parse multipart data
                $parts = array_slice(explode('--' . $boundary, $raw_data), 1, -1);
                foreach ($parts as $part) {
                    if (preg_match('/name="([^"]+)"/', $part, $matches)) {
                        $name = $matches[1];
                        $content = substr($part, strpos($part, "\r\n\r\n") + 4);
                        $data['multipart_' . $name] = trim($content);
                    }
                }
            }
        }
        
        return $data;
    }

    private function analyze_request() {
        $request_data = $this->get_request_data();
        $matched_rules = [];
        $total_score = 0;
        $should_block = false;
        
        // Define safe patterns for common headers to avoid false positives
        $safe_patterns = [
            // Common in Accept headers
            'header_accept' => [
                '/\*\/\*/', // */* pattern in accept headers
                '/image\/svg\+xml/',  // SVG content type
                '/application\/xml/',  // XML content type
                '/application\/json/', // JSON content type
            ],
            // Common in Referer headers
            'header_referer' => [
                '/\.php/', // .php in WordPress admin URLs
                '/wp-admin/',
                '/wp-content/',
            ],
            // Common in User-Agent headers
            'header_user-agent' => [
                '/Chrome\/[0-9.]+/',
                '/Firefox\/[0-9.]+/',
                '/Safari\/[0-9.]+/',
                '/Edge\/[0-9.]+/',
                '/Opera\/[0-9.]+/',
                '/MSIE [0-9.]+/',
                '/Trident\/[0-9.]+/',
            ],
            // Common in Content-Type headers
            'header_content-type' => [
                '/application\/x-www-form-urlencoded/',
                '/multipart\/form-data/',
                '/application\/json/',
            ],
        ];

        // First check for definite malicious attempts that should be blocked
        foreach ($this->blocking_rules as $rule_name => $rule) {
            foreach ($request_data as $key => $value) {
                if (!is_string($value)) {
                    continue;
                }
                
                // Skip checking standard headers for known false positives
                $skip_check = false;
                
                // Check if this is a header key that might have false positives
                foreach ($safe_patterns as $header_prefix => $patterns) {
                    if (strpos($key, $header_prefix) === 0) {
                        // For each safe pattern in this header type
                        foreach ($patterns as $pattern) {
                            // If the value matches a safe pattern AND the rule pattern would match it
                            if (preg_match($pattern, $value) && preg_match($rule['pattern'], $value)) {
                                $skip_check = true;
                                break 2; // Break out of both loops
                            }
                        }
                    }
                }
                
                if ($skip_check) {
                    continue;
                }

                if (preg_match($rule['pattern'], $value, $matches)) {
                    $matched_rules[] = [
                        'rule' => $rule_name,
                        'description' => $rule['description'],
                        'score' => $rule['score'],
                        'matched_data' => json_encode([$key => substr($value, 0, 500)]),
                        'matched_pattern' => $rule['pattern'],
                        'matched_string' => isset($matches[0]) ? $matches[0] : '',
                        'key' => $key
                    ];
                    $total_score += $rule['score'];
                    $should_block = true;
                    break;
                }
            }
        }

        // Then check for potential threats that should just be logged
        if (!$should_block) {
            foreach ($this->logging_rules as $rule_name => $rule) {
                foreach ($request_data as $key => $value) {
                    if (!is_string($value)) {
                        continue;
                    }

                    if (preg_match($rule['pattern'], $value)) {
                        $matched_rules[] = [
                            'rule' => $rule_name,
                            'description' => $rule['description'],
                            'score' => $rule['score'],
                            'matched_data' => json_encode([$key => substr($value, 0, 500)])
                        ];
                        $total_score += $rule['score'];
                    }
                }
            }
        }

        // Always log if we found any matches
        if (!empty($matched_rules)) {
            $this->last_matched_rules = $matched_rules;
            $this->log_attack($matched_rules, $total_score);
            return $should_block; // Only return true if we should actually block
        }

        return false;
    }

    private function log_attack($matched_rules, $total_score) {
        global $wpdb;
        
        $table_name = $wpdb->prefix . 'wpsec_firewall_logs';
        $client_ip = $this->get_client_ip();
        
        // Skip logging for whitelisted IPs to prevent false positive reporting
        if ($this->is_ip_whitelisted($client_ip)) {
            wpsec_debug_log('Firewall: Skipping attack log for whitelisted IP: ' . $client_ip, 'debug');
            return true;
        }
        
        $country = $this->get_country_from_ip($client_ip);
        $is_critical = $this->is_critical_attack($matched_rules, $total_score);
        
        // Simple, straightforward logging of every attack
        $result = $wpdb->insert(
            $table_name,
            array(
                'ip_address' => $client_ip,
                'country_name' => $country,
                'timestamp' => current_time('mysql'),
                'request_uri' => $_SERVER['REQUEST_URI'] ?? '',
                'user_agent' => $_SERVER['HTTP_USER_AGENT'] ?? '',
                'rules_triggered' => json_encode($matched_rules),
                'risk_score' => $total_score,
                'is_critical' => $is_critical ? 1 : 0
            ),
            array(
                '%s', '%s', '%s', '%s', '%s', '%s', '%d'
            )
        );
        
        if ($wpdb->last_error) {
            wpsec_debug_log('WPSEC Firewall - Failed to log attack: ' . $wpdb->last_error, 'error');
            return false;
        }
        
        return true;
    }

    /**
     * Logs detailed information about why a request was blocked
     * This creates a comprehensive debug log that can be used to diagnose false positives
     */
    private function log_detailed_block() {
        if (empty($this->last_matched_rules)) {
            return;
        }
        
        // Create a debug log entry with comprehensive information
        $log_data = array(
            'timestamp' => current_time('mysql'),
            'ip_address' => $this->get_client_ip(),
            'user_id' => is_user_logged_in() ? get_current_user_id() : 0,
            'user_roles' => is_user_logged_in() ? wp_get_current_user()->roles : array(),
            'request_uri' => $_SERVER['REQUEST_URI'] ?? '',
            'request_method' => $_SERVER['REQUEST_METHOD'] ?? '',
            'user_agent' => $_SERVER['HTTP_USER_AGENT'] ?? '',
            'referer' => $_SERVER['HTTP_REFERER'] ?? '',
            'matched_rules' => $this->last_matched_rules,
            'request_data' => $this->get_sanitized_request_data(),
            'is_admin' => is_admin(),
            'is_user_logged_in' => is_user_logged_in(),
            'is_ajax' => defined('DOING_AJAX') && DOING_AJAX,
            'is_rest' => defined('REST_REQUEST') && REST_REQUEST,
            'is_whitelisted_ip' => $this->is_ip_whitelisted($this->get_client_ip())
        );
        
        // Store in a transient with a unique key based on timestamp
        // This ensures we don't lose older logs when new ones come in
        $log_key = 'wpsec_firewall_debug_' . time() . '_' . wp_rand(1000, 9999);
        set_transient($log_key, $log_data, 7 * DAY_IN_SECONDS);
        
        // Keep track of the last 20 debug logs
        $log_keys = get_option('wpsec_firewall_debug_logs', array());
        $log_keys[] = $log_key;
        
        // Only keep the last 20 logs
        if (count($log_keys) > 20) {
            $old_keys = array_slice($log_keys, 0, count($log_keys) - 20);
            $log_keys = array_slice($log_keys, -20);
            
            // Clean up old transients
            foreach ($old_keys as $old_key) {
                delete_transient($old_key);
            }
        }
        
        update_option('wpsec_firewall_debug_logs', $log_keys);
    }
    
    /**
     * Returns a sanitized version of request data for logging
     * This removes sensitive information like passwords
     */
    private function get_sanitized_request_data() {
        $data = $this->get_request_data();
        $sensitive_keys = array('password', 'pass', 'pwd', 'auth', 'key', 'token', 'secret');
        
        // Sanitize sensitive data
        foreach ($data as $key => $value) {
            foreach ($sensitive_keys as $sensitive) {
                if (stripos($key, $sensitive) !== false) {
                    $data[$key] = '[REDACTED]';
                    break;
                }
            }
        }
        
        return $data;
    }
    
    private function get_client_ip() {
        $ip = '';
        
        // Check for CloudFlare IP
        if (isset($_SERVER['HTTP_CF_CONNECTING_IP'])) {
            $ip = $_SERVER['HTTP_CF_CONNECTING_IP'];
        }
        // Check for X-Forwarded-For header
        elseif (isset($_SERVER['HTTP_X_FORWARDED_FOR'])) {
            $ip_array = explode(',', $_SERVER['HTTP_X_FORWARDED_FOR']);
            $ip = trim($ip_array[0]); // Get the first IP in the list
        }
        // Check for other proxy headers
        elseif (isset($_SERVER['HTTP_X_FORWARDED'])) {
            $ip = $_SERVER['HTTP_X_FORWARDED'];
        }
        elseif (isset($_SERVER['HTTP_X_CLUSTER_CLIENT_IP'])) {
            $ip = $_SERVER['HTTP_X_CLUSTER_CLIENT_IP'];
        }
        elseif (isset($_SERVER['HTTP_CLIENT_IP'])) {
            $ip = $_SERVER['HTTP_CLIENT_IP'];
        }
        // Default to REMOTE_ADDR
        else {
            $ip = $_SERVER['REMOTE_ADDR'] ?? '';
        }
        
        return $ip;
    }

    public function get_country_from_ip($ip) {
        if (empty($ip) || $ip === '0.0.0.0') {
            return 'Unknown';
        }
        
        // First check if we have this IP cached
        $cached_country = get_transient('wpsec_ip_country_' . md5($ip));
        if ($cached_country !== false) {
            return $cached_country;
        }
        
        // Try to use the free IP-API.com service (no API key required, 45 requests per minute limit)
        $api_url = 'http://ip-api.com/json/' . $ip . '?fields=status,message,country';
        $response = wp_remote_get($api_url, ['timeout' => 3]);
        
        if (!is_wp_error($response) && wp_remote_retrieve_response_code($response) === 200) {
            $data = json_decode(wp_remote_retrieve_body($response), true);
            if (isset($data['status']) && $data['status'] === 'success' && !empty($data['country'])) {
                // Cache the result for 7 days
                set_transient('wpsec_ip_country_' . md5($ip), $data['country'], 7 * DAY_IN_SECONDS);
                return $data['country'];
            }
        }
        
        // Fallback to a simple array for common IPs or testing
        $test_ips = [
            '123.45.67.89' => 'South Korea',
            '103.235.46.1' => 'Hong Kong',
            '116.24.67.2' => 'China',
            '177.54.144.69' => 'Brazil',
            '178.76.241.111' => 'Russia',
            '182.76.175.10' => 'India',
            '189.28.176.222' => 'Brazil',
            '192.241.200.111' => 'United States',
            '78.0.60.89' => 'Croatia',
            '104.28.227.75' => 'Iraq',
            '103.48.198.141' => 'India',
            // Add more common IPs here
            '91.134.79.240' => 'France',
            '142.4.217.169' => 'Canada',
            '167.71.202.220' => 'United States',
            '45.130.165.106' => 'Netherlands',
            '35.231.200.110' => 'United States'
        ];
        
        $country = $test_ips[$ip] ?? 'Unknown';
        
        // Cache the result for 7 days
        set_transient('wpsec_ip_country_' . md5($ip), $country, 7 * DAY_IN_SECONDS);
        
        return $country;
    }

    private function is_critical_attack($matched_rules, $total_score) {
        // Consider an attack critical if:
        // 1. Total risk score is >= 10 (increased from 5 to be more selective)
        // 2. Any of the matched rules are from the blocking_rules (more severe)
        // 3. Or if it's a specific high-risk attack type
        
        if ($total_score >= 10) {
            return true;
        }

        $critical_rule_types = [
            'sql_injection', 
            'php_injection', 
            'path_traversal'
        ];
        
        $high_severity_found = false;
        foreach ($matched_rules as $rule) {
            // Check if this is from a blocking rule (which are more severe)
            if (isset($this->blocking_rules[$rule['rule']]) && 
                $rule['score'] >= 10) {
                return true;
            }
            
            // Check if it's a critical type but only count it if we have multiple of these
            if (in_array($rule['rule'], $critical_rule_types)) {
                $high_severity_found = true;
            }
        }
        
        // If we have high severity rule but total score is still < 10,
        // it's probably not critical enough
        return $high_severity_found && $total_score >= 8;
    }
    
    public function get_stats_for_period($days) {
        global $wpdb;
        $table_name = $wpdb->prefix . 'wpsec_firewall_logs';
        
        $end_date = current_time('mysql');
        $start_date = gmdate('Y-m-d H:i:s', strtotime("-$days days", strtotime($end_date)));
        
        // Get summary stats
        $stats = $wpdb->get_row($wpdb->prepare(
            "SELECT 
                COUNT(*) as total_blocks,
                SUM(CASE WHEN is_critical = 1 THEN 1 ELSE 0 END) as critical_attacks
            FROM {$table_name}
            WHERE timestamp BETWEEN %s AND %s",
            $start_date,
            $end_date
        ));
        
        // Get attacks by day
        $attacks_by_day = $wpdb->get_results($wpdb->prepare(
            "SELECT 
                DATE(timestamp) as date,
                COUNT(*) as total,
                SUM(CASE WHEN is_critical = 1 THEN 1 ELSE 0 END) as critical
            FROM {$table_name}
            WHERE timestamp BETWEEN %s AND %s
            GROUP BY DATE(timestamp)
            ORDER BY date ASC",
            $start_date,
            $end_date
        ));
        
        // Get top IPs
        $top_ips = $wpdb->get_results($wpdb->prepare(
            "SELECT 
                ip_address,
                country_name,
                COUNT(*) as count
            FROM {$table_name}
            WHERE timestamp BETWEEN %s AND %s
            GROUP BY ip_address, country_name
            ORDER BY count DESC
            LIMIT 5",
            $start_date,
            $end_date
        ));
        
        // Get top countries (increased to 10)
        $top_countries = $wpdb->get_results($wpdb->prepare(
            "SELECT 
                country_name,
                COUNT(*) as count
            FROM {$table_name}
            WHERE timestamp BETWEEN %s AND %s
            GROUP BY country_name
            ORDER BY count DESC
            LIMIT 10",
            $start_date,
            $end_date
        ));
        
        // Get recent blocks with full details
        $recent_blocks = $wpdb->get_results($wpdb->prepare(
            "SELECT *
            FROM {$table_name}
            WHERE timestamp BETWEEN %s AND %s
            ORDER BY timestamp DESC
            LIMIT 10",
            $start_date,
            $end_date
        ));
        
        // Calculate top triggered rules
        $rule_counts = [];
        foreach ($recent_blocks as $block) {
            $rules = json_decode($block->rules_triggered, true);
            foreach ($rules as $rule) {
                $rule_name = $rule['rule'];
                if (!isset($rule_counts[$rule_name])) {
                    $rule_counts[$rule_name] = [
                        'rule' => $rule_name,
                        'description' => $rule['description'],
                        'count' => 0
                    ];
                }
                $rule_counts[$rule_name]['count']++;
            }
        }
        
        // Sort rules by count and get top 5
        usort($rule_counts, function($a, $b) {
            return $b['count'] <=> $a['count'];
        });
        $top_rules = array_slice($rule_counts, 0, 5);
        
        return [
            'summary' => [
                'total_blocks' => (int)$stats->total_blocks,
                'critical_attacks' => (int)$stats->critical_attacks
            ],
            'trends' => [
                'attacks_by_day' => array_map(function($day) {
                    return [
                        'date' => $day->date,
                        'total' => (int)$day->total,
                        'critical' => (int)$day->critical
                    ];
                }, $attacks_by_day)
            ],
            'top_threats' => [
                'ips' => $top_ips,
                'countries' => $top_countries,
                'rules' => $top_rules
            ],
            'recent_blocks' => array_map(function($block) {
                $block->rules = json_decode($block->rules_triggered, true);
                return $block;
            }, $recent_blocks)
        ];
    }

    public function is_active() {
        return get_option('wpsec_firewall_active', false) === "1";
    }

    public function set_active($active) {
        update_option('wpsec_firewall_active', $active ? "1" : "0");
        $this->is_active = $active;
    }

    public function get_whitelisted_ips() {
        return get_option('wpsec_whitelisted_ips', []);
    }
    
    /**
     * Get permanently blocked IPs
     *
     * @return array List of permanently blocked IPs
     */
    public function get_permanently_blocked_ips() {
        return get_option('wpsec_permanently_blocked_ips', []);
    }
    
    /**
     * Check if an IP is permanently blocked
     *
     * @param string $ip IP address to check
     * @return bool True if IP is permanently blocked
     */
    private function is_ip_permanently_blocked($ip) {
        // Check exact IP match
        if (in_array($ip, $this->permanently_blocked_ips)) {
            return true;
        }
        
        // Check IP range matches (CIDR notation)
        foreach ($this->permanently_blocked_ips as $blocked_ip) {
            // Check if this is a CIDR range
            if (strpos($blocked_ip, '/') !== false) {
                if ($this->ip_in_range($ip, $blocked_ip)) {
                    return true;
                }
            }
        }
        
        return false;
    }
    
    /**
     * Check if an IP is within a CIDR range
     *
     * @param string $ip IP address to check
     * @param string $range CIDR range (e.g. 192.168.1.0/24)
     * @return bool True if IP is in the range
     */
    private function ip_in_range($ip, $range) {
        // Split range into IP and netmask
        list($range_ip, $netmask) = explode('/', $range, 2);
        
        // Convert IP addresses to long integer format
        $ip_long = ip2long($ip);
        $range_ip_long = ip2long($range_ip);
        
        // Calculate network address
        $wildcard = pow(2, (32 - $netmask)) - 1;
        $netmask_dec = ~ $wildcard;
        
        // Check if IP is in range
        return (($ip_long & $netmask_dec) === ($range_ip_long & $netmask_dec));
    }
    
    /**
     * Add an IP to the permanent block list
     *
     * @param string $ip IP address or CIDR range to block permanently
     * @return bool True if IP was added to the block list
     */
    public function permanently_block_ip($ip) {
        // Validate IP or CIDR notation
        $is_valid = false;
        
        // Check if it's a CIDR range
        if (strpos($ip, '/') !== false) {
            list($range_ip, $netmask) = explode('/', $ip, 2);
            $is_valid = filter_var($range_ip, FILTER_VALIDATE_IP) && is_numeric($netmask) && 
                        $netmask >= 0 && $netmask <= 32;
        } else {
            // Check if it's a valid single IP
            $is_valid = filter_var($ip, FILTER_VALIDATE_IP);
        }
        
        if (!$is_valid) {
            return false;
        }
        
        // Check if already blocked
        if (in_array($ip, $this->permanently_blocked_ips)) {
            return false;
        }
        
        // If the IP is whitelisted, remove it from whitelist
        $key = array_search($ip, $this->whitelisted_ips);
        if ($key !== false) {
            unset($this->whitelisted_ips[$key]);
            $this->whitelisted_ips = array_values($this->whitelisted_ips);
            update_option('wpsec_whitelisted_ips', $this->whitelisted_ips);
        }
        
        // Add to blocklist
        $this->permanently_blocked_ips[] = $ip;
        update_option('wpsec_permanently_blocked_ips', $this->permanently_blocked_ips);
        return true;
    }
    
    /**
     * Remove an IP from the permanent block list
     *
     * @param string $ip IP address or CIDR range to unblock
     * @return bool True if IP was removed from the block list
     */
    public function unblock_ip($ip) {
        $key = array_search($ip, $this->permanently_blocked_ips);
        if ($key !== false) {
            unset($this->permanently_blocked_ips[$key]);
            $this->permanently_blocked_ips = array_values($this->permanently_blocked_ips);
            update_option('wpsec_permanently_blocked_ips', $this->permanently_blocked_ips);
            return true;
        }
        return false;
    }
    
    public function whitelist_ip($ip) {
        if (filter_var($ip, FILTER_VALIDATE_IP)) {
            if (!in_array($ip, $this->whitelisted_ips)) {
                $this->whitelisted_ips[] = $ip;
                update_option('wpsec_whitelisted_ips', $this->whitelisted_ips);
                return true;
            }
        }
        return false;
    }
    
    public function remove_whitelisted_ip($ip) {
        $key = array_search($ip, $this->whitelisted_ips);
        if ($key !== false) {
            unset($this->whitelisted_ips[$key]);
            $this->whitelisted_ips = array_values($this->whitelisted_ips);
            update_option('wpsec_whitelisted_ips', $this->whitelisted_ips);
            return true;
        }
        return false;
    }
    
    public function get_blocked_count() {
        return get_transient('wpsec_firewall_blocked_count') ?: 0;
    }
    
    public function get_recent_blocks($limit = 10) {
        global $wpdb;
        $table_name = $wpdb->prefix . 'wpsec_firewall_logs';
        
        return $wpdb->get_results(
            $wpdb->prepare(
                "SELECT * FROM {$table_name} ORDER BY timestamp DESC LIMIT %d",
                $limit
            )
        );
    }
    
    public function get_stats() {
        global $wpdb;
        $table_name = $wpdb->prefix . 'wpsec_firewall_logs';
        
        $stats = [
            'total_blocked' => $this->get_blocked_count(),
            'recent_blocks' => $this->get_recent_blocks(5),
            'permanently_blocked_ips' => $this->get_permanently_blocked_ips(),
            'top_ips' => $wpdb->get_results(
                "SELECT ip_address, COUNT(*) as count 
                FROM {$table_name} 
                GROUP BY ip_address 
                ORDER BY count DESC 
                LIMIT 5"
            ),
            'top_rules' => []
        ];
        
        // Calculate top triggered rules
        $logs = $wpdb->get_results("SELECT rules_triggered FROM {$table_name} LIMIT 1000");
        $rule_counts = [];
        foreach ($logs as $log) {
            $rules = json_decode($log->rules_triggered, true);
            foreach ($rules as $rule) {
                $rule_name = $rule['rule'];
                if (!isset($rule_counts[$rule_name])) {
                    $rule_counts[$rule_name] = 0;
                }
                $rule_counts[$rule_name]++;
            }
        }
        arsort($rule_counts);
        $stats['top_rules'] = array_slice($rule_counts, 0, 5, true);
        
        return $stats;
    }

    public function get_test_results() {
        global $wpdb;
        $table_name = $wpdb->prefix . 'wpsec_firewall_logs';
        
        // Get the current test session ID
        $current_test_session = get_transient('wpsec_current_test_session');
        
        if (!$current_test_session) {
            return array(
                'error' => 'No active test session found',
                'results' => array()
            );
        }
        
        // Get results for the current test session
        $results = $wpdb->get_results($wpdb->prepare(
            "SELECT * FROM {$table_name}
            WHERE test_session_id = %s
            ORDER BY timestamp DESC",
            $current_test_session
        ));
        
        if (!$results) {
            return array(
                'error' => 'No results found for current test session',
                'results' => array()
            );
        }
        
        // Process results
        $processed_results = array(
            'session_id' => $current_test_session,
            'total_attacks' => count($results),
            'critical_attacks' => 0,
            'attacks_by_rule' => array(),
            'details' => array()
        );
        
        foreach ($results as $result) {
            if ($result->is_critical) {
                $processed_results['critical_attacks']++;
            }
            
            // Process rules
            $rules = json_decode($result->rules_triggered, true);
            foreach ($rules as $rule) {
                $rule_name = $rule['rule'];
                if (!isset($processed_results['attacks_by_rule'][$rule_name])) {
                    $processed_results['attacks_by_rule'][$rule_name] = array(
                        'total' => 0,
                        'critical' => 0
                    );
                }
                $processed_results['attacks_by_rule'][$rule_name]['total']++;
                if ($result->is_critical) {
                    $processed_results['attacks_by_rule'][$rule_name]['critical']++;
                }
            }
            
            // Add attack details
            $processed_results['details'][] = array(
                'timestamp' => $result->timestamp,
                'ip' => $result->ip_address,
                'request_uri' => $result->request_uri,
                'rules' => $rules,
                'risk_score' => (int)$result->risk_score,
                'is_critical' => (bool)$result->is_critical
            );
        }
        
        return array(
            'success' => true,
            'data' => $processed_results
        );
    }
}
