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

/**
 * ITK WAF — Web Application Firewall
 *
 * Loads rules from config/waf-rules.conf and inspects incoming requests
 * (GET params, REQUEST_URI, optionally POST, cookies, and User-Agent) for
 * known attack signatures.  On a match it logs, notifies the attacks API,
 * and either blocks or logs-only depending on the configured action.
 *
 * Options key: itk_waf (array)
 *   enabled         bool   Master switch.
 *   action          string 'block' (default) or 'log_only'.
 *   scan_post       bool   Also inspect $_POST values.
 *   scan_cookies    bool   Also inspect $_COOKIE values.
 *   scan_ua         bool   Also inspect HTTP_USER_AGENT.
 *   log_attacks     bool   Persist matches to the DB via ITK_Database::log_attack().
 *   response_code   string HTTP status code to send when blocking (e.g. '403', '301_custom').
 *   redirect_url    string Destination URL when response_code === '301_custom'.
 *   custom_message  string Body text sent with non-redirect block responses.
 *   block_{cat}     bool   Per-category enable flag (e.g. block_sqli, block_xss, …).
 */
class ITK_WAF {

    /** Transient key for the cached parsed ruleset. */
    const TRANSIENT_KEY = 'itk_waf_rules';

    /** Transient lifetime in seconds. */
    const CACHE_TTL = 300;

    /** Absolute path to the rules file. */
    private string $rules_file;

    /* ── Bootstrap ────────────────────────────────────────────── */

    public function __construct() {
        $options = get_option('itk_waf', []);
        if (empty($options['enabled'])) return;

        $this->rules_file = ITK_PATH . 'config/waf-rules.conf';

        add_action('init', [$this, 'inspect'], 0);
    }

    /* ── Main inspection hook ─────────────────────────────────── */

    /**
     * Runs at priority 0 on 'init'.  Inspects the current request against
     * all enabled WAF rules and calls handle_match() on the first hit.
     */
    public function inspect(): void {
        // Never block logged-in administrators or wp-admin non-AJAX requests.
        if (is_admin() && !wp_doing_ajax()) return;
        if (function_exists('current_user_can') && current_user_can('manage_options')) return;

        $options = get_option('itk_waf', []);
        $rules   = $this->load_rules();
        if (empty($rules)) return;

        // ── Build input map ────────────────────────────────────
        $inputs = [];

        // Always scan GET params.
        foreach ($_GET as $key => $value) {
            $inputs['GET'][(string)$key] = (string)$value;
        }

        // Always scan REQUEST_URI.
        $inputs['URI']['REQUEST_URI'] = (string)($_SERVER['REQUEST_URI'] ?? '');

        // Optionally scan POST.
        if (!empty($options['scan_post'])) {
            foreach ($_POST as $key => $value) {
                $inputs['POST'][(string)$key] = (string)$value;
            }
        }

        // Optionally scan cookies.
        if (!empty($options['scan_cookies'])) {
            foreach ($_COOKIE as $key => $value) {
                $inputs['COOKIE'][(string)$key] = (string)$value;
            }
        }

        // Optionally scan User-Agent.
        if (!empty($options['scan_ua'])) {
            $ua = (string)($_SERVER['HTTP_USER_AGENT'] ?? '');
            if ($ua !== '') {
                $inputs['UA']['HTTP_USER_AGENT'] = $ua;
            }
        }

        // ── Match loop ─────────────────────────────────────────
        foreach ($inputs as $source => $params) {
            foreach ($params as $key => $raw_value) {
                // Decode URL-encoding so encoded payloads are matched.
                $value = rawurldecode(urldecode($raw_value));

                foreach ($rules as $rule) {
                    // Skip categories disabled in options.
                    $opt_key = 'block_' . $rule['category'];
                    if (empty($options[$opt_key])) continue;

                    if (@preg_match($rule['pattern'], $value)) {
                        $this->handle_match($rule, $source, (string)$key, $raw_value);
                        return; // Stop on first match; handle_match() may exit().
                    }
                }
            }
        }
    }

    /* ── Rule loading ─────────────────────────────────────────── */

