"""
core/waf_detect.py — ReconNinja v7.0.0
Web Application Firewall (WAF) detection.

Checks response headers, cookies, and body for WAF fingerprints.
Optionally uses wafw00f if installed.

Knowing a WAF is present changes how you interpret dir scan and
nuclei results (false negatives are common behind WAFs).
"""

from __future__ import annotations

import urllib.request
import urllib.error
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

from utils.helpers import run_cmd, tool_exists, ensure_dir
from utils.logger import safe_print, log


# ── WAF fingerprint database ──────────────────────────────────────────────────

WAF_SIGNATURES: list[tuple[str, list[tuple[str, str]]]] = [
    # (WAF name, [(location, match_string)])
    # location: "header:<name>" | "cookie:<name>" | "body"
    ("Cloudflare",          [("header:server", "cloudflare"),
                              ("header:cf-ray", ""),
                              ("cookie:__cfduid", "")]),
    ("AWS WAF",             [("header:x-amzn-requestid", ""),
                              ("body", "aws-waf")]),
    ("Akamai Kona",         [("header:x-check-cacheable", ""),
                              ("body", "akamai")]),
    ("Imperva Incapsula",   [("header:x-iinfo", ""),
                              ("cookie:incap_ses", ""),
                              ("cookie:visid_incap", "")]),
    ("F5 BIG-IP ASM",       [("cookie:ts", ""),
                              ("header:x-waf-status", ""),
                              ("body", "request rejected by big-ip")]),
    ("ModSecurity",         [("header:server", "mod_security"),
                              ("body", "modsecurity"),
                              ("body", "mod_security")]),
    ("Sucuri",              [("header:x-sucuri-id", ""),
                              ("header:x-sucuri-cache", "")]),
    ("Barracuda",           [("cookie:barra_counter_session", ""),
                              ("header:x-barracuda-connect", "")]),
    ("Wordfence",           [("body", "wordfence"),
                              ("body", "generated by wordfence")]),
    ("DenyAll",             [("cookie:sessioncookie", ""),
                              ("body", "denyall")]),
    ("Fortinet FortiWeb",   [("header:set-cookie", "cookiesession1="),
                              ("body", "fortigate")]),
    ("Radware AppWall",     [("header:x-rdwr-pop", ""),
                              ("body", "rdwr")]),
    ("Sophos UTM",          [("body", "sophos web appliance")]),
    ("Netsparker",          [("header:x-protected-by", "netsparker")]),
    ("Varnish",             [("header:x-varnish", ""),
                              ("header:via", "varnish")]),
    ("nginx",               [("header:server", "nginx")]),    # not a WAF but useful
    ("Apache",              [("header:server", "apache")]),   # not a WAF but useful
]

# Payload that WAFs often block — tests active blocking
WAF_TEST_PAYLOAD = "/?test=<script>alert(1)</script>&id=1' OR 1=1--"


@dataclass
class WAFResult:
    url:          str
    waf_detected: bool
    waf_names:    list[str] = field(default_factory=list)
    headers:      dict      = field(default_factory=dict)
    blocked:      bool      = False   # returned 403/406 on test payload
    tool:         str       = "passive"   # passive | wafw00f

    def to_dict(self) -> dict:
        return {
            "url":          self.url,
            "waf_detected": self.waf_detected,
            "waf_names":    self.waf_names,
            "blocked":      self.blocked,
            "tool":         self.tool,
        }


# ── Passive header/cookie/body scan ──────────────────────────────────────────

