//! `purecrypto pkeyutl` — generic asymmetric encrypt/decrypt/sign/verify.

use crate::util::{Args, SentinelLock, die, read_input, write_output, write_output_with_mode};
use purecrypto::ec::sm2::DEFAULT_ID;
use purecrypto::ec::{BoxedEcdsaPrivateKey, Sm2PrivateKey, Sm2PublicKey, Sm2Signature};
use purecrypto::hash::{Sha256, Sha384, Sha512};
use purecrypto::key::{
    self, Algorithm as KeyAlg, Hash as KeyHash, PrivateKey as KeyPriv, SigEncoding, SignParams,
};
use purecrypto::lms::{HssPrivateKey, LmsPrivateKey};
use purecrypto::rng::OsRng;
use purecrypto::rsa::BoxedRsaPrivateKey;
use purecrypto::x509::AnyPublicKey;
use purecrypto::xmss::{XmssMtPrivateKey, XmssPrivateKey};

const USAGE: &str = "\
purecrypto pkeyutl <subcommand>

  encrypt -inkey FILE [-pubin] -pkeyopt OPT [-in FILE] [-out FILE]
  decrypt -inkey FILE          -pkeyopt OPT [-in FILE] [-out FILE]
  sign    -inkey FILE          [-pkeyopt OPT] -in FILE -out FILE
  verify  -inkey FILE [-pubin] [-pkeyopt OPT] -sigfile FILE -in FILE

-pkeyopt options:
  rsa_padding_mode:oaep        OAEP encrypt/decrypt
  rsa_padding_mode:pkcs1       PKCS#1 v1.5 encrypt/decrypt / signature
  rsa_padding_mode:pss         RSA-PSS signature
  rsa_oaep_md:sha256|sha384|sha512   OAEP hash (default sha256)
  rsa_oaep_label:HEX           OAEP label (default empty)
  digest:sha256|sha384|sha512|sha1   hash for sign/verify (default sha256)

SM2 (GB/T 32918 / RFC 8998): an SM2 key auto-routes to SM2-DSA (sign/verify,
DER Ecdsa-Sig-Value) and SM2-PKE (encrypt/decrypt). Use -id STR to override the
default signer identity (1234567812345678).";

/// Returns all `-pkeyopt key:value` entries.
fn pkeyopts(args: &Args) -> Vec<(String, String)> {
    let mut out = Vec::new();
    // Args treats every occurrence of `-pkeyopt`'s value: walk argv ourselves.
    let mut iter = args.tokens_iter();
    while let Some(t) = iter.next() {
        if (t == "-pkeyopt" || t == "--pkeyopt")
            && let Some(v) = iter.next()
        {
            let (k, val) = v.split_once(':').unwrap_or((v.as_str(), ""));
            out.push((k.to_string(), val.to_string()));
        }
    }
    out
}

#[derive(Default)]
struct Opts {
    padding: Option<String>,
    oaep_md: Option<String>,
    oaep_label: Vec<u8>,
    digest: Option<String>,
}

fn parse_opts(args: &Args) -> Opts {
    let mut opts = Opts::default();
    for (k, v) in pkeyopts(args) {
        match k.as_str() {
            "rsa_padding_mode" => opts.padding = Some(v),
            "rsa_oaep_md" => opts.oaep_md = Some(v),
            "rsa_oaep_label" => {
                opts.oaep_label =
                    crate::util::from_hex(&v).unwrap_or_else(|| die("rsa_oaep_label must be hex"));
            }
            "digest" => opts.digest = Some(v),
            other => die(format!("unknown -pkeyopt key: {other}")),
        }
    }
    opts
}

// The RSA-2048 / -3072 / -4096 variant dominates the enum size, but the value
// is short-lived inside `pkeyutl` — boxing every arm would be ceremony for no
// real benefit. (See the matching pattern in `pkey.rs`.)
//
// `encrypt`/`decrypt` only operate on RSA and SM2 keys (the only schemes with a
// public-key encryption mode), so this loader recognizes just those two. The
// `sign`/`verify` paths load keys through the unified `key::PrivateKey` facade
// instead (see `load_priv_dyn`).
#[allow(clippy::large_enum_variant)]
enum PrivKey {
    Rsa(BoxedRsaPrivateKey),
    Sm2(Sm2PrivateKey),
}

fn load_priv(path: &str) -> PrivKey {
    crate::util::warn_if_world_readable_key(path);
    let raw = std::fs::read(path).unwrap_or_else(|e| die(format!("cannot read {path}: {e}")));
    let pem = core::str::from_utf8(&raw).unwrap_or_else(|_| die("key file is not UTF-8 PEM"));
    if let Ok(k) = BoxedRsaPrivateKey::from_pkcs1_pem(pem) {
        return PrivKey::Rsa(k);
    }
    // SM2 keys share the SEC1 `EC PRIVATE KEY` PEM label with ECDSA keys; the
    // SM2 parser rejects every non-SM2 curve, routing SM2 to SM2-PKE.
    if let Ok(k) = Sm2PrivateKey::from_sec1_pem(pem) {
        return PrivKey::Sm2(k);
    }
    die("unrecognized private key for encrypt/decrypt (expected RSA PKCS#1 or SM2 SEC1 PEM)");
}

