#!/usr/bin/env bash
# =============================================================================
# NoTrust.now — open a post-quantum one-time secret from the command line.
#
# Fetches the encrypted blob (this CONSUMES a view) and decrypts it HERE; your
# private key never leaves this machine. Reads the v2 format that the web app
# and createpqc.sh write:
#
#   1. Verifies the ML-DSA-65 signature, if any, BEFORE decrypting. A signature
#      that does not verify, or a missing one where it is required, aborts.
#   2. Decapsulates ML-KEM-768 with your secret key.
#   3. Rebuilds the key: Scrypt(passphrase, salt) XOR shared secret.
#   4. Decrypts ChaCha20-Poly1305.
#
# The key files are checked before anything is fetched. Once the view is spent
# the blob is saved, and any failure prints the -i command that retries it.
#
# Usage:
#   ./decryptpqc.sh -k <you.kem.key> [-p <sender.dsa.pub>] <share link | token>
#   ./decryptpqc.sh -k <you.kem.key> [-p <sender.dsa.pub>] -i <response.json>
#
#   link      the share link (.../secret/<token>), the page it opens in the
#             browser (.../en/retrieve?token=<token>) or the bare token
#   -k FILE   your ML-KEM-768 secret key (keygenpqc.sh .kem.key, raw Base64 or
#             an OpenSSL PKCS#8 PEM that holds the expanded key)
#   -p FILE   the ML-DSA-65 public key you expect the sender to sign with.
#             With it, an unsigned secret or another signer is refused.
#   -i FILE   decrypt a saved GET /api/v1/secrets/{token} response instead of
#             fetching one. A blob this script saved is deleted once it opens.
#
# Environment:
#   NOTRUST_PASSPHRASE  the passphrase; prompted if unset
#   API_BASE            default https://notrust.now
#
# Prints the secret on stdout; the verification report goes to stderr.
# Requirements:  pip install pqcrypto cryptography   (and curl)
# =============================================================================
set -euo pipefail

API="${API_BASE:-https://notrust.now}"
KEM_KEY=""; PINNED=""; INPUT=""; SAVED=""; FETCHED=0

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

# A blob this script saved: notrust-<token>.XXXXXX in the temp dir (see below).
saved_blob() {
  local dir tmp
  [[ "$(basename -- "$1")" =~ ^notrust-[0-9a-f]{32}\.[A-Za-z0-9]{6}$ ]] || return 1
  dir="$(cd -- "$(dirname -- "$1")" 2>/dev/null && pwd -P)" &&
    tmp="$(cd -- "${TMPDIR:-/tmp}" 2>/dev/null && pwd -P)" && [ "$dir" = "$tmp" ]
}

while getopts ":k:p:i:h" opt; do
  case "$opt" in
    k) KEM_KEY="$OPTARG" ;;
    p) PINNED="$OPTARG" ;;
    i) INPUT="$OPTARG" ;;
    *) usage ;;
  esac
done
shift $((OPTIND - 1))

[ -n "$KEM_KEY" ] || { echo "ERROR: -k <you.kem.key> is required." >&2; usage; }
[ -f "$KEM_KEY" ] || { echo "ERROR: key not found: $KEM_KEY" >&2; exit 2; }
[ -z "$PINNED" ] || [ -f "$PINNED" ] || { echo "ERROR: sender key not found: $PINNED" >&2; exit 2; }
if [ -z "$INPUT" ]; then
  [ $# -ge 1 ] || { echo "ERROR: give a share link, a token or -i <response.json>." >&2; usage; }
  # The token= parameter of the browser URL, else the link's last path segment.
  case "$1" in
    *[?\&]token=*) TOKEN="${1#*[?&]token=}"; TOKEN="${TOKEN%%[&#]*}" ;;
    *)             TOKEN="${1%%[?#]*}"; TOKEN="${TOKEN%/}"; TOKEN="${TOKEN##*/}" ;;
  esac
  # The server's tokens are 32 lowercase hex characters (a UUID without dashes).
  if ! [[ "$TOKEN" =~ ^[0-9a-f]{32}$ ]]; then
    echo "ERROR: no secret token in: $1" >&2
    echo "       Give the share link (.../secret/<token>), the .../retrieve?token=<token> URL" >&2
    echo "       or the token itself (32 lowercase hex characters). Nothing was fetched." >&2
    exit 2
  fi
