challenge 0.1.0

use anyhow::{Result, bail};
use clap::{Parser, Subcommand};
use rand::RngExt;
use serde_json::Value;

use challenge::{
    CreateConfig, SolveConfig, create, parse_data_args, parse_server_payload, print_json,
    read_input, read_json, solve, unix_now, verify_pair, verify_register_payload,
    verify_server_payload,
};

#[derive(Debug, Parser)]
#[command(name = "challenge")]
#[command(about = "A lightweight CLI for ALTCHA Proof-of-Work v2 challenges")]
#[command(version)]
struct Cli {
    /// Emit compact JSON instead of pretty-printed JSON.
    #[arg(long, global = true)]
    compact: bool,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Debug, Subcommand)]
enum Commands {
    /// Create a new ALTCHA challenge and print it as JSON.
    Create(CreateArgs),

    /// Solve an ALTCHA challenge JSON document.
    Solve(SolveArgs),

    /// Verify a challenge + solution pair.
    Verify(VerifyArgs),

    /// Verify a `/register` JSON body containing altcha.challenge + altcha.solution.
    VerifyPayload(VerifyPayloadArgs),

    /// Verify an ALTCHA Sentinel server-signature payload.
    VerifyServer(VerifyServerArgs),
}

#[derive(Debug, Parser)]
struct CreateArgs {
    /// KDF algorithm string, e.g. PBKDF2/SHA-256, SHA-256, SCRYPT, ARGON2ID.
    #[arg(long, default_value = "PBKDF2/SHA-256")]
    algorithm: String,

    /// Algorithm cost: iteration count, time cost, or scrypt N.
    #[arg(long, default_value_t = 5_000)]
    cost: u32,

    /// Fixed counter for deterministic mode.
    #[arg(long)]
    counter: Option<u32>,

    /// Generate a random deterministic counter using --counter-min..--counter-max.
    #[arg(long)]
    random_counter: bool,

    /// Minimum counter used with --random-counter.
    #[arg(long, default_value_t = 5_000)]
    counter_min: u32,

    /// Maximum counter used with --random-counter.
    #[arg(long, default_value_t = 10_000)]
    counter_max: u32,

    /// Expire the challenge N seconds from now.
    #[arg(long)]
    expires_in: Option<u64>,

    /// Absolute Unix timestamp in seconds when the challenge expires.
    #[arg(long)]
    expires_at: Option<u64>,

    /// Secret used to sign challenge parameters.
    #[arg(long, env = "ALTCHA_HMAC_SECRET")]
    hmac_secret: Option<String>,

    /// Allow creating an unsigned challenge. Not recommended for production.
    #[arg(long)]
    unsigned: bool,

    /// Secret used to sign the derived key in deterministic mode.
    #[arg(long, env = "ALTCHA_HMAC_KEY_SECRET")]
    hmac_key_secret: Option<String>,

    /// Output key length in bytes.
    #[arg(long)]
    key_length: Option<usize>,

    /// Required hex prefix the derived key must start with.
    #[arg(long)]
    key_prefix: Option<String>,

    /// Bytes used as prefix in deterministic mode.
    #[arg(long)]
    key_prefix_length: Option<usize>,

    /// Memory cost in KiB for Argon2id, or r for scrypt.
    #[arg(long)]
    memory_cost: Option<u32>,

    /// Parallelism factor for Argon2id or scrypt.
    #[arg(long)]
    parallelism: Option<u32>,

    /// Metadata to embed and sign, as key=value. Values are parsed as JSON when possible.
    #[arg(long = "data")]
    data: Vec<String>,
}

#[derive(Debug, Parser)]
struct SolveArgs {
    /// Path to challenge JSON, or '-' for stdin.
    #[arg(long, short = 'c', default_value = "-")]
    challenge: String,

    /// First counter value to try.
    #[arg(long, default_value_t = 0)]
    counter_start: u32,

    /// Counter increment per attempt.
    #[arg(long, default_value_t = 1)]
    counter_step: u32,

    /// Maximum solve time in milliseconds.
    #[arg(long, default_value_t = 90_000)]
    timeout_ms: u64,
}

#[derive(Debug, Parser)]
struct VerifyArgs {
    /// Path to the original challenge JSON, or '-' for stdin.
    #[arg(long, short = 'c')]
    challenge: String,

    /// Path to the submitted solution JSON, or '-' for stdin.
    #[arg(long, short = 's')]
    solution: String,

    /// Secret used when the challenge was created.
    #[arg(long, env = "ALTCHA_HMAC_SECRET")]
    secret: String,

    /// Secret used for deterministic-mode key-signature verification.
    #[arg(long, env = "ALTCHA_HMAC_KEY_SECRET")]
    key_secret: Option<String>,

    /// Return a non-zero exit code when verification returns verified=false.
    #[arg(long)]
    fail_on_invalid: bool,
}

#[derive(Debug, Parser)]
struct VerifyPayloadArgs {
    /// Path to `/register` JSON body, or '-' for stdin.
    #[arg(long, short = 'p', default_value = "-")]
    payload: String,

