#!/usr/bin/env bash
# =============================================================================
# NoTrust.now — create a post-quantum one-time secret from the command line.
#
# Everything cryptographic runs HERE; the server only receives the encrypted
# blob. The payload is the same v2 format the web app writes, so the recipient
# can open the link in the browser (vault unlocked) or with decryptpqc.sh:
#
#   KEM:   ML-KEM-768 (FIPS 203), encapsulated to the RECIPIENT's public key
#   Key:   Scrypt(passphrase, salt; N=2^16, r=8, p=1) XOR ML-KEM shared secret
#   AEAD:  ChaCha20-Poly1305 (RFC 8439)
#   Sign:  ML-DSA-65 (FIPS 204), optional, with YOUR signing key
#
# Opening the secret takes both the recipient's private key and the passphrase.
#
# Usage:
#   ./createpqc.sh -r <recipient.kem.pub> [options] ["<secret>" | -]
#
#   -r FILE   recipient's ML-KEM-768 public key (required). Raw Base64, the file
#             the vault exports ("Export public key"), a keygenpqc.sh .kem.pub
#             or an OpenSSL PEM. To keep a secret for yourself, use your own.
#   -s FILE   your ML-DSA-65 secret key, to sign (needs -S too)
#   -S FILE   your ML-DSA-65 public key
#   -t MIN    time to live in minutes, 1 to 60 (default 15)
#   -v N      maximum views, 1 to 10 (default 1)
#             (the server applies these limits to every API call today)
#   -b        print the POST body instead of sending it (no API key needed)
#   secret    the text to protect; omitted or "-" reads it from stdin
#
# Environment:
#   NOTRUST_API_KEY     your API key (sk_live_...), required to publish.
#                       Create one in your account at https://notrust.now/en/account
#   NOTRUST_PASSPHRASE  the passphrase (min. 6 characters); prompted if unset
#   API_BASE            default https://notrust.now
#
# Prints the share link on stdout; progress goes to stderr.
# Requirements:  pip install pqcrypto cryptography   (and curl)
# =============================================================================
set -euo pipefail

API="${API_BASE:-https://notrust.now}"
RECIPIENT=""; SIGN_SK=""; SIGN_PK=""; TTL=15; VIEWS=1; BODY_ONLY=0

usage() { awk 'NR > 2 && /^# ====/ { exit } NR > 2' "$0" >&2; exit 2; }

while getopts ":r:s:S:t:v:bh" opt; do
  case "$opt" in
    r) RECIPIENT="$OPTARG" ;;
    s) SIGN_SK="$OPTARG" ;;
    S) SIGN_PK="$OPTARG" ;;
    t) TTL="$OPTARG" ;;
    v) VIEWS="$OPTARG" ;;
    b) BODY_ONLY=1 ;;
    *) usage ;;
  esac
done
shift $((OPTIND - 1))

[ -n "$RECIPIENT" ] || { echo "ERROR: -r <recipient.kem.pub> is required: a secret is always encrypted to a recipient's key." >&2; usage; }
[ -f "$RECIPIENT" ] || { echo "ERROR: recipient key not found: $RECIPIENT" >&2; exit 2; }
if [ -n "$SIGN_SK" ] || [ -n "$SIGN_PK" ]; then
  # ML-DSA cannot derive the public key from the secret key, and the payload
  # carries the public key so the recipient can verify.
  [ -f "$SIGN_SK" ] && [ -f "$SIGN_PK" ] || { echo "ERROR: signing needs both -s <you.dsa.key> and -S <you.dsa.pub>." >&2; exit 2; }
fi
# The server's limits for API calls (it does not read the key owner's plan
# yet). The length check keeps an absurd number from overflowing [ ].
case "$TTL" in ''|*[!0-9]*) echo "ERROR: -t must be a whole number of minutes." >&2; exit 2 ;; esac
case "$VIEWS" in ''|*[!0-9]*) echo "ERROR: -v must be a whole number." >&2; exit 2 ;; esac
[ "${#TTL}" -le 4 ] && [ "$TTL" -ge 1 ] && [ "$TTL" -le 60 ] || { echo "ERROR: -t must be 1 to 60 minutes." >&2; exit 2; }
[ "${#VIEWS}" -le 4 ] && [ "$VIEWS" -ge 1 ] && [ "$VIEWS" -le 10 ] || { echo "ERROR: -v must be 1 to 10 views." >&2; exit 2; }
if [ "$BODY_ONLY" -eq 0 ] && [ -z "${NOTRUST_API_KEY:-}" ]; then
  echo "ERROR: NOTRUST_API_KEY is not set. Creating secrets through the API requires an API key:" >&2
  echo "       sign in at https://notrust.now/en/login, create one at https://notrust.now/en/account" >&2
  echo "       and export NOTRUST_API_KEY=sk_live_..." >&2
  exit 2
fi

python3 -c 'import pqcrypto.kem.ml_kem_768, pqcrypto.sign.ml_dsa_65
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305' 2>/dev/null \
  || { echo "ERROR: missing dependencies. Install:  pip install pqcrypto cryptography" >&2; exit 1; }