else
  [ -f "$INPUT" ] || { echo "ERROR: file not found: $INPUT" >&2; exit 2; }
  if saved_blob "$INPUT"; then SAVED="$INPUT"; fi
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 Python side: "check" reads the key files, "decrypt" opens the blob ---
pqc() {
  NT_MODE="$1" NT_INPUT="$INPUT" NT_PASS="${PASS:-}" NT_KEM_KEY="$KEM_KEY" NT_PINNED="$PINNED" python3 - <<'PY'
import base64, hashlib, json, os, sys

KEM_SK, KEM_PK, KEM_CT, DSA_PK, DSA_SK, DSA_SIG = 2400, 1184, 1088, 1952, 4032, 3309
# What a key file of the wrong kind usually is, by its raw size.
KNOWN = {KEM_PK: "an ML-KEM-768 public key (.kem.pub)", KEM_SK: "an ML-KEM-768 secret key (.kem.key)",
         DSA_PK: "an ML-DSA-65 public key (.dsa.pub)", DSA_SK: "an ML-DSA-65 secret key (.dsa.key)"}

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

# A malformed payload ends in one ERROR line, not a traceback (exit status 1).
sys.excepthook = lambda kind, value, tb: print(
    f"ERROR: could not process the secret ({kind.__name__}: {value}).", file=sys.stderr)

def b64d(value, size=None, label="field"):
    try:
        raw = base64.b64decode(value, validate=True)
    except Exception:
        fail(f"{label} is not valid Base64.")
    if size is not None and len(raw) != size:
        fail(f"{label} has {len(raw)} bytes; expected {size}.")
    return raw

def fingerprint(pk):
    return hashlib.sha256(bytes(pk)).hexdigest()[:16].upper()

def read_key(path, size, label, secret, flag):
    """Raw key bytes from raw Base64, a vault/keygen export or an OpenSSL PEM."""
    try:
        lines = [l.strip() for l in open(path).read().splitlines()]
    except (OSError, UnicodeDecodeError):
        fail(f"{label} ({flag} {path}) is not a readable key file.")
    body = "".join(l for l in lines
                   if l and not l.startswith(("#", "-----")) and ":" not in l)
    raw = b64d(body, None, f"{label} ({flag} {path})")
    if len(raw) == size:
        return raw
    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:]
    what = KNOWN.get(len(raw))
    fail(f"{label} ({flag} {path}) has {len(raw)} bytes; expected {size}" + (f": it looks like {what}." if what else "."))

kem_sk = read_key(os.environ["NT_KEM_KEY"], KEM_SK, "Your ML-KEM-768 secret key", True, "-k")
# A FIPS 203 decapsulation key is dk_PKE | ek | H(ek) | z; the embedded public
# key ek sits at bytes 1152..2336 and its SHA3-256 right after it.
kem_ek = kem_sk[1152:1152 + KEM_PK]
if hashlib.sha3_256(kem_ek).digest() != kem_sk[1152 + KEM_PK:1152 + KEM_PK + 32]:
    fail(f"Your ML-KEM-768 secret key (-k {os.environ['NT_KEM_KEY']}) is damaged: its embedded public key fails its hash.")
pinned = read_key(os.environ["NT_PINNED"], DSA_PK, "Expected sender key", False, "-p") if os.environ["NT_PINNED"] else None
if os.environ["NT_MODE"] == "check":
    sys.exit(0)

try:  # pqcrypto >= 1.0
    from pqcrypto.kem.ml_kem_768 import decaps as kem_decaps
    from pqcrypto.sign.ml_dsa_65 import verify as dsa_verify_raw
except ImportError:  # pqcrypto 0.x
    from pqcrypto.kem.ml_kem_768 import decrypt as kem_decaps
    from pqcrypto.sign.ml_dsa_65 import verify as dsa_verify_raw
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt

def dsa_verify(pk, msg, sig):
    try:
        result = dsa_verify_raw(pk, msg, sig)
        return True if result is None else bool(result)  # 1.x returns None, 0.x a bool
    except Exception:
        return False

try:
    response = json.load(open(os.environ["NT_INPUT"]))
    data = response.get("encryptedData", response) if isinstance(response, dict) else None
    if isinstance(data, str):
        data = json.loads(data)
except Exception:
    fail("the response is not JSON, or its encryptedData is not.")
if not isinstance(data, dict) or not all(isinstance(data.get(f), str) for f in ("ciphertext", "iv", "salt")):
    fail("no encrypted payload in the response.")

kyber = data.get("kyberData") or {}
md = data.get("metadata") or {}
if not kyber.get("ciphertext") or not (data.get("keyDerivation") == "Scrypt+KEM" or (md.get("version") or 0) >= 2):
    fail("this is not a v2 post-quantum payload (ML-KEM + passphrase).")

if kyber.get("publicKey") and b64d(kyber["publicKey"], None, "kyberData.publicKey")[-KEM_PK:] != kem_ek:
    fail("this secret was encrypted to a different ML-KEM key than yours "
         f"(recipient fingerprint {fingerprint(b64d(kyber['publicKey'])[-KEM_PK:])}, "
         f"yours {fingerprint(kem_ek)}).")