/// Loads a private key as a boxed unified [`key::PrivateKey`] trait object for
/// the `sign` path. The format-specific parsers the generic PKCS#8 decoder does
/// not cover are tried first: PKCS#1 (`RSA PRIVATE KEY`), SM2 SEC1 (routed to
/// SM2-DSA rather than ECDSA-over-the-SM2-curve), and plain ECDSA SEC1; the
/// generic decoder then covers Ed25519/Ed448/ML-DSA/SLH-DSA and PKCS#8
/// RSA/ECDSA.
fn load_priv_dyn(path: &str) -> Box<dyn KeyPriv> {
    crate::util::warn_if_world_readable_key(path);
    let raw = std::fs::read(path).unwrap_or_else(|e| die(format!("cannot read {path}: {e}")));
    let pem = core::str::from_utf8(&raw).unwrap_or_else(|_| die("key file is not UTF-8 PEM"));
    if let Ok(k) = BoxedRsaPrivateKey::from_pkcs1_pem(pem) {
        return Box::new(k);
    }
    if let Ok(k) = Sm2PrivateKey::from_sec1_pem(pem) {
        return Box::new(k);
    }
    if let Ok(k) = BoxedEcdsaPrivateKey::from_sec1_pem(pem) {
        return Box::new(k);
    }
    if let Ok(k) = key::private_key_from_pkcs8_pem(pem) {
        return k;
    }
    die("unrecognized private key (expected RSA PKCS#1, EC SEC1, or PKCS#8 PEM)");
}

/// Maps the RSA `-pkeyopt digest:` choice onto a [`KeyHash`], reproducing the
/// previous per-path validation: SHA-1 is rejected for PSS and (when signing)
/// warned for PKCS#1 v1.5; unknown digests die.
fn rsa_digest(opts: &Opts, pss: bool, warn_sha1: bool) -> KeyHash {
    match opts
        .digest
        .as_deref()
        .unwrap_or("sha256")
        .to_ascii_lowercase()
        .as_str()
    {
        "sha256" => KeyHash::Sha256,
        "sha384" => KeyHash::Sha384,
        "sha512" => KeyHash::Sha512,
        "sha1" if !pss => {
            if warn_sha1 {
                eprintln!(
                    "purecrypto: warning: digest:sha1 is collision-broken for \
                     signing; use digest:sha256 or stronger"
                );
            }
            KeyHash::Sha1
        }
        other if pss => die(format!("unsupported RSA-PSS digest: {other}")),
        other => die(format!("unsupported RSA digest: {other}")),
    }
}

/// Builds the [`SignParams`] for a `sign`/`verify` call from the key's
/// algorithm and the `-pkeyopt` flags, reproducing the previous dispatch:
/// curve-implied hash and DER `Ecdsa-Sig-Value` for ECDSA, the SM2 signer id +
/// DER for SM2, opt-driven hash/padding for RSA, and scheme defaults (ignored
/// params) for EdDSA / ML-DSA / SLH-DSA. `id` is only read for SM2; `warn_sha1`
/// is set on the signing path. Used for both `sign` and `verify` (`id` is empty
/// for verify, where SM2 is handled separately).
fn build_params<'a>(
    alg: KeyAlg,
    opts: &Opts,
    id: &'a [u8],
    pss: bool,
    warn_sha1: bool,
) -> SignParams<'a> {
    let p = SignParams::new();
    match alg {
        KeyAlg::Rsa => {
            let p = p.hash(rsa_digest(opts, pss, warn_sha1));
            if pss { p } else { p.pkcs1v15() }
        }
        KeyAlg::P256 | KeyAlg::Secp256k1 => p.hash(KeyHash::Sha256).sig_encoding(SigEncoding::Der),
        KeyAlg::P384 => p.hash(KeyHash::Sha384).sig_encoding(SigEncoding::Der),
        KeyAlg::P521 => p.hash(KeyHash::Sha512).sig_encoding(SigEncoding::Der),
        KeyAlg::Sm2 => p.context(id).sig_encoding(SigEncoding::Der),
        _ => p,
    }
}

