#!/usr/bin/env python3
"""
Sign mirror-thin-v1 artefacts using an Ed25519 key and emit DSSE + TUF signatures.

Usage:
  python scripts/mirror/sign_thin_bundle.py \
    --key out/mirror/thin/tuf/keys/mirror-ed25519-test-1.pem \
    --manifest out/mirror/thin/mirror-thin-v1.manifest.json \
    --tar out/mirror/thin/mirror-thin-v1.tar.gz \
    --tuf-dir out/mirror/thin/tuf \
    --time-anchor out/mirror/thin/stage-v1/layers/time-anchor.json

Writes:
  - mirror-thin-v1.manifest.dsse.json
  - mirror-thin-v1.bundle.dsse.json (optional, when --bundle is provided)
  - updates signatures in root.json, targets.json, snapshot.json, timestamp.json
"""
import argparse, base64, json, pathlib, hashlib
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

def b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()

def load_key(path: pathlib.Path) -> Ed25519PrivateKey:
    return serialization.load_pem_private_key(path.read_bytes(), password=None)

def keyid_from_pub(pub_path: pathlib.Path) -> str:
    raw = pub_path.read_bytes()
    return hashlib.sha256(raw).hexdigest()

def sign_bytes(key: Ed25519PrivateKey, data: bytes) -> bytes:
    return key.sign(data)

def write_json(path: pathlib.Path, obj):
    path.write_text(json.dumps(obj, indent=2, sort_keys=True) + "\n")

def sign_tuf(path: pathlib.Path, keyid: str, key: Ed25519PrivateKey):
    data = path.read_bytes()
    sig = sign_bytes(key, data)
    obj = json.loads(data)
    obj["signatures"] = [{"keyid": keyid, "sig": b64url(sig)}]
    write_json(path, obj)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--key", required=True, type=pathlib.Path)
    ap.add_argument("--manifest", required=True, type=pathlib.Path)
    ap.add_argument("--tar", required=True, type=pathlib.Path)
    ap.add_argument("--tuf-dir", required=True, type=pathlib.Path)
    ap.add_argument("--bundle", required=False, type=pathlib.Path)
    ap.add_argument("--time-anchor", required=False, type=pathlib.Path)
    args = ap.parse_args()

    key = load_key(args.key)
    pub_path = args.key.with_suffix(".pub")
    keyid = keyid_from_pub(pub_path)

    manifest_bytes = args.manifest.read_bytes()
    sig = sign_bytes(key, manifest_bytes)
    dsse = {
        "payloadType": "application/vnd.stellaops.mirror.manifest+json",
        "payload": b64url(manifest_bytes),
        "signatures": [{"keyid": keyid, "sig": b64url(sig)}],
    }
    dsse_path = args.manifest.with_suffix(".dsse.json")
    write_json(dsse_path, dsse)

    if args.bundle:
        bundle_bytes = args.bundle.read_bytes()
        bundle_sig = sign_bytes(key, bundle_bytes)
        bundle_dsse = {
            "payloadType": "application/vnd.stellaops.mirror.bundle+json",
            "payload": b64url(bundle_bytes),
            "signatures": [{"keyid": keyid, "sig": b64url(bundle_sig)}],
        }
        bundle_dsse_path = args.bundle.with_suffix(".dsse.json")
        write_json(bundle_dsse_path, bundle_dsse)

    anchor_dsse_path = None
    if args.time_anchor:
        anchor_bytes = args.time_anchor.read_bytes()
        anchor_sig = sign_bytes(key, anchor_bytes)
        anchor_dsse = {
            "payloadType": "application/vnd.stellaops.time-anchor+json",
            "payload": b64url(anchor_bytes),
            "signatures": [{"keyid": keyid, "sig": b64url(anchor_sig)}],
        }
        anchor_dsse_path = args.time_anchor.with_suffix(".dsse.json")
        write_json(anchor_dsse_path, anchor_dsse)

    # update TUF metadata
    for name in ["root.json", "targets.json", "snapshot.json", "timestamp.json"]:
        sign_tuf(args.tuf_dir / name, keyid, key)

    parts = [f"manifest DSSE -> {dsse_path}"]
    if args.bundle:
        parts.append(f"bundle DSSE -> {bundle_dsse_path}")
    if anchor_dsse_path:
        parts.append(f"time anchor DSSE -> {anchor_dsse_path}")
    parts.append("TUF metadata updated")
    print(f"Signed DSSE + TUF using keyid {keyid}; " + ", ".join(parts))

if __name__ == "__main__":
    main()