bgpkit-parser 0.23.0

use itertools::Itertools;
use serde_json::json;
use std::io::{IsTerminal, Write};
use std::net::IpAddr;
use std::path::PathBuf;

use bgpkit_parser::render::text::Style;
use bgpkit_parser::{BgpElem, BgpkitParser, Elementor, RecoveryConfig, RecoveryEvent, RecoveryGap};
use clap::{Parser, ValueEnum};
use ipnet::IpNet;

/// Output format for the parser
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, ValueEnum)]
enum OutputFormat {
    /// Default pipe-separated format (elem-level)
    #[default]
    Default,
    /// JSON format (one object per line, elem-level)
    Json,
    /// Pretty-printed JSON format (elem-level)
    JsonPretty,
    /// PSV format with header (elem-level)
    Psv,
    /// Layered human-readable text, one block per MRT record. Always uses
    /// record-level output (implies `--level records`); the other formats
    /// follow `--level` and default to elems. Inspired by bgpdump's
    /// human-readable output.
    Text,
}

/// Output level granularity
#[derive(Debug, Clone, Copy, Default, ValueEnum)]
enum OutputLevel {
    /// Output BGP elements (per-prefix)
    #[default]
    Elems,
    /// Output MRT records
    Records,
}

/// When to colorize `--format text` output
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, ValueEnum)]
enum ColorChoice {
    /// Color only when stdout is a terminal (default)
    #[default]
    Auto,
    /// Color even when output is piped or redirected
    Always,
    /// Never color, even on a terminal
    Never,
}

/// bgpkit-parser is a simple CLI tool for parsing MRT/BGP/BMP files.
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Opts {
    /// File path to a MRT file, local or remote.
    #[clap(name = "FILE")]
    file_path: PathBuf,

    /// Set the cache directory for caching remote files. Default behavior does not enable caching.
    #[clap(short, long)]
    cache_dir: Option<PathBuf>,

    /// Output format
    #[clap(short = 'F', long, value_enum, default_value_t = OutputFormat::Default)]
    format: OutputFormat,

    /// Output level: elems (per-prefix) or records (MRT records)
    #[clap(short = 'L', long, value_enum, default_value_t = OutputLevel::Elems)]
    level: OutputLevel,

    /// Output as JSON objects (shorthand for --format json)
    #[clap(long)]
    json: bool,

    /// Pretty-print JSON output (shorthand for --format json-pretty)
    #[clap(long)]
    pretty: bool,

    /// Output as full PSV entries with header (shorthand for --format psv)
    #[clap(long)]
    psv: bool,

    /// Include each record's raw bytes as hex: a `HEX:` line in text
    /// blocks, a `hex` field in JSON records. Record-level only — implies
    /// `--level records`; requires `--format text`, `json`, or
    /// `json-pretty`. Count-only runs (-e/-r) ignore this flag.
    #[clap(long)]
    hex: bool,

    /// Colorize `--format text` output: `auto` colors only when stdout is a
    /// terminal, `NO_COLOR` disables coloring, and `CLICOLOR_FORCE` forces it.
    /// The other formats are never colored.
    #[clap(long, value_enum, default_value = "auto")]
    color: ColorChoice,

    /// Count BGP elems
    #[clap(short, long)]
    elems_count: bool,

    /// Count MRT records
    #[clap(short, long)]
    records_count: bool,

    /// Recover after damaged MRT framing and report skipped byte ranges on stderr
    #[clap(long)]
    recover: bool,

    #[clap(flatten)]
    filters: Filters,
}

#[derive(Parser, Debug)]
struct Filters {
    /// Filter by origin AS Number
    #[clap(short = 'o', long)]
    origin_asn: Option<u32>,

    /// Generic filter expression (can be used multiple times)
    /// Format: "key=value" for positive match, "key!=value" for negative match
    /// Examples: --filter "origin_asn!=13335" --filter "peer_ip!=192.168.1.1"
    /// For multi-value filters: --filter "origin_asns!=13335,15169" excludes both ASNs
    /// Supported keys: origin_asn, origin_asns, prefix, prefixes, peer_ip, peer_ips, peer_asn, peer_asns, type, as_path, community, ip_version
    #[clap(short = 'f', long = "filter")]
    filters: Vec<String>,

    /// Filter by network prefix
    #[clap(short = 'p', long)]
    prefix: Option<IpNet>,