// ---------------------------------------------------------------------------
// Stateful hash-based signatures (LMS / HSS / XMSS / XMSS^MT)
//
// These keys carry a one-time-key index that advances on every signature. The
// CLI stores the raw `to_bytes()` serialization (NOT PEM). On `sign` we:
//   1. load the private key from the FILE (stdin is rejected — we must write
//      the advanced state back to a real path),
//   2. produce the signature (which advances the in-memory index),
//   3. ATOMICALLY rewrite the key file (temp + rename) with the advanced
//      `to_bytes()` BEFORE emitting the signature, and
//   4. warn loudly on stderr that the key has advanced and the old copy must
//      never be reused.
// `verify` derives the public key from the same raw key file (deriving the
// public key does NOT advance state) and checks the signature.
// ---------------------------------------------------------------------------

/// A loaded stateful signing key, tagged by scheme.
enum StatefulKey {
    Lms(LmsPrivateKey),
    Hss(HssPrivateKey),
    Xmss(XmssPrivateKey),
    XmssMt(XmssMtPrivateKey),
}

/// Attempts to parse `raw` as one of the stateful private-key serializations.
/// The four `from_bytes` parsers are mutually exclusive (distinct magic /
/// length / typecode framing), so order is immaterial.
fn parse_stateful(raw: &[u8]) -> Option<StatefulKey> {
    if let Ok(k) = LmsPrivateKey::from_bytes(raw) {
        return Some(StatefulKey::Lms(k));
    }
    if let Ok(k) = HssPrivateKey::from_bytes(raw) {
        return Some(StatefulKey::Hss(k));
    }
    if let Ok(k) = XmssPrivateKey::from_bytes(raw) {
        return Some(StatefulKey::Xmss(k));
    }
    if let Ok(k) = XmssMtPrivateKey::from_bytes(raw) {
        return Some(StatefulKey::XmssMt(k));
    }
    None
}

/// Atomically replaces `path`'s contents with `data` (write a sibling temp
/// file, fsync, rename over the original, then fsync the containing
/// directory). On Unix the temp file is created mode 0o600. Dies on any I/O
/// failure — we must NOT emit a signature if persisting the advanced key
/// failed.
fn atomic_overwrite(path: &str, data: &[u8]) {
    use std::io::Write;
    let tmp = format!("{path}.tmp.{}", std::process::id());
    {
        use std::fs::OpenOptions;
        #[cfg(unix)]
        use std::os::unix::fs::OpenOptionsExt;
        let mut opts = OpenOptions::new();
        // `create_new` refuses to clobber a pre-existing file or symlink at the
        // temp path (defense in depth). If a stale temp survives a previous
        // crashed run we remove it once and retry, then fail hard.
        opts.create_new(true).write(true);
        #[cfg(unix)]
        opts.mode(0o600);
        let mut f = match opts.open(&tmp) {
            Ok(f) => f,
            Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
                std::fs::remove_file(&tmp).unwrap_or_else(|e| {
                    die(format!("cannot remove stale temp key file {tmp}: {e}"))
                });
                opts.open(&tmp)
                    .unwrap_or_else(|e| die(format!("cannot create temp key file {tmp}: {e}")))
            }
            Err(e) => die(format!("cannot create temp key file {tmp}: {e}")),
        };
        f.write_all(data)
            .unwrap_or_else(|e| die(format!("cannot write temp key file {tmp}: {e}")));
        f.sync_all()
            .unwrap_or_else(|e| die(format!("cannot fsync temp key file {tmp}: {e}")));
    }
    std::fs::rename(&tmp, path).unwrap_or_else(|e| {
        let _ = std::fs::remove_file(&tmp);
        die(format!("cannot atomically replace key file {path}: {e}"))
    });
    // rename(2) only becomes durable once the containing directory's entry
    // reaches disk. Without an fsync of the directory, a power loss after the
    // signature is emitted can roll the key file back to the PREVIOUS one-time
    // index — exactly the OTS-reuse class this whole function exists to
    // prevent. So fsync the parent directory before returning (and before the
    // caller releases the signature).
    let dir = std::path::Path::new(path)
        .parent()
        .filter(|p| !p.as_os_str().is_empty())
        .unwrap_or_else(|| std::path::Path::new("."));
    #[cfg(unix)]
    {
        // On Unix a directory can be opened read-only and fsync'd. This must
        // succeed: dying here keeps the "no signature unless the advanced key
        // is durable" contract.
        let d = std::fs::File::open(dir).unwrap_or_else(|e| {
            die(format!(
                "cannot open directory {} to fsync key rename: {e}",
                dir.display()
            ))
        });
        d.sync_all().unwrap_or_else(|e| {
            die(format!(
                "cannot fsync directory {} after key rename: {e}",
                dir.display()
            ))
        });
    }
    #[cfg(not(unix))]
    {
        // Windows cannot fsync a directory through std (opening a directory
        // requires FILE_FLAG_BACKUP_SEMANTICS, and FlushFileBuffers on a
        // directory handle is not supported). Best-effort only; NTFS metadata
        // journaling gives the rename reasonable ordering guarantees.
        if let Ok(d) = std::fs::File::open(dir) {
            let _ = d.sync_all();
        }
    }
}

