#!/usr/bin/env python3
from __future__ import annotations

import hashlib
import json
import sys
import zipfile
from pathlib import Path
from typing import Any


def read_text(zf: zipfile.ZipFile, name: str) -> str:
    return zf.read(name).decode("utf-8")


def read_json(zf: zipfile.ZipFile, name: str) -> Any:
    return json.loads(read_text(zf, name))


def read_jsonl(zf: zipfile.ZipFile, name: str) -> list[dict[str, Any]]:
    rows = []
    for line in read_text(zf, name).splitlines():
        line = line.strip()
        if line:
            rows.append(json.loads(line))
    return rows


def derive_key(raw_answers: list[str]) -> bytes:
    return hashlib.sha256("|".join(raw_answers).encode("utf-8")).digest()


def xor_stream(key: bytes, length: int) -> bytes:
    stream = bytearray()
    counter = 0
    while len(stream) < length:
        stream.extend(hashlib.sha256(key + counter.to_bytes(4, "big")).digest())
        counter += 1
    return bytes(stream[:length])


def decrypt_closeout(note: bytes, raw_answers: list[str]) -> dict[str, Any]:
    wrapped = json.loads(note.decode("utf-8"))
    key = derive_key(raw_answers)
    nonce = bytes.fromhex(wrapped["nonce"])
    ciphertext = bytes.fromhex(wrapped["ciphertext"])
    expected_mac = hashlib.sha256(key + nonce + ciphertext).hexdigest()
    if wrapped["mac"] != expected_mac:
        raise ValueError("closeout note MAC mismatch; one of Q01-Q09 is wrong")
    stream = xor_stream(hashlib.sha256(key + nonce).digest(), len(ciphertext))
    plaintext = bytes(a ^ b for a, b in zip(ciphertext, stream))
    return json.loads(plaintext.decode("utf-8"))


def wrap_answer(value: str) -> str:
    if value.startswith("caplag{"):
        return value
    return f"caplag{{{value}}}"


def recover_raw_answers(zf: zipfile.ZipFile) -> list[str]:
    raw: list[str] = []

    audit = read_jsonl(zf, "github/audit_log.jsonl")
    issue_event = next(
        row
        for row in audit
        if row.get("event") == "issue_comment.created"
        and "release-triage" in row.get("labels", [])
        and "agent_route" in row
    )
    raw.append(issue_event["agent_route"]["source_key"])

    agent = read_jsonl(zf, "agent/triage_assistant.jsonl")
    tool_call = next(
        row
        for row in agent
        if row.get("event") == "tool_call"
        and row.get("tool") == "deploy_plan.render"
        and row.get("input", {}).get("from_issue") == raw[0]
    )
    raw.append(tool_call["tool_call_id"])

    run = read_json(zf, "github/actions/run_8716443.json")
    if run["source"]["tool_call_id"] != raw[1]:
        raise ValueError("run source does not match recovered tool call")
    raw.append(run["tuple"])

    lockfile = read_json(zf, "repo/package-lock.after.json")
    suspicious_pkg = None
    for package_path, package_info in lockfile["packages"].items():
        if not package_path.startswith("node_modules/"):
            continue
        if package_info.get("hasInstallScript"):
            package_name = package_path.removeprefix("node_modules/")
            suspicious_pkg = f"{package_name}@{package_info['version']}"
            break
    if suspicious_pkg is None:
        raise ValueError("no install-script package found")
    raw.append(suspicious_pkg)

    npm_entry = read_json(zf, "npm/_cacache/index-v5/ledger-seal-entry.json")
    if f"{npm_entry['package']}@{npm_entry['version']}" != raw[3]:
        raise ValueError("npm cache entry does not match lockfile package")
    raw.append(npm_entry["cache_path"])

    runner_events = read_jsonl(zf, "runner/process_events.jsonl")
    run_id = int(raw[2].split("/", 1)[0].removeprefix("run-"))
    check_run_id = int(raw[2].split("/", 1)[1].removeprefix("check-"))
    runner_event = next(
        row
        for row in runner_events
        if row.get("run_id") == run_id
        and row.get("check_run_id") == check_run_id
        and raw[4] in row.get("cmd", "")
    )
    raw.append(runner_event["host"])

    oidc_tokens = read_jsonl(zf, "cloud/oidc_tokens.jsonl")
    job_id = run["jobs"][0]["job_id"]
    oidc = next(
        row
        for row in oidc_tokens
        if row.get("run_id") == run_id
        and row.get("check_run_id") == check_run_id
        and row.get("job_id") == job_id
    )
    raw.append(oidc["sub"])

    cloudtrail = read_jsonl(zf, "cloud/cloudtrail.jsonl")
    put_object = next(
        row
        for row in cloudtrail
        if row.get("eventName") == "PutObject"
        and row.get("requestParameters", {}).get("x-amz-meta-run-id") == str(run_id)
        and row.get("requestParameters", {}).get("x-amz-meta-check-run-id") == str(check_run_id)
    )
    raw.append(put_object["responseElements"]["object_locator"])

    kube_audit = read_jsonl(zf, "kubernetes/audit.jsonl")
    deploy_patch = next(
        row
        for row in kube_audit
        if row.get("verb") == "patch"
        and row.get("objectRef", {}).get("resource") == "deployments"
        and row.get("objectRef", {}).get("name") == "recon-ledger"
        and row.get("requestObject", {})
        .get("metadata", {})
        .get("annotations", {})
        .get("northstar.io/source-object")
        == raw[7]
    )
    image = deploy_patch["requestObject"]["spec"]["template"]["spec"]["containers"][0]["image"]
    raw.append(image)

    return raw


def solve_bundle(bundle_path: str | Path) -> list[str]:
    with zipfile.ZipFile(bundle_path) as zf:
        raw = recover_raw_answers(zf)
        closeout = decrypt_closeout(zf.read("cloud/final_closeout.note"), raw)
    return [wrap_answer(value) for value in raw] + [wrap_answer(closeout["final_flag"])]


def main(argv: list[str]) -> int:
    if len(argv) != 2:
        print(f"usage: {Path(argv[0]).name} public/runner_ashfall_case.zip", file=sys.stderr)
        return 2
    answers = solve_bundle(argv[1])
    for idx, answer in enumerate(answers, 1):
        print(f"Q{idx:02d}: {answer}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))