    /// Include super-prefix when filtering
    #[clap(short = 's', long)]
    include_super: bool,

    /// Include sub-prefix when filtering
    #[clap(short = 'S', long)]
    include_sub: bool,

    /// Filter by IPv4 only
    #[clap(short = '4', long)]
    ipv4_only: bool,

    /// Filter by IPv6 only
    #[clap(short = '6', long)]
    ipv6_only: bool,

    /// Filter by peer IP address
    #[clap(short = 'j', long)]
    peer_ip: Vec<IpAddr>,

    /// Filter by peer ASN
    #[clap(short = 'J', long)]
    peer_asn: Option<u32>,

    /// Filter by elem type: announce (a) or withdraw (w)
    #[clap(short = 'm', long)]
    elem_type: Option<String>,

    /// Filter by start unix timestamp inclusive
    #[clap(short = 't', long)]
    start_ts: Option<f64>,

    /// Filter by end unix timestamp inclusive
    #[clap(short = 'T', long)]
    end_ts: Option<f64>,

    /// Filter by AS path regex string
    #[clap(short = 'a', long)]
    as_path: Option<String>,

    /// Filter by community regular expression (e.g. 64512:100 or no-export)
    #[clap(short = 'C', long)]
    community: Option<String>,
}

fn main() {
    let opts: Opts = Opts::parse();

    env_logger::init();

    let file_path = opts.file_path.to_str().unwrap();

    let parser_opt = match opts.cache_dir {
        None => BgpkitParser::new(file_path),
        Some(c) => BgpkitParser::new_cached(file_path, c.to_str().unwrap()),
    };

    let mut parser = match parser_opt {
        Ok(p) => p,
        Err(err) => {
            eprintln!("{err}");
            std::process::exit(1);
        }
    };

    if let Some(v) = opts.filters.as_path {
        parser = parser.add_filter("as_path", v.as_str()).unwrap();
    }
    if let Some(v) = opts.filters.community {
        parser = parser.add_filter("community", v.as_str()).unwrap();
    }
    if let Some(v) = opts.filters.origin_asn {
        parser = parser
            .add_filter("origin_asn", v.to_string().as_str())
            .unwrap();
    }
    if let Some(v) = opts.filters.prefix {
        let filter_type = match (opts.filters.include_super, opts.filters.include_sub) {
            (false, false) => "prefix",
            (true, false) => "prefix_super",
            (false, true) => "prefix_sub",
            (true, true) => "prefix_super_sub",
        };
        parser = parser
            .add_filter(filter_type, v.to_string().as_str())
            .unwrap();
    }
    if !opts.filters.peer_ip.is_empty() {
        let v = opts.filters.peer_ip.iter().map(|p| p.to_string()).join(",");
        parser = parser.add_filter("peer_ips", v.as_str()).unwrap();
    }
    if let Some(v) = opts.filters.peer_asn {
        parser = parser
            .add_filter("peer_asn", v.to_string().as_str())
            .unwrap();
    }
    if let Some(v) = opts.filters.elem_type {
        parser = parser.add_filter("type", v.as_str()).unwrap();
    }
    if let Some(v) = opts.filters.start_ts {
        parser = parser
            .add_filter("start_ts", v.to_string().as_str())
            .unwrap();
    }
    if let Some(v) = opts.filters.end_ts {
        parser = parser.add_filter("end_ts", v.to_string().as_str()).unwrap();
    }

    // Process generic filter expressions
    for filter_expr in &opts.filters.filters {
        match parse_filter_expression(filter_expr) {
            Ok((filter_type, filter_value)) => {
                parser = match parser.add_filter(&filter_type, &filter_value) {
                    Ok(p) => p,
                    Err(e) => {
                        eprintln!("Error adding filter '{}': {}", filter_expr, e);
                        std::process::exit(1);
                    }
                };
            }
            Err(e) => {
                eprintln!("Invalid filter expression '{}': {}", filter_expr, e);
                std::process::exit(1);
            }
        }
    }

    match (opts.filters.ipv4_only, opts.filters.ipv6_only) {
        (true, true) => {
            eprintln!("Error: --ipv4-only and --ipv6-only cannot be used together");
            std::process::exit(1);
        }
        (false, false) => {
            // no filters on IP version, skip
        }
        (true, false) => {
            parser = parser.add_filter("ip_version", "ipv4").unwrap();
        }
        (false, true) => {
            parser = parser.add_filter("ip_version", "ipv6").unwrap();
        }
    }

    // Determine final output format (shorthand flags override --format)
    let output_format = if opts.pretty {
        OutputFormat::JsonPretty
    } else if opts.json {
        OutputFormat::Json
    } else if opts.psv {
        OutputFormat::Psv
    } else {
        opts.format
    };

    // Count-only runs emit no hex, so the flag is ignored there (and the
    // format/recover constraints do not apply).
    let counting = opts.elems_count || opts.records_count;
    if opts.hex && !counting {
        if !matches!(
            output_format,
            OutputFormat::Json | OutputFormat::JsonPretty | OutputFormat::Text
        ) {
            eprintln!("Error: --hex requires --format text, json, or json-pretty");
            std::process::exit(1);
        }
        if opts.recover {
            eprintln!("Error: --hex does not support --recover");
            std::process::exit(1);
        }
    }

    let style = if resolve_color(opts.color, std::io::stdout().is_terminal()) {
        Style::ansi()
    } else {
        Style::plain()
    };

    let recovery_config = RecoveryConfig::default();
    // Element-level runs (element output or counting only elements) use the elem
    // iterators, which apply filters per element; everything else stays at the record
    // level. Counting both (-e -r) iterates records and converts once per record.
    let use_elem_stream = ((opts.elems_count && !opts.records_count)
        || (!opts.elems_count && !opts.records_count && matches!(opts.level, OutputLevel::Elems)))
        && output_format != OutputFormat::Text;

    // The hex path always works on original wire bytes from the
    // raw-record pipeline; it cannot share run_records, which holds
    // parsed records. Count-only runs emit no hex, so they keep the
    // normal elem/record counting pipelines and their semantics
    // (per-elem filtering for -e).
    let result = if opts.hex && !counting {
        run_hex_records(parser.into_filtered_raw_record_iter(), output_format, style)
    } else {
        match (opts.recover, use_elem_stream) {
            (true, true) => run_elems(
                parser
                    .into_recovering_elem_iter(recovery_config)
                    .map(|event| event.map_err(|error| error.to_string())),
                output_format,
                opts.elems_count,
                true,
            ),
            (false, true) => run_elems(
                parser
                    .into_elem_iter()
                    .map(|elem| Ok(RecoveryEvent::Item(elem))),
                output_format,
                opts.elems_count,
                false,
            ),
            (true, false) => run_records(
                parser
                    .into_recovering_record_iter(recovery_config)
                    .map(|event| event.map_err(|error| error.to_string())),
                output_format,
                style,
                opts.elems_count,
                opts.records_count,
                true,
            ),
            (false, false) => run_records(
                parser
                    .into_record_iter()
                    .map(|record| Ok(RecoveryEvent::Item(record))),
                output_format,
                style,
                opts.elems_count,
                opts.records_count,
                false,
            ),
        }
    };
    if let Err(error) = result {
        eprintln!("{error}");
        std::process::exit(1);
    }
}