    /// Secret used when the challenge was created.
    #[arg(long, env = "ALTCHA_HMAC_SECRET")]
    secret: String,

    /// Secret used for deterministic-mode key-signature verification.
    #[arg(long, env = "ALTCHA_HMAC_KEY_SECRET")]
    key_secret: Option<String>,

    /// Return a non-zero exit code when verification returns verified=false.
    #[arg(long)]
    fail_on_invalid: bool,
}

#[derive(Debug, Parser)]
struct VerifyServerArgs {
    /// Path to ServerSignaturePayload JSON, or '-' for stdin.
    #[arg(long, short = 'p', default_value = "-")]
    payload: String,

    /// Treat input as the base64-encoded JSON form field from ALTCHA Sentinel.
    #[arg(long)]
    base64: bool,

    /// HMAC secret shared with ALTCHA Sentinel.
    #[arg(long, env = "ALTCHA_HMAC_SECRET")]
    secret: String,

    /// Return a non-zero exit code when verification returns verified=false.
    #[arg(long)]
    fail_on_invalid: bool,
}

fn main() -> Result<()> {
    let cli = Cli::parse();

    match cli.command {
        Commands::Create(args) => {
            let challenge = create(build_create_config(args)?)?;
            print_json(&challenge, cli.compact)?;
        }
        Commands::Solve(args) => {
            let challenge = read_json(&args.challenge)?;
            let solution = solve(
                &challenge,
                SolveConfig {
                    counter_start: args.counter_start,
                    counter_step: args.counter_step,
                    timeout_ms: args.timeout_ms,
                },
            )?;
            print_json(&solution, cli.compact)?;
        }
        Commands::Verify(args) => {
            validate_stdin_pair(&args.challenge, &args.solution)?;
            let challenge = read_json(&args.challenge)?;
            let solution = read_json(&args.solution)?;
            let output = verify_pair(&challenge, &solution, &args.secret, args.key_secret)?;
            let verified = output.verified;
            print_json(&output, cli.compact)?;
            maybe_fail_on_invalid(verified, args.fail_on_invalid)?;
        }
        Commands::VerifyPayload(args) => {
            let payload: Value = read_json(&args.payload)?;
            let output = verify_register_payload(&payload, &args.secret, args.key_secret)?;
            let verified = output.verified;
            print_json(&output, cli.compact)?;
            maybe_fail_on_invalid(verified, args.fail_on_invalid)?;
        }
        Commands::VerifyServer(args) => {
            let raw = read_input(&args.payload)?;
            let payload = parse_server_payload(&raw, args.base64)?;
            let output = verify_server_payload(&payload, &args.secret)?;
            let verified = output.verified;
            print_json(&output, cli.compact)?;
            maybe_fail_on_invalid(verified, args.fail_on_invalid)?;
        }
    }

    Ok(())
}

fn build_create_config(args: CreateArgs) -> Result<CreateConfig> {
    if args.counter.is_some() && args.random_counter {
        bail!("use either --counter or --random-counter, not both");
    }

    if args.expires_in.is_some() && args.expires_at.is_some() {
        bail!("use either --expires-in or --expires-at, not both");
    }

    if args.unsigned && args.hmac_secret.is_some() {
        bail!("use either --hmac-secret or --unsigned, not both");
    }

    if !args.unsigned && args.hmac_secret.is_none() {
        bail!(
            "--hmac-secret is required for signed production challenges; use --unsigned only for local demos"
        );
    }

    let deterministic = args.counter.is_some() || args.random_counter;
    if deterministic && args.hmac_key_secret.is_none() {
        bail!("--hmac-key-secret is required with --counter or --random-counter");
    }

    let counter = if args.random_counter {
        if args.counter_min > args.counter_max {
            bail!("--counter-min must be <= --counter-max");
        }
        Some(rand::rng().random_range(args.counter_min..=args.counter_max))
    } else {
        args.counter
    };

    let expires_at = match (args.expires_in, args.expires_at) {
        (Some(seconds), None) => Some(unix_now()? + seconds),
        (None, Some(timestamp)) => Some(timestamp),
        (None, None) => None,
        (Some(_), Some(_)) => unreachable!("validated above"),
    };

    Ok(CreateConfig {
        algorithm: args.algorithm,
        cost: args.cost,
        counter,
        data: parse_data_args(&args.data)?,
        expires_at,
        hmac_signature_secret: args.hmac_secret,
        hmac_key_signature_secret: args.hmac_key_secret,
        key_length: args.key_length,
        key_prefix: args.key_prefix,
        key_prefix_length: args.key_prefix_length,
        memory_cost: args.memory_cost,
        parallelism: args.parallelism,
    })
}

fn validate_stdin_pair(first: &str, second: &str) -> Result<()> {
    if first == "-" && second == "-" {
        bail!("only one input can be read from stdin; provide a file for either input");
    }
    Ok(())
}

fn maybe_fail_on_invalid(verified: bool, fail_on_invalid: bool) -> Result<()> {
    if fail_on_invalid && !verified {
        bail!("ALTCHA verification failed");
    }
    Ok(())
}