/// Stateful `sign`: load, sign (advancing the index), persist the advanced key
/// back to `key_path` atomically, warn, then return the signature.
///
/// The entire load → sign → persist sequence runs under an inter-process
/// [`SentinelLock`] (`{key_path}.lock`). Without it, two concurrent
/// invocations would both read index `N`, both emit signatures under the
/// same one-time key, and the on-disk state would still end self-consistent
/// — the catastrophic OTS reuse would be invisible after the fact.
fn stateful_sign(key_path: &str, msg: &[u8]) -> Vec<u8> {
    let _lock = SentinelLock::acquire(
        std::path::PathBuf::from(format!("{key_path}.lock")),
        "`purecrypto pkeyutl sign`",
    );
    crate::util::warn_if_world_readable_key(key_path);
    let raw =
        std::fs::read(key_path).unwrap_or_else(|e| die(format!("cannot read {key_path}: {e}")));
    let key = parse_stateful(&raw)
        .unwrap_or_else(|| die("not a recognized LMS/HSS/XMSS/XMSS^MT private key"));

    // Sign, then capture the ADVANCED serialization to persist before we hand
    // the signature back to the caller.
    let (sig, advanced) = match key {
        // UFCS on the concrete type to reach the inherent stateful `sign`
        // (`&mut self`): the in-scope `key::PrivateKey::sign` (`&self`) would
        // otherwise win method resolution at the shared-ref autoref step.
        StatefulKey::Lms(mut k) => {
            let s = LmsPrivateKey::sign(&mut k, &mut OsRng, msg)
                .unwrap_or_else(|e| die(format!("LMS sign failed: {e:?}")));
            (s, k.to_bytes())
        }
        StatefulKey::Hss(mut k) => {
            let s = HssPrivateKey::sign(&mut k, &mut OsRng, msg)
                .unwrap_or_else(|e| die(format!("HSS sign failed: {e:?}")));
            (s, k.to_bytes())
        }
        StatefulKey::Xmss(mut k) => {
            let s = XmssPrivateKey::sign(&mut k, msg)
                .unwrap_or_else(|e| die(format!("XMSS sign failed: {e:?}")));
            (s, k.to_bytes())
        }
        StatefulKey::XmssMt(mut k) => {
            let s = XmssMtPrivateKey::sign(&mut k, msg)
                .unwrap_or_else(|e| die(format!("XMSS^MT sign failed: {e:?}")));
            (s, k.to_bytes())
        }
    };

    // Persist the advanced state BEFORE the signature is emitted. If this fails,
    // `atomic_overwrite` exits non-zero and the signature is never written.
    atomic_overwrite(key_path, &advanced);
    eprintln!(
        "purecrypto: warning: stateful key {key_path} has ADVANCED to its next \
         one-time index and been rewritten in place. The previous key state is \
         GONE — never restore or reuse an older copy of this file, or signatures \
         will reuse a one-time key (catastrophic)."
    );
    sig
}

/// Stateful `verify`: derive the public key from the raw private-key file and
/// check `sig` over `msg`. Returns `true` if the file is a stateful key (and
/// sets `*ok`); `false` if it is not a stateful key at all.
fn stateful_verify(raw: &[u8], msg: &[u8], sig: &[u8]) -> Option<bool> {
    let key = parse_stateful(raw)?;
    let ok = match key {
        StatefulKey::Lms(k) => k.public_key().verify(msg, sig),
        StatefulKey::Hss(k) => k.public_key().verify(msg, sig),
        StatefulKey::Xmss(k) => k.public_key().verify(msg, sig),
        StatefulKey::XmssMt(k) => k.public_key().verify(msg, sig),
    };
    Some(ok)
}

/// The SM2 signer identity `Z_A` (GB/T 32918.2). Defaults to the standard
/// `1234567812345678`; `-id STR` overrides it with raw UTF-8 bytes.
fn sm2_id(args: &Args) -> Vec<u8> {
    match args.value("-id").or_else(|| args.value("--id")) {
        Some(s) => s.as_bytes().to_vec(),
        None => DEFAULT_ID.to_vec(),
    }
}

fn load_spki(path: &str) -> AnyPublicKey {
    let raw = std::fs::read(path).unwrap_or_else(|e| die(format!("cannot read {path}: {e}")));
    let pem = core::str::from_utf8(&raw).unwrap_or_else(|_| die("pubkey is not UTF-8 PEM"));
    AnyPublicKey::from_spki_pem(pem).unwrap_or_else(|e| die(format!("cannot parse SPKI PEM: {e}")))
}