/// Per-gap reporting and end-of-run summary for `--recover`.
struct RecoveryStats {
    enabled: bool,
    gap_count: usize,
    skipped_bytes: u64,
}

impl RecoveryStats {
    fn new(enabled: bool) -> Self {
        Self {
            enabled,
            gap_count: 0,
            skipped_bytes: 0,
        }
    }

    fn observe(&mut self, gap: &RecoveryGap) {
        self.gap_count += 1;
        self.skipped_bytes += gap.skipped_bytes();
        eprintln!(
            "recovered MRT framing: skipped bytes {}..{} ({} bytes, {:?}, {} confirming records): {}",
            gap.start_offset,
            gap.end_offset,
            gap.skipped_bytes(),
            gap.evidence,
            gap.confirmed_records,
            gap.cause
        );
    }

    fn print_summary(&self) {
        if self.enabled {
            eprintln!(
                "recovery summary: {} gaps, {} bytes skipped",
                self.gap_count, self.skipped_bytes
            );
        }
    }
}

fn run_elems<I>(
    events: I,
    output_format: OutputFormat,
    count_requested: bool,
    report_recovery: bool,
) -> Result<(), String>
where
    I: IntoIterator<Item = Result<RecoveryEvent<BgpElem>, String>>,
{
    let mut stdout = std::io::stdout();
    let mut elems_count = 0usize;
    let mut elem_index = 0usize;
    let mut stats = RecoveryStats::new(report_recovery);
    let mut terminal_error = None;

    for event in events {
        match event {
            Err(error) => {
                terminal_error = Some(error);
                break;
            }
            Ok(RecoveryEvent::Gap(gap)) => stats.observe(&gap),
            Ok(RecoveryEvent::Item(elem)) => {
                elems_count += 1;
                if count_requested {
                    continue;
                }
                let output = format_elem(&elem, output_format, elem_index);
                elem_index += 1;
                if !write_output(&mut stdout, &output)? {
                    return Ok(());
                }
            }
        }
    }

    if count_requested {
        println!("total elems: {elems_count}");
    }
    stats.print_summary();
    terminal_error.map_or(Ok(()), Err)
}

