#!/usr/bin/env -S uv run
"""Generate or derive Chrome extension key and ID.

Usage:
  python scripts/gen_extension_key.py                  # generate new key pair
  python scripts/gen_extension_key.py --from-manifest  # derive ID from extension/manifest.json
  python scripts/gen_extension_key.py --key <base64>   # derive ID from given public key
"""
import argparse, hashlib
import base64, json, sys
from pathlib import Path

def public_key_to_extension_id(pub_key_der:bytes) -> str:
  digest = hashlib.sha256(pub_key_der).hexdigest()
  return "".join(chr(ord("a") + int(c, 16)) for c in digest[:32])

def derive_from_key_b64(key_b64:str) -> tuple[str, str]:
  der = base64.b64decode(key_b64)
  ext_id = public_key_to_extension_id(der)
  return key_b64, ext_id

def generate_new_key() -> tuple[str, str, str]:
  try:
    from cryptography.hazmat.primitives.asymmetric import rsa
    from cryptography.hazmat.primitives.serialization import (
      Encoding,
      PublicFormat,
      PrivateFormat,
      NoEncryption,
    )
  except ImportError:
    print("Install 'cryptography' to generate new keys: pip install cryptography", file=sys.stderr)
    sys.exit(1)

  private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
  pub_der = private_key.public_key().public_bytes(Encoding.DER, PublicFormat.SubjectPublicKeyInfo)
  priv_pem = private_key.private_bytes(Encoding.PEM, PrivateFormat.TraditionalOpenSSL, NoEncryption()).decode()

  key_b64 = base64.b64encode(pub_der).decode()
  ext_id = public_key_to_extension_id(pub_der)
  return key_b64, ext_id, priv_pem

if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="Chrome extension key/ID tool")
  group = parser.add_mutually_exclusive_group()
  group.add_argument("--from-manifest", action="store_true", help="Derive ID from extension/manifest.json")
  group.add_argument("--key", metavar="BASE64", help="Derive ID from given base64 public key")
  args = parser.parse_args()

  if args.from_manifest:
    manifest_path = Path(__file__).parent.parent / "extension" / "manifest.json"
    manifest = json.loads(manifest_path.read_text())
    key_b64 = manifest.get("key")
    if not key_b64:
      print("No 'key' field in manifest.json", file=sys.stderr)
      sys.exit(1)
    key_b64, ext_id = derive_from_key_b64(key_b64)
    print(f"Extension ID: {ext_id}")
    print(f"Key (b64):  {key_b64}")

  elif args.key:
    key_b64, ext_id = derive_from_key_b64(args.key)
    print(f"Extension ID: {ext_id}")

  else:
    key_b64, ext_id, priv_pem = generate_new_key()
    print(f"Extension ID: {ext_id}")
    print(f"Key (b64):  {key_b64}")
    print()
    print("Add this to extension/manifest.json:")
    print(f'  "key": "{key_b64}"')
    print()
    print("Private key (keep secret, needed to re-derive same ID):")
    print(priv_pem)