/// Attempts to load an SM2 public key from `path`: a `PUBLIC KEY` SPKI PEM
/// (preferred for `-pubin`), or the public half of an SM2 SEC1 private key.
/// Returns `None` if the file is not an SM2 key.
fn try_load_sm2_public(path: &str) -> Option<Sm2PublicKey> {
    let raw = std::fs::read(path).ok()?;
    let pem = core::str::from_utf8(&raw).ok()?;
    if let Ok(der) = purecrypto::der::pem_decode(pem, "PUBLIC KEY")
        && let Ok(pk) = Sm2PublicKey::from_spki_der(&der)
    {
        return Some(pk);
    }
    Sm2PrivateKey::from_sec1_pem(pem)
        .ok()
        .map(|sk| sk.public_key())
}

/// Emits the stderr warning for PKCS#1 v1.5 encryption padding (the
/// OpenSSL-compatible default for `encrypt`/`decrypt`). The padding itself is
/// legacy, and the decrypt failure modes form a Bleichenbacher oracle the
/// moment the command is driven by attacker-supplied ciphertexts, so the
/// decrypt variant spells that risk out.
fn warn_pkcs1_padding(decrypt: bool) {
    if decrypt {
        eprintln!(
            "purecrypto: warning: rsa_padding_mode:pkcs1 (the default) decryption is a \
             Bleichenbacher padding oracle; do not expose this command to untrusted \
             ciphertexts in a loop. Prefer -pkeyopt rsa_padding_mode:oaep"
        );
    } else {
        eprintln!(
            "purecrypto: warning: rsa_padding_mode:pkcs1 (the default) is legacy; \
             prefer -pkeyopt rsa_padding_mode:oaep"
        );
    }
}

fn run_encrypt(args: Args) {
    let in_path = args.value("-in").or_else(|| args.value("--in"));
    let pt = read_input(in_path);
    let opts = parse_opts(&args);
    let padding = opts.padding.as_deref().unwrap_or("pkcs1");
    let inkey = args
        .value("-inkey")
        .or_else(|| args.value("--inkey"))
        .unwrap_or_else(|| die("missing -inkey"));

    // SM2 hybrid PKE (GB/T 32918.4 / RFC 8998): detect an SM2 key and route
    // before the RSA-only padding logic.
    if let Some(pk) = try_load_sm2_public(inkey) {
        let ct = pk
            .encrypt(&pt, &mut OsRng)
            .unwrap_or_else(|e| die(format!("SM2 encrypt failed: {e:?}")));
        let out_path = args.value("-out").or_else(|| args.value("--out"));
        write_output(out_path, &ct);
        return;
    }

    // From here on the key is RSA, so the padding mode actually applies.
    if padding == "pkcs1" {
        warn_pkcs1_padding(/* decrypt = */ false);
    }
    let ct = if args.flag("-pubin") || args.flag("--pubin") {
        let any = load_spki(inkey);
        let rsa = match any {
            AnyPublicKey::Rsa(k) => k,
            _ => die("RSA encrypt requires an RSA SPKI"),
        };
        match padding {
            "oaep" => {
                let md = opts.oaep_md.as_deref().unwrap_or("sha256");
                match md.to_ascii_lowercase().as_str() {
                    "sha256" => rsa
                        .encrypt_oaep::<Sha256, _>(&pt, &opts.oaep_label, &mut OsRng)
                        .unwrap_or_else(|e| die(format!("OAEP encrypt failed: {e}"))),
                    "sha384" => rsa
                        .encrypt_oaep::<Sha384, _>(&pt, &opts.oaep_label, &mut OsRng)
                        .unwrap_or_else(|e| die(format!("OAEP encrypt failed: {e}"))),
                    "sha512" => rsa
                        .encrypt_oaep::<Sha512, _>(&pt, &opts.oaep_label, &mut OsRng)
                        .unwrap_or_else(|e| die(format!("OAEP encrypt failed: {e}"))),
                    _ => die(format!("unsupported rsa_oaep_md: {md}")),
                }
            }
            "pkcs1" => rsa
                .encrypt_pkcs1v15(&pt, &mut OsRng)
                .unwrap_or_else(|e| die(format!("PKCS1 encrypt failed: {e}"))),
            other => die(format!("unsupported rsa_padding_mode for encrypt: {other}")),
        }
    } else {
        let key = load_priv(inkey);
        let rsa = match key {
            PrivKey::Rsa(k) => k,
            _ => die("RSA encrypt requires an RSA key"),
        };
        let pub_k = rsa.public_key();
        match padding {
            "oaep" => {
                let md = opts.oaep_md.as_deref().unwrap_or("sha256");
                match md.to_ascii_lowercase().as_str() {
                    "sha256" => pub_k
                        .encrypt_oaep::<Sha256, _>(&pt, &opts.oaep_label, &mut OsRng)
                        .unwrap_or_else(|e| die(format!("OAEP encrypt failed: {e}"))),
                    "sha384" => pub_k
                        .encrypt_oaep::<Sha384, _>(&pt, &opts.oaep_label, &mut OsRng)
                        .unwrap_or_else(|e| die(format!("OAEP encrypt failed: {e}"))),
                    "sha512" => pub_k
                        .encrypt_oaep::<Sha512, _>(&pt, &opts.oaep_label, &mut OsRng)
                        .unwrap_or_else(|e| die(format!("OAEP encrypt failed: {e}"))),
                    _ => die(format!("unsupported rsa_oaep_md: {md}")),
                }
            }
            "pkcs1" => pub_k
                .encrypt_pkcs1v15(&pt, &mut OsRng)
                .unwrap_or_else(|e| die(format!("PKCS1 encrypt failed: {e}"))),
            other => die(format!("unsupported rsa_padding_mode: {other}")),
        }
    };
    let out_path = args.value("-out").or_else(|| args.value("--out"));
    write_output(out_path, &ct);
}