fn run_records<I>(
    events: I,
    output_format: OutputFormat,
    style: Style,
    elems_count_requested: bool,
    records_count_requested: bool,
    report_recovery: bool,
) -> Result<(), String>
where
    I: IntoIterator<Item = Result<RecoveryEvent<bgpkit_parser::MrtRecord>, String>>,
{
    let mut stdout = std::io::stdout();
    let mut elementor = Elementor::new();
    let mut records_count = 0usize;
    let mut elems_count = 0usize;
    let mut stats = RecoveryStats::new(report_recovery);
    let mut terminal_error = None;

    for event in events {
        match event {
            Err(error) => {
                terminal_error = Some(error);
                break;
            }
            Ok(RecoveryEvent::Gap(gap)) => stats.observe(&gap),
            Ok(RecoveryEvent::Item(record)) => {
                records_count += 1;
                if elems_count_requested {
                    // Counting both (-e -r): count every element of records that passed
                    // record-level filtering, matching the historical CLI behavior.
                    elems_count += elementor.record_to_elems(record).len();
                    continue;
                }
                if records_count_requested {
                    continue;
                }
                let output = format_record(&record, output_format, None, style);
                if !write_output(&mut stdout, &output)? {
                    return Ok(());
                }
            }
        }
    }

    match (elems_count_requested, records_count_requested) {
        (true, true) => {
            println!("total records: {records_count}");
            println!("total elems:   {elems_count}");
        }
        (false, true) => println!("total records: {records_count}"),
        (true, false) => println!("total elems: {elems_count}"),
        (false, false) => {}
    }
    stats.print_summary();
    terminal_error.map_or(Ok(()), Err)
}

/// Hex-augmented output over the raw-record pipeline: hex comes from the
/// original wire bytes (never a re-encoding), the block/JSON body from the
/// parsed record.
fn run_hex_records(
    records: impl Iterator<Item = (bgpkit_parser::RawMrtRecord, bgpkit_parser::MrtRecord)>,
    output_format: OutputFormat,
    style: Style,
) -> Result<(), String> {
    let mut stdout = std::io::stdout();

    for (raw, record) in records {
        let hex = bgpkit_parser::render::hex::encode(raw.raw_bytes().as_ref());
        let output = format_record(&record, output_format, Some(&hex), style);
        if !write_output(&mut stdout, &output)? {
            return Ok(());
        }
    }
    Ok(())
}

fn write_output(stdout: &mut std::io::Stdout, output: &str) -> Result<bool, String> {
    if let Err(error) = writeln!(stdout, "{output}") {
        if error.kind() == std::io::ErrorKind::BrokenPipe {
            return Ok(false);
        }
        return Err(error.to_string());
    }
    Ok(true)
}

fn format_elem(elem: &BgpElem, format: OutputFormat, index: usize) -> String {
    match format {
        OutputFormat::Json => {
            let val = json!(elem);
            val.to_string()
        }
        OutputFormat::JsonPretty => {
            let val = json!(elem);
            serde_json::to_string_pretty(&val).unwrap()
        }
        OutputFormat::Psv => {
            if index == 0 {
                format!("{}\n{}", BgpElem::get_psv_header(), elem.to_psv())
            } else {
                elem.to_psv()
            }
        }
        OutputFormat::Default => elem.to_string(),
        // The dispatch above routes text output to the record pipeline;
        // elem-level text rendering does not exist.
        OutputFormat::Text => unreachable!("text format renders records, not elems"),
    }
}