# --- Sender authentication, before any plaintext exists -----------------------
sig = data.get("signatureData") or {}
signer = None
if sig.get("signature"):
    transcript = "\n".join([
        "pqc-sig-v1",
        f"id={'' if md.get('id') is None else md.get('id')}",
        f"ttl={'' if md.get('ttl') is None else md.get('ttl')}",
        f"maxViews={'' if md.get('maxViews') is None else md.get('maxViews')}",
        f"version={'' if md.get('version') is None else md.get('version')}",
        f"ciphertext={data['ciphertext']}",
        f"iv={data['iv']}",
        f"salt={data['salt']}",
        f"kemCiphertext={kyber.get('ciphertext') or ''}",
        f"recipientPublicKey={kyber.get('publicKey') or ''}",
    ]).encode()
    # A pinned key is the only requirement the server cannot strip or swap.
    vk = pinned if pinned is not None else (b64d(sig["publicKey"], None, "signer key")[-DSA_PK:] if sig.get("publicKey") else None)
    if vk is None:
        fail("the secret is signed but carries no signer key and you pinned none (-p).")
    if not dsa_verify(vk, transcript, b64d(sig["signature"], DSA_SIG, "signature")):
        fail("signature verification FAILED: the secret was altered or not signed by "
             + ("the sender you expected." if pinned is not None else "the key it claims."))
    signer = vk
elif pinned is not None:
    fail("you required a signature from a specific sender (-p), but the secret is unsigned.")
elif md.get("requireSignature") is True:
    fail("the secret was created with a mandatory signature, but none is present.")

# --- Decrypt ------------------------------------------------------------------
params = data.get("scryptParams")
if not isinstance(params, dict):
    fail("the payload has no scryptParams.")
n, r, p, dk_len = (params.get(k) for k in ("N", "r", "p", "dkLen"))
# Bounds keep a hostile payload from asking for gigabytes of scrypt memory.
if not (isinstance(n, int) and 1024 <= n <= 1 << 20 and n & (n - 1) == 0
        and isinstance(r, int) and 1 <= r <= 32 and isinstance(p, int) and 1 <= p <= 16 and dk_len == 32):
    fail(f"unexpected scrypt parameters {params}.")

shared = bytes(kem_decaps(kem_sk, b64d(kyber["ciphertext"], KEM_CT, "kyberData.ciphertext")))
scrypt_key = Scrypt(salt=b64d(data["salt"], None, "salt"), length=32, n=n, r=r, p=p).derive(os.environ["NT_PASS"].encode())
key = bytes(a ^ b for a, b in zip(scrypt_key, shared[:32]))
try:
    plaintext = ChaCha20Poly1305(key).decrypt(b64d(data["iv"], 12, "iv"), b64d(data["ciphertext"], None, "ciphertext"), None)
except Exception:
    fail("could not decrypt: wrong passphrase (or wrong key).")

if signer is not None:
    who = "the sender you pinned" if pinned is not None else "the key embedded in the secret; compare it with your contact's"
    print(f"Signature valid: signed by {fingerprint(signer)} ({who}).", file=sys.stderr)
else:
    print("Unsigned secret: nothing proves who sent it.", file=sys.stderr)
sys.stdout.write(plaintext.decode("utf-8", errors="replace"))
if sys.stdout.isatty():
    sys.stdout.write("\n")
PY
}

# A wrong -k or -p file must fail here, before the GET spends the view.
pqc check || { [ -n "$INPUT" ] || echo "Nothing was fetched, so no view was spent." >&2; exit 2; }

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: " PASS </dev/tty; echo >&2
fi

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

if [ -z "$INPUT" ]; then
  echo "Fetching the secret from $API (this consumes a view) ..." >&2
  HTTP=$(curl -sS -o "$WORK/resp.json" -w '%{http_code}' "$API/api/v1/secrets/$TOKEN") \
    || { echo "ERROR: could not reach $API" >&2; exit 1; }
  case "$HTTP" in
    200) ;;
    404) echo "ERROR: secret not found: it never existed or was already read." >&2; exit 1 ;;
    410) echo "ERROR: the secret expired or has no views left." >&2; exit 1 ;;
    *)   echo "ERROR: HTTP $HTTP from the server." >&2; exit 1 ;;
  esac
  # The view is spent: keep the blob so a mistyped passphrase can be retried
  # with -i instead of losing the secret.
  SAVED="$(mktemp "${TMPDIR:-/tmp}/notrust-$TOKEN.XXXXXX")"
  cp "$WORK/resp.json" "$SAVED"
  INPUT="$SAVED"; FETCHED=1
fi

if ! pqc decrypt; then
  if [ -n "$SAVED" ]; then
    # Spelled with bash and a path so it runs even when the script is not
    # executable or was started as `bash decryptpqc.sh`.
    SELF="$0"; case "$SELF" in */*) ;; *) SELF="./$SELF" ;; esac
    RETRY="bash $(printf '%q -k %q' "$SELF" "$KEM_KEY")"
    [ -z "$PINNED" ] || RETRY+="$(printf ' -p %q' "$PINNED")"
    echo "The fetched secret is saved in $SAVED. Retry without spending another view:" >&2
    echo "  $RETRY -i $(printf '%q' "$SAVED")" >&2
    [ -z "${NOTRUST_PASSPHRASE:-}" ] || echo "  (the passphrase came from NOTRUST_PASSPHRASE: fix or unset it before retrying)" >&2
  fi
  exit 1
fi

# Decrypted: the saved blob has served its purpose.
if [ -n "$SAVED" ]; then
  rm -f -- "$SAVED"
  [ "$FETCHED" -eq 1 ] || echo "Deleted $SAVED: the blob this script saved is no longer needed." >&2
fi