fn run_decrypt(args: Args) {
    let in_path = args.value("-in").or_else(|| args.value("--in"));
    let ct = read_input(in_path);
    let opts = parse_opts(&args);
    let padding = opts.padding.as_deref().unwrap_or("pkcs1");
    let inkey = args
        .value("-inkey")
        .or_else(|| args.value("--inkey"))
        .unwrap_or_else(|| die("missing -inkey"));
    let key = load_priv(inkey);
    // SM2 hybrid PKE decrypt routes to the SM2 private key.
    if let PrivKey::Sm2(sk) = &key {
        // Fixed failure string (see the RSA paths below): the C3 tag check vs
        // malformed-C1 vs length causes must stay indistinguishable.
        let pt = sk.decrypt(&ct).unwrap_or_else(|_| die("decrypt failed"));
        let out_path = args.value("-out").or_else(|| args.value("--out"));
        // Recovered plaintext is typically a wrapped key: write it like the
        // kem/kex secrets (0600, create_new, refuse a TTY) rather than a
        // world-readable 0644 file (mirrors `enc -d` for AES-KW/KWP).
        write_output_with_mode(out_path, &pt, /* private = */ true);
        return;
    }
    let rsa = match key {
        PrivKey::Rsa(k) => k,
        _ => die("RSA decrypt requires an RSA key"),
    };
    if padding == "pkcs1" {
        warn_pkcs1_padding(/* decrypt = */ true);
    }
    // All decryption failures collapse into one fixed string with no
    // underlying detail: distinguishing bad-padding from bad-length (etc.) is
    // exactly the Bleichenbacher / Manger oracle, and anyone wrapping this
    // one-shot CLI in a network-facing loop would hand it to the attacker.
    let pt = match padding {
        "oaep" => {
            let md = opts.oaep_md.as_deref().unwrap_or("sha256");
            match md.to_ascii_lowercase().as_str() {
                "sha256" => rsa
                    .decrypt_oaep::<Sha256>(&ct, &opts.oaep_label)
                    .unwrap_or_else(|_| die("decrypt failed")),
                "sha384" => rsa
                    .decrypt_oaep::<Sha384>(&ct, &opts.oaep_label)
                    .unwrap_or_else(|_| die("decrypt failed")),
                "sha512" => rsa
                    .decrypt_oaep::<Sha512>(&ct, &opts.oaep_label)
                    .unwrap_or_else(|_| die("decrypt failed")),
                _ => die(format!("unsupported rsa_oaep_md: {md}")),
            }
        }
        "pkcs1" => rsa
            .decrypt_pkcs1v15(&ct)
            .unwrap_or_else(|_| die("decrypt failed")),
        other => die(format!("unsupported rsa_padding_mode: {other}")),
    };
    let out_path = args.value("-out").or_else(|| args.value("--out"));
    // Same as the SM2 path above: the recovered plaintext is key-grade.
    write_output_with_mode(out_path, &pt, /* private = */ true);
}

fn run_sign(args: Args) {
    let in_path = args
        .value("-in")
        .or_else(|| args.value("--in"))
        .unwrap_or_else(|| die("missing -in FILE"));
    let msg = std::fs::read(in_path).unwrap_or_else(|e| die(format!("cannot read {in_path}: {e}")));
    let opts = parse_opts(&args);
    let inkey = args
        .value("-inkey")
        .or_else(|| args.value("--inkey"))
        .unwrap_or_else(|| die("missing -inkey"));
    if inkey == "-" {
        die("stateful keys (LMS/XMSS) and signing require -inkey FILE, not stdin");
    }

    // Stateful hash-based signature? Detect from the raw key bytes and route to
    // the special path that rewrites the advanced key file before emitting the
    // signature. (Reading the file twice is fine; the authoritative parse is in
    // `stateful_sign`.)
    if let Ok(raw) = std::fs::read(inkey)
        && parse_stateful(&raw).is_some()
    {
        let sig = stateful_sign(inkey, &msg);
        let out_path = args
            .value("-out")
            .or_else(|| args.value("--out"))
            .unwrap_or_else(|| die("missing -out FILE"));
        write_output(Some(out_path), &sig);
        return;
    }

    // Dispatch through the unified `key::PrivateKey` facade: the per-algorithm
    // hash / padding / encoding choices that used to be a big `match` are now
    // encoded in `SignParams`, keyed off the loaded key's `algorithm()`.
    let key = load_priv_dyn(inkey);
    let pss = matches!(opts.padding.as_deref(), Some("pss"));
    let id = sm2_id(&args);
    let params = build_params(key.algorithm(), &opts, &id, pss, /* warn_sha1 */ true);
    let sig = key
        .sign(&msg, &params, &mut OsRng)
        .unwrap_or_else(|e| die(format!("sign failed: {e}")));
    let out_path = args
        .value("-out")
        .or_else(|| args.value("--out"))
        .unwrap_or_else(|| die("missing -out FILE"));
    write_output(Some(out_path), &sig);
}