# --- The secret: argument, or stdin (keeps it out of the process list) -------
if [ $# -ge 1 ] && [ "$1" != "-" ]; then
  SECRET="$1"
else
  [ -t 0 ] && echo "Type the secret, then Ctrl-D:" >&2
  SECRET="$(cat)"
fi
[ -n "$SECRET" ] || { echo "ERROR: the secret is empty." >&2; exit 2; }

# --- The passphrase: second factor, mandatory ---------------------------------
PASS="${NOTRUST_PASSPHRASE:-}"
if [ -z "$PASS" ]; then
  [ -r /dev/tty ] || { echo "ERROR: set NOTRUST_PASSPHRASE (no terminal to prompt on)." >&2; exit 2; }
  read -rs -p "Passphrase (min. 6 characters): " PASS </dev/tty; echo >&2
  read -rs -p "Repeat the passphrase: " PASS2 </dev/tty; echo >&2
  [ "$PASS" = "$PASS2" ] || { echo "ERROR: the passphrases do not match." >&2; exit 2; }
fi
[ "${#PASS}" -ge 6 ] || { echo "ERROR: the passphrase needs at least 6 characters." >&2; exit 2; }

WORK="$(mktemp -d)"; chmod 700 "$WORK"
trap 'rm -rf "$WORK"' EXIT

# Secret and passphrase reach Python through its environment, never its argv.
NT_SECRET="$SECRET" NT_PASS="$PASS" NT_RECIPIENT="$RECIPIENT" NT_SIGN_SK="$SIGN_SK" \
NT_SIGN_PK="$SIGN_PK" NT_TTL="$TTL" NT_VIEWS="$VIEWS" NT_OUT="$WORK/body.json" python3 - <<'PY'
import base64, hashlib, json, os, sys, time, uuid
try:  # pqcrypto >= 1.0
    from pqcrypto.kem.ml_kem_768 import encaps as kem_encaps
    from pqcrypto.sign.ml_dsa_65 import sign as dsa_sign
except ImportError:  # pqcrypto 0.x
    from pqcrypto.kem.ml_kem_768 import encrypt as kem_encaps
    from pqcrypto.sign.ml_dsa_65 import sign as dsa_sign
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt

KEM_PK, DSA_PK, DSA_SK = 1184, 1952, 4032
# The web app's parameters (pqc-crypto.ts SCRYPT_PARAMS); stored in the payload.
SCRYPT = {"N": 65536, "r": 8, "p": 1, "dkLen": 32}

def fail(msg):
    print("ERROR: " + msg, file=sys.stderr)
    sys.exit(1)

def b64(data):
    return base64.b64encode(bytes(data)).decode()

def fingerprint(pk):
    # Same value the vault shows next to each key and contact.
    return hashlib.sha256(bytes(pk)).hexdigest()[:16].upper()

def read_key(path, size, label, secret):
    """Raw key bytes from raw Base64, a vault/keygen export or an OpenSSL PEM."""
    lines = [l.strip() for l in open(path).read().splitlines()]
    body = "".join(l for l in lines
                   if l and not l.startswith(("#", "-----")) and ":" not in l)
    try:
        raw = base64.b64decode(body, validate=True)
    except Exception:
        fail(f"{label} ({path}) is not Base64.")
    if len(raw) == size:
        return raw
    # DER wrappers keep the key at the end: SubjectPublicKeyInfo ends in the
    # BIT STRING, PKCS#8 ends in the expanded key's OCTET STRING (04 82 len).
    if raw[:1] == b"\x30" and len(raw) > size:
        if not secret and len(raw) - size <= 64:
            return raw[-size:]
        if secret and raw[-size - 4:-size] == b"\x04\x82" + size.to_bytes(2, "big"):
            return raw[-size:]
    fail(f"{label} ({path}) has {len(raw)} bytes; expected {size}.")

recipient_pk = read_key(os.environ["NT_RECIPIENT"], KEM_PK, "Recipient ML-KEM-768 public key", False)
signing = bool(os.environ["NT_SIGN_SK"])
if signing:
    dsa_sk = read_key(os.environ["NT_SIGN_SK"], DSA_SK, "Your ML-DSA-65 secret key", True)
    dsa_pk = read_key(os.environ["NT_SIGN_PK"], DSA_PK, "Your ML-DSA-65 public key", False)
    # FIPS 204 stores tr = SHAKE256(pk, 64) at sk[64:128]: a mismatched pair
    # would sign fine and fail verification on the recipient's side.
    if hashlib.shake_256(dsa_pk).digest(64) != dsa_sk[64:128]:
        fail("-s and -S are not the same ML-DSA-65 key pair.")

ttl_minutes = int(os.environ["NT_TTL"])
max_views = int(os.environ["NT_VIEWS"])
secret = os.environ["NT_SECRET"].encode()
passphrase = os.environ["NT_PASS"].encode()

# Metadata as the web writes it; ttl is in SECONDS here (the request's ttl is minutes).
metadata = {
    "ttl": ttl_minutes * 60,
    "maxViews": max_views,
    "requireSignature": signing,
    "id": str(uuid.uuid4()),
    "clientSidePQC": True,
    "timestamp": int(time.time() * 1000),
    "version": 2,
}

# 1) ML-KEM-768 encapsulation to the recipient's public key (possession factor).
kem_ct, shared = kem_encaps(recipient_pk)
# 2) Scrypt of the passphrase (knowledge factor), XORed with the shared secret.
salt, iv = os.urandom(32), os.urandom(12)
scrypt_key = Scrypt(salt=salt, length=SCRYPT["dkLen"], n=SCRYPT["N"],
                    r=SCRYPT["r"], p=SCRYPT["p"]).derive(passphrase)
key = bytes(a ^ b for a, b in zip(scrypt_key, bytes(shared)[:32]))
# 3) ChaCha20-Poly1305, no associated data (as the web).
ciphertext = ChaCha20Poly1305(key).encrypt(iv, secret, None)

kyber_data = {"publicKey": b64(recipient_pk), "ciphertext": b64(kem_ct)}
signature_data = None
if signing:
    # The web's canonical transcript (buildSignatureTranscript): it binds the
    # metadata, the AEAD fields, the KEM ciphertext and the recipient key.
    transcript = "\n".join([
        "pqc-sig-v1",
        f"id={metadata['id']}",
        f"ttl={metadata['ttl']}",
        f"maxViews={metadata['maxViews']}",
        f"version={metadata['version']}",
        f"ciphertext={b64(ciphertext)}",
        f"iv={b64(iv)}",
        f"salt={b64(salt)}",
        f"kemCiphertext={kyber_data['ciphertext']}",
        f"recipientPublicKey={kyber_data['publicKey']}",
    ]).encode()
    signature_data = {"signature": b64(dsa_sign(dsa_sk, transcript)),
                      "publicKey": b64(dsa_pk), "algorithm": "ML-DSA-65"}

payload = {
    "ciphertext": b64(ciphertext),
    "iv": b64(iv),
    "salt": b64(salt),
    "algorithm": "ChaCha20-Poly1305",
    "keyDerivation": "Scrypt+KEM",
    "scryptParams": SCRYPT,
    "ttl": metadata["ttl"],
    "maxViews": max_views,
    "hasPassphrase": True,
    "metadata": metadata,
    "kyberData": kyber_data,
    "signatureData": signature_data,
}
body = {"secret": json.dumps(payload, separators=(",", ":")), "ttl": ttl_minutes, "maxViews": max_views}
with open(os.environ["NT_OUT"], "w") as fh:
    json.dump(body, fh)

print(f"Encrypted locally for recipient key {fingerprint(recipient_pk)}"
      + (f", signed with {fingerprint(dsa_pk)}" if signing else ", unsigned"), file=sys.stderr)
PY

if [ "$BODY_ONLY" -eq 1 ]; then
  cat "$WORK/body.json"; echo
  exit 0
fi

echo "Publishing the encrypted blob to $API ..." >&2
HTTP=$(curl -sS -o "$WORK/resp.json" -w '%{http_code}' -X POST "$API/api/v1/secrets" \
  -H "X-API-Key: $NOTRUST_API_KEY" \
  -H 'Content-Type: application/json' \
  --data-binary @"$WORK/body.json") || { echo "ERROR: could not reach $API" >&2; exit 1; }

if [ "$HTTP" != "201" ]; then
  # Errors come as {"error": CODE, "message": text}; line 1 is the code.
  ERR=$(python3 -c 'import json,sys
try:
    d = json.load(open(sys.argv[1])); print(d.get("error") or ""); print(d.get("message") or d)
except Exception:
    print(); print(open(sys.argv[1], errors="replace").read()[:200])' "$WORK/resp.json")
  echo "ERROR: HTTP $HTTP — ${ERR#*$'\n'}" >&2
  ACCOUNT="https://notrust.now/en/account"
  case "$HTTP ${ERR%%$'\n'*}" in
    "401 INVALID_API_KEY")  echo "       The server does not know this key (mistyped or deleted). Copy it again or create one at $ACCOUNT" >&2 ;;
    "401 REVOKED_API_KEY")  echo "       This key was revoked or deactivated. Create a new one at $ACCOUNT" >&2 ;;
    "401 EXPIRED_API_KEY")  echo "       This key has expired. Create a new one at $ACCOUNT" >&2 ;;
    "401 API_KEY_REQUIRED_FOR_PQC") echo "       The server saw no API key: check that NOTRUST_API_KEY is not blank and no proxy strips X-API-Key." >&2 ;;
    401*)                   echo "       Check NOTRUST_API_KEY: it must be a valid key from $ACCOUNT" >&2 ;;
  esac
  exit 1
fi

python3 -c 'import json,sys
d = json.load(open(sys.argv[1]))
print("Created. Token %s, expires %s" % (d.get("token"), d.get("expiresAt")), file=sys.stderr)
print(d.get("secretUrl") or d.get("token"))' "$WORK/resp.json"
echo "Send the passphrase through a different channel than the link." >&2