fn format_record(
    record: &bgpkit_parser::MrtRecord,
    format: OutputFormat,
    record_hex: Option<&str>,
    style: Style,
) -> String {
    match format {
        OutputFormat::Json => {
            let mut val = json!(record);
            if let (Some(hex), Some(obj)) = (&record_hex, val.as_object_mut()) {
                obj.insert("hex".to_string(), json!(hex));
            }
            val.to_string()
        }
        OutputFormat::JsonPretty => {
            let mut val = json!(record);
            if let (Some(hex), Some(obj)) = (&record_hex, val.as_object_mut()) {
                obj.insert("hex".to_string(), json!(hex));
            }
            serde_json::to_string_pretty(&val).unwrap()
        }
        OutputFormat::Text => match record_hex {
            Some(hex) => {
                bgpkit_parser::render::text::format_record_with_hex_and_style(record, hex, style)
            }
            None => bgpkit_parser::render::text::format_record_with_style(record, style),
        },
        OutputFormat::Psv | OutputFormat::Default => {
            // Use the Display implementation for MrtRecord
            format!("{}", record)
        }
    }
}

/// Parse a filter expression in the format "key=value" or "key!=value"
/// Returns (filter_type, filter_value) where filter_value may be prefixed with "!" for negation
///
/// For multi-value filters (e.g., "origin_asns!=13335,15169"), the negation is distributed
/// to each value: ("origin_asns", "!13335,!15169")
fn parse_filter_expression(expr: &str) -> Result<(String, String), String> {
    // Multi-value filter types that support comma-separated values
    let multi_value_filters = [
        "origin_asns",
        "prefixes",
        "prefixes_super",
        "prefixes_sub",
        "prefixes_super_sub",
        "peer_ips",
        "peer_asns",
    ];

    // Check for "!=" (negative filter) first
    if let Some(pos) = expr.find("!=") {
        let key = expr[..pos].trim();
        let value = expr[pos + 2..].trim();
        if key.is_empty() {
            return Err("filter key cannot be empty".to_string());
        }
        if value.is_empty() {
            return Err("filter value cannot be empty".to_string());
        }

        // For multi-value filters, prefix each value with "!"
        if multi_value_filters.contains(&key) {
            let negated_values: Vec<String> =
                value.split(',').map(|v| format!("!{}", v.trim())).collect();
            Ok((key.to_string(), negated_values.join(",")))
        } else {
            // For single-value filters, prefix the value with "!"
            Ok((key.to_string(), format!("!{}", value)))
        }
    }
    // Check for "=" (positive filter)
    else if let Some(pos) = expr.find('=') {
        let key = expr[..pos].trim();
        let value = expr[pos + 1..].trim();
        if key.is_empty() {
            return Err("filter key cannot be empty".to_string());
        }
        if value.is_empty() {
            return Err("filter value cannot be empty".to_string());
        }
        Ok((key.to_string(), value.to_string()))
    } else {
        Err("filter expression must contain '=' or '!=' (e.g., 'origin_asn=13335' or 'origin_asn!=13335')".to_string())
    }
}

/// Resolve `--color` against stdout and the `NO_COLOR` / `CLICOLOR_FORCE`
/// conventions.
fn resolve_color(choice: ColorChoice, stdout_is_terminal: bool) -> bool {
    match choice {
        ColorChoice::Always => true,
        ColorChoice::Never => false,
        ColorChoice::Auto => auto_color(
            stdout_is_terminal,
            std::env::var_os("NO_COLOR").is_some_and(|value| !value.is_empty()),
            std::env::var_os("CLICOLOR_FORCE")
                .is_some_and(|value| !value.is_empty() && value != "0"),
        ),
    }
}

/// The `auto` policy, split from the environment lookup so it can be tested
/// directly.
fn auto_color(stdout_is_terminal: bool, no_color: bool, clicolor_force: bool) -> bool {
    if no_color {
        false
    } else if clicolor_force {
        true
    } else {
        stdout_is_terminal
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn explicit_color_choices_are_honored() {
        assert!(resolve_color(ColorChoice::Always, false));
        assert!(!resolve_color(ColorChoice::Never, true));
    }

    #[test]
    fn auto_color_follows_the_conventions() {
        // NO_COLOR wins, CLICOLOR_FORCE forces color, otherwise stdout decides
        assert!(!auto_color(true, true, true));
        assert!(auto_color(false, false, true));
        assert!(auto_color(true, false, false));
        assert!(!auto_color(false, false, false));
    }
}