    /**
     * Reads and parses config/waf-rules.conf.
     * Results are cached in a transient for CACHE_TTL seconds.
     *
     * @return array<int, array{category: string, pattern: string, desc: string}>
     */
    public function load_rules(): array {
        $cached = get_transient(self::TRANSIENT_KEY);
        if (is_array($cached)) return $cached;

        $rules = [];

        if (!file_exists($this->rules_file)) {
            set_transient(self::TRANSIENT_KEY, $rules, self::CACHE_TTL);
            return $rules;
        }

        $lines = file($this->rules_file, FILE_IGNORE_NEW_LINES | FILE_SKIP_EMPTY_LINES);
        foreach ($lines as $line) {
            $line = trim($line);

            // Skip comment and header lines.
            if ($line === '' || $line[0] === '#') continue;

            $parts = explode('|', $line, 3);
            if (count($parts) < 3) continue;

            [$category, $pattern, $desc] = $parts;
            $category = trim($category);
            $pattern  = trim($pattern);
            $desc     = trim($desc);

            if ($category === '' || $pattern === '') continue;

            // Wrap bare patterns (no delimiter) with # delimiters so they are
            // valid PCRE.  Patterns that already start with a delimiter are
            // used as-is.
            $delimiters = ['/', '#', '~', '!', '@', '%'];
            if (!in_array($pattern[0], $delimiters, true)) {
                $pattern = '#' . $pattern . '#';
            }

            // Validate the regex before adding it to the ruleset.
            if (@preg_match($pattern, '') === false) continue;

            $rules[] = [
                'category' => $category,
                'pattern'  => $pattern,
                'desc'     => $desc,
            ];
        }

        set_transient(self::TRANSIENT_KEY, $rules, self::CACHE_TTL);
        return $rules;
    }

    /* ── Match handler ────────────────────────────────────────── */

    /**
     * Handles a rule match: logs the event, queues it for the attacks API,
     * then either returns (log-only mode) or terminates the request.
     *
     * @param array  $rule   Parsed rule: category, pattern, desc.
     * @param string $source Input source label: GET, POST, COOKIE, URI, UA.
     * @param string $key    Parameter name that triggered the rule.
     * @param string $value  Raw parameter value.
     */
    public function handle_match(array $rule, string $source, string $key, string $value): void {
        $options = get_option('itk_waf', []);

        $event = [
            'ip'          => $this->get_ip(),
            'attack_type' => $rule['category'],
            'rule_desc'   => $rule['desc'],
            'source'      => $source,
            'param'       => $key,
            'payload'     => substr($value, 0, 500),
            'uri'         => (string)($_SERVER['REQUEST_URI']  ?? ''),
            'method'      => (string)($_SERVER['REQUEST_METHOD'] ?? ''),
            'ua'          => (string)($_SERVER['HTTP_USER_AGENT'] ?? ''),
        ];

        // Persist to local DB if logging is enabled.
        if (!empty($options['log_attacks']) && class_exists('ITK_Database') &&
            method_exists('ITK_Database', 'log_attack')) {
            ITK_Database::log_attack($event);
        }

        // Queue for the attacks API.
        if (class_exists('ITK_Attacks_API') && method_exists('ITK_Attacks_API', 'queue')) {
            ITK_Attacks_API::queue($event);
        }

        // Log-only mode: record the event but do not block the request.
        $action = $options['action'] ?? 'block';
        if ($action === 'log_only') return;

        // ── Block the request ──────────────────────────────────
        $response_code  = $options['response_code']  ?? '403';
        $redirect_url   = $options['redirect_url']   ?? '';
        $custom_message = $options['custom_message'] ?? 'Access denied.';

        if ($response_code === '301_custom' && !empty($redirect_url)) {
            header('Location: ' . esc_url_raw($redirect_url), true, 301);
        } else {
            status_header((int)$response_code ?: 403);
            echo esc_html($custom_message);
        }
        exit;
    }

    /* ── IP resolution ────────────────────────────────────────── */

    /**
     * Returns the best available client IP address, checking proxy headers
     * in order of trustworthiness.  Mirrors the pattern used by other ITK
     * classes.
     */
    public function get_ip(): string {
        $keys = [
            'HTTP_CLIENT_IP',
            'HTTP_X_FORWARDED_FOR',
            'HTTP_X_FORWARDED',
            'HTTP_X_CLUSTER_CLIENT_IP',
            'HTTP_FORWARDED_FOR',
            'HTTP_FORWARDED',
            'REMOTE_ADDR',
        ];

        foreach ($keys as $key) {
            if (empty($_SERVER[$key])) continue;
            // X-Forwarded-For may contain a comma-separated list; take first.
            $ip = trim(explode(',', $_SERVER[$key])[0]);
            if (filter_var($ip, FILTER_VALIDATE_IP)) return $ip;
        }

        return 'UNKNOWN';
    }

    /* ── Cache management ─────────────────────────────────────── */

    /**
     * Deletes the cached ruleset transient so the next request re-parses
     * the rules file.  Call this after saving updated rules.
     */
    public function invalidate_cache(): void {
        delete_transient(self::TRANSIENT_KEY);
    }
}