def _passive_detect(url: str, timeout: int = 10) -> WAFResult:
    result = WAFResult(url=url, waf_detected=False)

    headers_map: dict[str, str] = {}
    cookies:     dict[str, str] = {}
    body = ""

    try:
        req = urllib.request.Request(
            url,
            headers={"User-Agent": "Mozilla/5.0 (ReconNinja/7.0.0)"},
        )
        with urllib.request.urlopen(req, timeout=timeout) as r:
            for k, v in r.headers.items():
                headers_map[k.lower()] = v.lower()
                # Parse Set-Cookie
                if k.lower() == "set-cookie":
                    cookie_name = v.split("=")[0].strip().lower()
                    cookies[cookie_name] = v
            body = r.read(4096).decode(errors="ignore").lower()
            result.headers = {k: v for k, v in headers_map.items()}
    except urllib.error.HTTPError as e:
        result.headers = {k.lower(): v.lower() for k, v in e.headers.items()}
    except Exception as e:
        log.debug(f"WAF detect fetch error {url}: {e}")
        return result

    # Match fingerprints
    for waf_name, sigs in WAF_SIGNATURES:
        matched = False
        for location, match_str in sigs:
            if location.startswith("header:"):
                hdr_name = location[7:]
                hdr_val  = headers_map.get(hdr_name, "")
                if hdr_name in headers_map:
                    if not match_str or match_str in hdr_val:
                        matched = True
                        break
            elif location.startswith("cookie:"):
                ck_name = location[7:]
                if ck_name in cookies:
                    matched = True
                    break
            elif location == "body":
                if match_str and match_str in body:
                    matched = True
                    break
        if matched:
            result.waf_names.append(waf_name)

    # Test payload — does it return 403/406/429?
    try:
        test_url = url.rstrip("/") + WAF_TEST_PAYLOAD
        test_req = urllib.request.Request(
            test_url,
            headers={"User-Agent": "Mozilla/5.0 (ReconNinja/7.0.0)"},
        )
        with urllib.request.urlopen(test_req, timeout=timeout) as r:
            pass
    except urllib.error.HTTPError as e:
        if e.code in (403, 406, 429, 503):
            result.blocked = True
    except Exception:
        pass

    result.waf_detected = bool(result.waf_names) or result.blocked
    result.tool = "passive"
    return result


# ── wafw00f integration ───────────────────────────────────────────────────────

def _wafw00f_detect(url: str, out_folder: Path, timeout: int = 30) -> Optional[WAFResult]:
    if not tool_exists("wafw00f"):
        return None

    out_file = out_folder / "wafw00f.txt"
    rc, stdout, stderr = run_cmd(
        ["wafw00f", url, "-o", str(out_file), "--format", "json"],
        timeout=timeout,
    )

    result = WAFResult(url=url, waf_detected=False, tool="wafw00f")

    # Try parsing JSON output
    if out_file.exists():
        try:
            import json
            data = json.loads(out_file.read_text())
            for entry in data.get("results", []):
                if entry.get("detected"):
                    result.waf_detected = True
                    waf = entry.get("firewall", "Unknown WAF")
                    if waf not in result.waf_names:
                        result.waf_names.append(waf)
        except Exception:
            pass

    # Fallback: parse plain-text output
    if not result.waf_names and stdout:
        for line in stdout.splitlines():
            if "is behind" in line.lower():
                parts = line.split("is behind")
                if len(parts) > 1:
                    waf_name = parts[-1].strip().strip("'\"")
                    if waf_name not in result.waf_names:
                        result.waf_names.append(waf_name)
                    result.waf_detected = True
            elif "no waf detected" in line.lower():
                result.waf_detected = False

    return result


# ── Public API ────────────────────────────────────────────────────────────────

def detect_waf(
    targets: list[str],
    out_folder: Path,
    timeout: int = 10,
) -> list[WAFResult]:
    """
    Detect WAF presence on a list of web targets.

    Uses wafw00f if available, falls back to passive header/cookie detection.

    Args:
        targets:    list of URLs to check
        out_folder: output directory
        timeout:    per-request timeout

    Returns:
        list of WAFResult
    """
    ensure_dir(out_folder)
    results: list[WAFResult] = []

    use_wafw00f = tool_exists("wafw00f")
    safe_print(
        f"[info]▶ WAF Detection — {'wafw00f + ' if use_wafw00f else ''}passive headers "
        f"on {len(targets)} target(s)[/]"
    )

    for url in targets[:20]:
        if use_wafw00f:
            wf_result = _wafw00f_detect(url, out_folder, timeout=timeout * 3)
            if wf_result:
                result = wf_result
            else:
                result = _passive_detect(url, timeout)
        else:
            result = _passive_detect(url, timeout)

        results.append(result)

        if result.waf_detected:
            names = ", ".join(result.waf_names) if result.waf_names else "Unknown WAF"
            blocked_str = " (active blocking confirmed)" if result.blocked else ""
            safe_print(
                f"  [warning]WAF detected @ {url}: [cyan]{names}[/]{blocked_str}[/]"
            )
        else:
            safe_print(f"  [dim]No WAF detected @ {url}[/]")

    detected = [r for r in results if r.waf_detected]
    safe_print(
        f"[success]✔ WAF Detection: {len(detected)}/{len(results)} target(s) behind a WAF[/]"
    )
    return results
