#!/usr/bin/env bash
# =============================================================================
# NoTrust.now — generate your post-quantum key pairs locally.
#
#   ML-KEM-768 (FIPS 203)  confidentiality: others encrypt secrets TO this key
#   ML-DSA-65  (FIPS 204)  identity: you sign the secrets you send with it
#
# Nothing touches the network. Share the two .pub files; keep the .key files.
# A web user adds you as a contact by importing your .kem.pub (and .dsa.pub to
# check your signature) in the vault; you give them nothing else.
#
# Usage:
#   ./keygenpqc.sh [prefix]          (default prefix: ./notrust)
#
# Writes, refusing to overwrite:
#   <prefix>.kem.pub   ML-KEM-768 public key  (1184 B, Base64)  -> share it
#   <prefix>.kem.key   ML-KEM-768 secret key  (2400 B, Base64)  -> decryptpqc.sh -k
#   <prefix>.dsa.pub   ML-DSA-65 public key   (1952 B, Base64)  -> share it
#   <prefix>.dsa.key   ML-DSA-65 secret key   (4032 B, Base64)  -> createpqc.sh -s
#
# Requirements:  pip install pqcrypto
# =============================================================================
set -euo pipefail

PREFIX="${1:-./notrust}"
case "$PREFIX" in -h|--help) awk 'NR > 2 && /^# ====/ { exit } NR > 2' "$0"; exit 0 ;; esac

for f in "$PREFIX.kem.pub" "$PREFIX.kem.key" "$PREFIX.dsa.pub" "$PREFIX.dsa.key"; do
  if [ -e "$f" ]; then
    echo "ERROR: $f already exists. Pick another prefix; overwriting a secret key loses every secret sent to it." >&2
    exit 2
  fi
done

python3 -c 'import pqcrypto.kem.ml_kem_768, pqcrypto.sign.ml_dsa_65' 2>/dev/null \
  || { echo "ERROR: missing dependency. Install:  pip install pqcrypto" >&2; exit 1; }

# Secret keys are created 0600 from the start, not chmod-ed after being written.
umask 077

NT_PREFIX="$PREFIX" python3 - <<'PY'
import base64, hashlib, os, datetime
try:  # pqcrypto >= 1.0
    from pqcrypto.kem.ml_kem_768 import keygen as kem_keygen
    from pqcrypto.sign.ml_dsa_65 import keygen as dsa_keygen
except ImportError:  # pqcrypto 0.x
    from pqcrypto.kem.ml_kem_768 import generate_keypair as kem_keygen
    from pqcrypto.sign.ml_dsa_65 import generate_keypair as dsa_keygen

prefix = os.environ["NT_PREFIX"]
b64 = lambda b: base64.b64encode(bytes(b)).decode()
# Same fingerprint the vault shows: first 16 hex chars of SHA-256(raw public key).
fingerprint = lambda pk: hashlib.sha256(bytes(pk)).hexdigest()[:16].upper()
now = datetime.datetime.now(datetime.timezone.utc).isoformat()

def write(path, algorithm, kind, key, fp):
    # Same header as the vault's own export, so the vault importer reads it as is.
    header = (f"# {algorithm} {kind} Key\n# Fingerprint: {fp}\n"
              f"# Generated: {now}\n# Format: Base64\n\n")
    with open(path, "x") as fh:
        fh.write(header + b64(key) + "\n")

kem_pk, kem_sk = kem_keygen()
dsa_pk, dsa_sk = dsa_keygen()
kem_fp, dsa_fp = fingerprint(kem_pk), fingerprint(dsa_pk)

write(prefix + ".kem.pub", "ML-KEM-768", "Public", kem_pk, kem_fp)
write(prefix + ".kem.key", "ML-KEM-768", "Secret", kem_sk, kem_fp)
write(prefix + ".dsa.pub", "ML-DSA-65", "Public", dsa_pk, dsa_fp)
write(prefix + ".dsa.key", "ML-DSA-65", "Secret", dsa_sk, dsa_fp)

print(f"ML-KEM-768  {prefix}.kem.pub / .kem.key   fingerprint {kem_fp}")
print(f"ML-DSA-65   {prefix}.dsa.pub / .dsa.key   fingerprint {dsa_fp}")
print("Share the .pub files. The .key files never leave this machine.")
PY