fn run_verify(args: Args) {
    let in_path = args
        .value("-in")
        .or_else(|| args.value("--in"))
        .unwrap_or_else(|| die("missing -in FILE"));
    let msg = std::fs::read(in_path).unwrap_or_else(|e| die(format!("cannot read {in_path}: {e}")));
    let sig_path = args
        .value("-sigfile")
        .or_else(|| args.value("--sigfile"))
        .unwrap_or_else(|| die("missing -sigfile FILE"));
    let sig =
        std::fs::read(sig_path).unwrap_or_else(|e| die(format!("cannot read {sig_path}: {e}")));
    let inkey = args
        .value("-inkey")
        .or_else(|| args.value("--inkey"))
        .unwrap_or_else(|| die("missing -inkey"));
    let opts = parse_opts(&args);
    let pss = matches!(opts.padding.as_deref(), Some("pss"));

    // `-inkey` may be a raw LMS/HSS/XMSS or SM2/SEC1 PRIVATE key file
    // (the public key is derived from it). Warn if it is group/world-
    // readable, matching the other secret-key readers.
    crate::util::warn_if_world_readable_key(inkey);
    let raw_key = std::fs::read(inkey).unwrap_or_else(|e| die(format!("cannot read {inkey}: {e}")));

    // Stateful hash-based signatures: the public key is derived from the raw
    // private-key file (deriving the public key does NOT advance state).
    if let Some(ok) = stateful_verify(&raw_key, &msg, &sig) {
        report_verify(ok);
        return;
    }

    // SM2 verification: accept either a `PUBLIC KEY` SPKI PEM (the SM2 named
    // curve) or an SM2 SEC1 private key (deriving its public key — which does
    // not expose the secret). The generic `AnyPublicKey` parser does not handle
    // SM2, so detect it here and route to SM2-DSA (`Ecdsa-Sig-Value` signature).
    if let Ok(pem) = core::str::from_utf8(&raw_key) {
        let pk = purecrypto::der::pem_decode(pem, "PUBLIC KEY")
            .ok()
            .and_then(|der| Sm2PublicKey::from_spki_der(&der).ok())
            .or_else(|| {
                Sm2PrivateKey::from_sec1_pem(pem)
                    .ok()
                    .map(|sk| sk.public_key())
            });
        if let Some(pk) = pk {
            let id = sm2_id(&args);
            let ok = Sm2Signature::from_der(&sig)
                .map(|parsed| pk.verify(&msg, &parsed, &id).is_ok())
                .unwrap_or(false);
            report_verify(ok);
            return;
        }
    }

    // Unified verify: decode the SPKI public key and check through the
    // `key::PublicKey` facade. The hash (curve-implied for ECDSA) and DER
    // signature encoding are selected from the key's `algorithm()`, reproducing
    // the previous per-type dispatch. SM2 and the stateful schemes are handled
    // above; everything else flows through here. A malformed signature (wrong
    // length, bad DER) verifies as a plain failure.
    let pem = core::str::from_utf8(&raw_key).unwrap_or_else(|_| die("pubkey is not UTF-8 PEM"));
    let pub_key = key::public_key_from_spki_pem(pem)
        .unwrap_or_else(|e| die(format!("cannot parse SPKI: {e}")));
    let params = build_params(
        pub_key.algorithm(),
        &opts,
        &[],
        pss,
        /* warn_sha1 */ false,
    );
    let ok = pub_key.verify(&msg, &sig, &params).is_ok();
    report_verify(ok);
}

/// Prints the OpenSSL-style verification result and exits non-zero on failure.
fn report_verify(ok: bool) {
    if ok {
        println!("Signature verified");
    } else {
        println!("Signature verification failure");
        std::process::exit(1);
    }
}

