bee-check 0.4.0

//! `bee-check` — retrievability checker for Swarm references.
//!
//! Multi-vantage probe of `GET /stewardship/{ref}` across one or more
//! Bee nodes, with optional per-chunk drill-down and re-seed.

use anyhow::{Context, Result};
use clap::{Parser, ValueEnum};

use bee_check::{
    DEFAULT_GATEWAY, OutputFormat, ParsedInput, ReseedRequest, annotate_target_overlay,
    check_gateways, check_multi_vantage, check_stamp, drill_down, merge_gateways, parse_input,
    render_report, render_stamp_status, reseed, resolve_feed,
};

#[derive(Parser, Debug)]
#[command(
    name = "bee-check",
    version,
    about = "Retrievability checker for Swarm references",
    long_about = "Probe one or more Bee nodes to determine whether a Swarm \
                  reference is retrievable from the network. Supports \
                  per-chunk drill-down and one-shot re-seed via stewardship."
)]
struct Cli {
    /// Swarm reference (64- or 128-hex) or feed reference
    /// `feed:OWNER:TOPIC` (40-hex owner, 64-hex topic). Feed inputs are
    /// resolved via `GET /feeds/{owner}/{topic}` on the first --bee
    /// before probing.
    #[arg(value_name = "INPUT")]
    reference: String,

    /// Bee API URL(s) to probe. Repeat for multi-vantage. Defaults to
    /// $BEE_API_URL or http://localhost:1633.
    #[arg(short = 'b', long = "bee", value_name = "URL")]
    bee: Vec<String>,

    /// Public gateway URL(s) to HEAD-probe via `{gw}/bzz/{ref}/`.
    /// Repeat for multiple. Default: api.gateway.ethswarm.org unless
    /// --no-gateway is set.
    #[arg(long = "gateway", value_name = "URL", conflicts_with = "no_gateway")]
    gateway: Vec<String>,

    /// Skip public-gateway probing entirely.
    #[arg(long)]
    no_gateway: bool,

    /// Compute and display the proximity order between each vantage
    /// and this target overlay (hex). Vantages are re-sorted closest-
    /// first. Useful for "from a node near neighborhood X, is this
    /// retrievable?" questions.
    #[arg(long, value_name = "HEX")]
    target_overlay: Option<String>,

    /// Walk the manifest and probe each chunk per vantage.
    #[arg(long)]
    per_chunk: bool,

    /// After probing, re-upload the reference via PUT /stewardship/{ref}.
    /// Requires --stamp.
    #[arg(long)]
    reseed: bool,

    /// Postage batch ID for re-seed.
    #[arg(long, value_name = "ID", requires = "reseed")]
    stamp: Option<String>,

    /// Per-call timeout in seconds.
    #[arg(long, default_value_t = 60)]
    timeout: u64,

    /// Max concurrent chunk probes during drill-down.
    #[arg(long, default_value_t = 8)]
    concurrency: usize,

    /// Output format.
    #[arg(long, value_enum, default_value_t = OutputKind::Text)]
    output: OutputKind,
}

#[derive(Copy, Clone, Debug, ValueEnum)]
enum OutputKind {
    Text,
    Json,
}

impl From<OutputKind> for OutputFormat {
    fn from(k: OutputKind) -> Self {
        match k {
            OutputKind::Text => OutputFormat::Text,
            OutputKind::Json => OutputFormat::Json,
        }
    }
}

fn default_bees() -> Vec<String> {
    std::env::var("BEE_API_URL")
        .ok()
        .into_iter()
        .chain(std::iter::once("http://localhost:1633".to_string()))
        .take(1)
        .collect()
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();

    let bees: Vec<String> = if cli.bee.is_empty() {
        default_bees()
    } else {
        cli.bee.clone()
    };

    let timeout = std::time::Duration::from_secs(cli.timeout);

    // Resolve the input. A `feed:OWNER:TOPIC` is turned into a bare
    // reference via the first vantage; everything else is taken as a
    // reference verbatim.
    let (reference, resolution) = match parse_input(&cli.reference) {
        ParsedInput::Reference(r) => (r, None),
        ParsedInput::Feed { owner, topic } => {
            let first_bee = bees
                .first()
                .context("feed resolution requires at least one --bee URL")?;
            let (r, res) = resolve_feed(first_bee, &owner, &topic, timeout)
                .await
                .context("feed resolution failed")?;
            eprintln!("resolved feed -> {r}");
            (r, Some(res))
        }
    };

    let gateways: Vec<String> = if cli.no_gateway {
        Vec::new()
    } else if cli.gateway.is_empty() {
        vec![DEFAULT_GATEWAY.to_string()]
    } else {
        cli.gateway.clone()
    };

    let (vantage_report, gateway_results) = tokio::join!(
        check_multi_vantage(&reference, &bees, timeout),
        check_gateways(&reference, &gateways, timeout),
    );
    let mut report = vantage_report.context("multi-vantage check failed")?;
    let gateways_out = gateway_results.context("gateway probe failed")?;
    report = merge_gateways(report, gateways_out);
    report.resolution = resolution;

    let report = if cli.per_chunk {
        drill_down(report, &bees, timeout, cli.concurrency)
            .await
            .context("per-chunk drill-down failed")?
    } else {
        report
    };

    let report = match cli.target_overlay.as_deref() {
        Some(target) => annotate_target_overlay(report, target),
        None => report,
    };

    print!("{}", render_report(&report, cli.output.into()));

    if cli.reseed {
        let stamp = cli
            .stamp
            .as_ref()
            .context("--reseed requires --stamp <batch-id>")?;
        let target_bee = bees.first().context("no bee URL for --reseed")?;
        let status = check_stamp(target_bee, stamp, timeout)
            .await
            .context("stamp pre-flight failed")?;
        eprint!("{}", render_stamp_status(&status));
        if !status.exists || !status.usable {
            anyhow::bail!("refusing to re-seed: stamp is not usable; see warnings above");
        }
        let req = ReseedRequest {
            reference: reference.clone(),
            bee_url: target_bee.clone(),
            batch_id: stamp.clone(),
            timeout,
        };
        reseed(req).await.context("re-seed failed")?;
        eprintln!("re-seeded {} via {}", reference, target_bee);
    }

    let any_retrievable = report.vantages.iter().any(|v| v.retrievable == Some(true))
        || report.gateways.iter().any(|g| g.retrievable == Some(true));
    if !any_retrievable {
        std::process::exit(2);
    }
    Ok(())
}