pub(crate) fn run(args: Args) {
    let pos = args.positionals(&[
        "-inkey",
        "--inkey",
        "-in",
        "--in",
        "-out",
        "--out",
        "-sigfile",
        "--sigfile",
        "-pkeyopt",
        "--pkeyopt",
        "-id",
        "--id",
    ]);
    let sub = pos.first().copied().unwrap_or("");
    match sub {
        "encrypt" => run_encrypt(args),
        "decrypt" => run_decrypt(args),
        "sign" => run_sign(args),
        "verify" => run_verify(args),
        "" => die(USAGE),
        other => die(format!("unknown pkeyutl subcommand '{other}'\n\n{USAGE}")),
    }
}

#[cfg(test)]
mod tests {
    //! Regression coverage for the stateful-sign inter-process lock: two
    //! concurrent `pkeyutl sign` invocations against the same LMS key file
    //! must never both sign with the same one-time-key index. The `run_*`
    //! entry points shell out to `die()` (process exit), so we exercise
    //! `stateful_sign` directly against a scratch key file.
    use super::*;
    use purecrypto::lms::{LmotsType, LmsType};
    use std::collections::HashSet;
    use std::path::PathBuf;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicU64, Ordering};

    /// Counter used to mint unique scratch directory names (no `tempfile`
    /// dev-dependency — Cargo.toml is off-limits; same pattern as
    /// `ca::tests`).
    static SCRATCH_COUNTER: AtomicU64 = AtomicU64::new(0);

    /// Auto-cleaning scratch directory. Drop removes the tree on a best-
    /// effort basis so a panicking test does not leak permanent files.
    struct ScratchDir(PathBuf);

    impl ScratchDir {
        fn new(tag: &str) -> Self {
            let n = SCRATCH_COUNTER.fetch_add(1, Ordering::Relaxed);
            let pid = std::process::id();
            let nanos = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map(|d| d.subsec_nanos())
                .unwrap_or(0);
            let path =
                std::env::temp_dir().join(format!("purecrypto-pkeyutl-{tag}-{pid}-{nanos}-{n}"));
            std::fs::create_dir_all(&path).expect("mkdir scratch");
            ScratchDir(path)
        }
    }

    impl Drop for ScratchDir {
        fn drop(&mut self) {
            let _ = std::fs::remove_dir_all(&self.0);
        }
    }

    /// Threads racing `stateful_sign` against the same LMS key file must
    /// each consume a DISTINCT one-time-key index `q` (the first 4 bytes of
    /// an RFC 8554 LMS signature), and the on-disk key must end advanced by
    /// exactly the number of signatures issued. Without the `{key}.lock`
    /// sentinel held across load → sign → persist, two racers both read
    /// index N and emit two signatures under the same OTS key.
    #[test]
    fn concurrent_stateful_sign_never_reuses_an_ots_index() {
        let td = ScratchDir::new("statefulsign");
        let key_path = td.0.join("lms.key");
        let key_path_str = key_path.to_str().expect("utf-8 path").to_string();

        let key = LmsPrivateKey::generate(LmsType::Sha256M32H5, LmotsType::Sha256N32W4, &mut OsRng);
        std::fs::write(&key_path, key.to_bytes()).expect("seed key file");

        const THREADS: u32 = 4;
        const PER_THREAD: u32 = 3;
        let barrier = Arc::new(std::sync::Barrier::new(THREADS as usize));
        let mut handles = Vec::new();
        for _ in 0..THREADS {
            let barrier = Arc::clone(&barrier);
            let path = key_path_str.clone();
            handles.push(std::thread::spawn(move || {
                barrier.wait();
                let mut qs = Vec::with_capacity(PER_THREAD as usize);
                for _ in 0..PER_THREAD {
                    let sig = stateful_sign(&path, b"ots reuse regression");
                    // LMS signature = u32(q) || LM-OTS sig || u32(type) || path.
                    let q = u32::from_be_bytes(sig[..4].try_into().unwrap());
                    qs.push(q);
                }
                qs
            }));
        }
        let mut all: Vec<u32> = Vec::new();
        for h in handles {
            all.extend(h.join().expect("signer thread"));
        }

        // Every signature consumed a distinct OTS index, with no gaps.
        let set: HashSet<u32> = all.iter().copied().collect();
        assert_eq!(
            set.len(),
            all.len(),
            "one-time-key index reused under concurrency: {all:?}"
        );
        let expected: HashSet<u32> = (0..THREADS * PER_THREAD).collect();
        assert_eq!(set, expected, "unexpected index set: {all:?}");

        // The persisted key advanced by exactly the number of signatures
        // (H5 = 32 leaves total).
        let advanced = std::fs::read(&key_path).expect("read advanced key");
        let on_disk = LmsPrivateKey::from_bytes(&advanced).expect("parse advanced key");
        assert_eq!(on_disk.remaining(), 32 - (THREADS * PER_THREAD) as u64);

        // The sentinel lock was released.
        assert!(!PathBuf::from(format!("{key_path_str}.lock")).exists());
    }
}
