moshpits 0.8.19

// Copyright (c) 2025 moshpit developers
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or https://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use std::{collections::BTreeSet, ffi::OsString, io::Cursor, sync::LazyLock};

use clap::{ArgAction, ArgMatches, CommandFactory, Parser, parser::ValueSource};
use config::{ConfigError, Map, Source, Value, ValueKind};
use getset::{CopyGetters, Getters};
use libmoshpit::PathDefaults;
use vergen_pretty::{Pretty, vergen_pretty_env};

static LONG_VERSION: LazyLock<String> = LazyLock::new(|| {
    let pretty = Pretty::builder().env(vergen_pretty_env!()).build();
    let mut cursor = Cursor::new(vec![]);
    let mut output = env!("CARGO_PKG_VERSION").to_string();
    output.push_str("\n\n");
    pretty
        .display(&mut cursor)
        .expect("writing to Vec never fails");
    output += &String::from_utf8_lossy(cursor.get_ref());
    output
});

#[derive(Clone, CopyGetters, Debug, Getters, Parser)]
#[command(author, version, about, long_version = LONG_VERSION.as_str(), long_about = None)]
pub(crate) struct Cli {
    /// Set logging verbosity.  More v's, more verbose.
    #[clap(
        short,
        long,
        action = ArgAction::Count,
        help = "Turn up logging verbosity (multiple will turn it up more)",
        conflicts_with = "quiet",
    )]
    #[getset(get_copy = "pub(crate)")]
    verbose: u8,
    /// Set logging quietness.  More q's, more quiet.
    #[clap(
        short,
        long,
        action = ArgAction::Count,
        help = "Turn down logging verbosity (multiple will turn it down more)",
        conflicts_with = "verbose",
    )]
    #[getset(get_copy = "pub(crate)")]
    quiet: u8,
    /// Enable logging to stdout/stderr in additions to the tracing output file
    /// * NOTE * - This should not be used when running as a daemon/service
    #[clap(short, long, help = "Enable logging to stdout/stderr")]
    enable_std_output: bool,
    /// The absolute path to a non-standard config file
    #[clap(short, long, help = "Specify the absolute path to the config file")]
    #[getset(get = "pub(crate)")]
    config_absolute_path: Option<String>,
    /// The absolute path to a non-standard tracing output file
    #[clap(
        short,
        long,
        help = "Specify the absolute path to the tracing output file"
    )]
    #[getset(get = "pub(crate)")]
    tracing_absolute_path: Option<String>,
    /// An absolute path to a non-standard private key file
    #[clap(
        short,
        long,
        help = "Specify the absolute path to the private key file"
    )]
    #[getset(get = "pub(crate)")]
    private_key_path: Option<String>,
    /// An absolute path to a non-standard public key file
    #[clap(
        short = 'k',
        long,
        help = "Specify the absolute path to the public key file"
    )]
    #[getset(get = "pub(crate)")]
    public_key_path: Option<String>,
    /// Additional delay in milliseconds after the client's first UDP datagram
    /// is received, before bulk terminal data is sent.  Provides extra margin
    /// for NAT bindings on slow NAT devices when clients use `--nat-warmup`.
    #[clap(
        long,
        value_name = "MILLIS",
        help = "Extra delay (ms) after peer discovery before sending terminal data"
    )]
    #[getset(get_copy = "pub(crate)")]
    warmup_delay_ms: Option<u64>,
    /// Minimum delay in microseconds between consecutive diff packets from the same
    /// PTY read batch.  Spreads back-to-back packets over time to prevent burst loss
    /// on stateful NAT devices.  Set to 0 to disable.
    #[clap(
        long,
        value_name = "MICROS",
        help = "Min inter-packet delay (µs) between diff chunks [default: 1000]"
    )]
    #[getset(get_copy = "pub(crate)")]
    pacing_delay_us: Option<u64>,
    /// Minimum wire protocol version this server will accept.  Clients that cannot
    /// speak at least this version are rejected during key exchange.  Raise this to
    /// retire older protocols (e.g. after a security fix); clamped to the range this
    /// build implements.
    #[clap(
        long,
        value_name = "VERSION",
        help = "Minimum wire protocol version to accept (older clients are rejected)"
    )]
    #[getset(get_copy = "pub(crate)")]
    min_protocol_version: Option<u16>,
    /// TERM environment variable to set for spawned shells
    #[clap(
        long,
        value_name = "TERM",
        default_value = "xterm-256color",
        help = "TERM environment variable for spawned shells"
    )]
    #[getset(get = "pub(crate)")]
    term_type: String,
    /// Ordered KEX algorithms to prefer (comma-separated).
    /// Example: `--kex-algos ml-kem-768-sha256,x25519-sha256`
    #[clap(
        long,
        value_name = "ALGOS",
        help = "Ordered KEX algorithms to prefer, comma-separated [supported: x25519-sha256 (default), ml-kem-768-sha256, ml-kem-512-sha256, ml-kem-1024-sha256, p384-sha384, p256-sha256]"
    )]
    #[getset(get = "pub(crate)")]
    kex_algos: Option<String>,
    /// Ordered AEAD algorithms to prefer (comma-separated).
    /// Example: `--aead-algos chacha20-poly1305,aes256-gcm-siv`
    #[clap(
        long,
        value_name = "ALGOS",
        help = "Ordered AEAD algorithms to prefer, comma-separated [supported: aes256-gcm-siv (default), aes256-gcm, chacha20-poly1305, aes128-gcm-siv]"
    )]
    #[getset(get = "pub(crate)")]
    aead_algos: Option<String>,
    /// Ordered MAC algorithms to prefer (comma-separated).
    /// Example: `--mac-algos hmac-sha256`
    #[clap(
        long,
        value_name = "ALGOS",
        help = "Ordered MAC algorithms to prefer, comma-separated [supported: hmac-sha512 (default), hmac-sha256]"
    )]
    #[getset(get = "pub(crate)")]
    mac_algos: Option<String>,
    /// Ordered KDF algorithms to prefer (comma-separated).
    /// Example: `--kdf-algos hkdf-sha512`
    #[clap(
        long,
        value_name = "ALGOS",
        help = "Ordered KDF algorithms to prefer, comma-separated [supported: hkdf-sha256 (default), hkdf-sha384, hkdf-sha512]"
    )]
    #[getset(get = "pub(crate)")]
    kdf_algos: Option<String>,
    /// Set of clap argument ids the user actually supplied on the command line
    /// (`ValueSource::CommandLine`), populated by [`Cli::parse_argv`].  This lets
    /// [`Source::collect`] emit only user-provided values so clap defaults no
    /// longer clobber the config file or environment.  Not a real CLI argument.
    #[clap(skip)]
    #[getset(get = "pub(crate)")]
    explicit_args: BTreeSet<String>,
}

impl Cli {
    /// Parse `argv` into a [`Cli`], recording which arguments the user actually
    /// supplied on the command line.
    ///
    /// clap's derived struct cannot distinguish a value the user typed from a
    /// `default_value`, so we parse twice: once into the typed struct and once
    /// into raw [`ArgMatches`], whose [`ValueSource`] reveals the true origin of
    /// each argument.  The resulting [`Cli::explicit_args`] set drives
    /// [`Source::collect`] so clap defaults no longer clobber the config file or
    /// environment.
    ///
    /// # Errors
    /// Returns a clap error if `argv` fails to parse.
    pub(crate) fn parse_argv<I, T>(argv: I) -> clap::error::Result<Self>
    where
        I: IntoIterator<Item = T> + Clone,
        T: Into<OsString> + Clone,
    {
        let mut cli = Cli::try_parse_from(argv.clone())?;
        let matches = <Cli as CommandFactory>::command().try_get_matches_from(argv)?;
        cli.explicit_args = explicit_command_line_ids(&matches);
        Ok(cli)
    }
}

/// Collect the ids of all arguments whose value originated from the command
/// line (as opposed to a default value or being unset).
fn explicit_command_line_ids(matches: &ArgMatches) -> BTreeSet<String> {
    matches
        .ids()
        .filter(|id| matches.value_source(id.as_str()) == Some(ValueSource::CommandLine))
        .map(|id| id.as_str().to_string())
        .collect()
}

fn build_algo_table(
    kex: Option<&str>,
    aead: Option<&str>,
    mac: Option<&str>,
    kdf: Option<&str>,
) -> Option<Map<String, Value>> {
    let mut table = Map::new();
    let parse = |s: &str| -> Vec<Value> {
        s.split(',')
            .map(|a| Value::new(None, ValueKind::String(a.trim().to_string())))
            .collect()
    };
    for (key, opt) in [("kex", kex), ("aead", aead), ("mac", mac), ("kdf", kdf)] {
        if let Some(s) = opt {
            let _old = table.insert(
                key.to_string(),
                Value::new(None, ValueKind::Array(parse(s))),
            );
        }
    }
    (!table.is_empty()).then_some(table)
}

impl Source for Cli {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<Map<String, Value>, ConfigError> {
        let mut map = Map::new();
        let origin = String::from("command line");
        // Only emit values the user actually typed on the command line.  Without
        // this gate, clap's defaults would be emitted here and override the
        // config file and environment (see `parse_argv` / `explicit_args`).
        let on = |id: &str| self.explicit_args.contains(id);

        if on("verbose") {
            let _old = map.insert(
                "verbose".to_string(),
                Value::new(Some(&origin), ValueKind::U64(u8::into(self.verbose))),
            );
        }
        if on("quiet") {
            let _old = map.insert(
                "quiet".to_string(),
                Value::new(Some(&origin), ValueKind::U64(u8::into(self.quiet))),
            );
        }
        if on("enable_std_output") {
            let _old = map.insert(
                "enable_std_output".to_string(),
                Value::new(Some(&origin), ValueKind::Boolean(self.enable_std_output)),
            );
        }
        if on("config_absolute_path")
            && let Some(config_path) = &self.config_absolute_path
        {
            let _old = map.insert(
                "config_path".to_string(),
                Value::new(Some(&origin), ValueKind::String(config_path.clone())),
            );
        }
        if on("tracing_absolute_path")
            && let Some(tracing_path) = &self.tracing_absolute_path
        {
            let _old = map.insert(
                "tracing_path".to_string(),
                Value::new(Some(&origin), ValueKind::String(tracing_path.clone())),
            );
        }
        if on("private_key_path")
            && let Some(private_key_path) = &self.private_key_path
        {
            let _old = map.insert(
                "private_key_path".to_string(),
                Value::new(Some(&origin), ValueKind::String(private_key_path.clone())),
            );
        }
        if on("public_key_path")
            && let Some(public_key_path) = &self.public_key_path
        {
            let _old = map.insert(
                "public_key_path".to_string(),
                Value::new(Some(&origin), ValueKind::String(public_key_path.clone())),
            );
        }
        if on("warmup_delay_ms")
            && let Some(warmup_delay_ms) = self.warmup_delay_ms
        {
            let _old = map.insert(
                "warmup_delay_ms".to_string(),
                Value::new(Some(&origin), ValueKind::U64(warmup_delay_ms)),
            );
        }
        if on("pacing_delay_us")
            && let Some(pacing_delay_us) = self.pacing_delay_us
        {
            let _old = map.insert(
                "pacing_delay_us".to_string(),
                Value::new(Some(&origin), ValueKind::U64(pacing_delay_us)),
            );
        }
        if on("min_protocol_version")
            && let Some(min_protocol_version) = self.min_protocol_version
        {
            let _old = map.insert(
                "min_protocol_version".to_string(),
                Value::new(
                    Some(&origin),
                    ValueKind::U64(u64::from(min_protocol_version)),
                ),
            );
        }
        if on("term_type") {
            let _old = map.insert(
                "term_type".to_string(),
                Value::new(Some(&origin), ValueKind::String(self.term_type.clone())),
            );
        }
        if let Some(table) = build_algo_table(
            self.kex_algos.as_deref().filter(|_| on("kex_algos")),
            self.aead_algos.as_deref().filter(|_| on("aead_algos")),
            self.mac_algos.as_deref().filter(|_| on("mac_algos")),
            self.kdf_algos.as_deref().filter(|_| on("kdf_algos")),
        ) {
            let _old = map.insert(
                "preferred_algorithms".to_string(),
                Value::new(Some(&origin), ValueKind::Table(table)),
            );
        }
        Ok(map)
    }
}

impl PathDefaults for Cli {
    fn env_prefix(&self) -> String {
        env!("CARGO_PKG_NAME").to_ascii_uppercase()
    }

    fn config_absolute_path(&self) -> Option<String> {
        self.config_absolute_path.clone()
    }

    fn default_file_path(&self) -> String {
        env!("CARGO_PKG_NAME").to_string()
    }

    fn default_file_name(&self) -> String {
        env!("CARGO_PKG_NAME").to_string()
    }

    fn tracing_absolute_path(&self) -> Option<String> {
        self.tracing_absolute_path.clone()
    }

    fn default_tracing_path(&self) -> String {
        format!("{}/logs", env!("CARGO_PKG_NAME"))
    }

    fn default_tracing_file_name(&self) -> String {
        env!("CARGO_PKG_NAME").to_string()
    }
}

#[cfg(test)]
mod test {
    use config::Source as _;

    use super::Cli;

    fn parse(args: &[&str]) -> Cli {
        Cli::parse_argv(args).expect("args parse")
    }

    #[test]
    fn cli_defaults() {
        let cli = parse(&["mps"]);
        assert_eq!(cli.verbose(), 0);
        assert_eq!(cli.quiet(), 0);
        assert!(!cli.enable_std_output);
        assert!(cli.config_absolute_path().is_none());
        assert!(cli.tracing_absolute_path().is_none());
        assert!(cli.private_key_path().is_none());
        assert!(cli.public_key_path().is_none());
        assert_eq!(cli.term_type(), "xterm-256color");
    }

    #[test]
    fn cli_verbose() {
        let cli = parse(&["mps", "-vv"]);
        assert_eq!(cli.verbose(), 2);
    }

    #[test]
    fn cli_quiet() {
        let cli = parse(&["mps", "-qq"]);
        assert_eq!(cli.quiet(), 2);
    }

    #[test]
    fn cli_private_key_path() {
        let cli = parse(&["mps", "-p", "/tmp/key"]);
        assert_eq!(cli.private_key_path().as_deref(), Some("/tmp/key"));
    }

    #[test]
    fn cli_public_key_path() {
        let cli = parse(&["mps", "-k", "/tmp/key.pub"]);
        assert_eq!(cli.public_key_path().as_deref(), Some("/tmp/key.pub"));
    }

    #[test]
    fn cli_term_type_default() {
        let cli = parse(&["mps"]);
        assert_eq!(cli.term_type(), "xterm-256color");
    }

    #[test]
    fn cli_term_type_custom() {
        let cli = parse(&["mps", "--term-type", "screen-256color"]);
        assert_eq!(cli.term_type(), "screen-256color");
    }

    #[test]
    fn cli_term_type_various_values() {
        let test_cases = vec!["xterm", "screen", "tmux-256color", "linux", "vt100"];
        for term in test_cases {
            let cli = parse(&["mps", "--term-type", term]);
            assert_eq!(cli.term_type(), term);
        }
    }

    #[test]
    fn cli_source_collect() {
        let cli = parse(&["mps"]);
        let map = cli.collect().expect("collect should succeed");
        // Defaulted values are NOT emitted (no clap-default pollution that would
        // clobber the config file / environment).
        assert!(!map.contains_key("verbose"));
        assert!(!map.contains_key("quiet"));
        assert!(!map.contains_key("enable_std_output"));
        assert!(!map.contains_key("term_type"));
        // Optional keys absent when not provided
        assert!(!map.contains_key("private_key_path"));
        assert!(!map.contains_key("public_key_path"));
        assert!(!map.contains_key("config_path"));
        assert!(!map.contains_key("tracing_path"));
    }

    #[test]
    fn cli_source_collect_emits_explicit_flags() {
        let cli = parse(&["mps", "-vv", "--enable-std-output", "--term-type", "screen"]);
        let map = cli.collect().expect("collect should succeed");
        assert!(map.contains_key("verbose"));
        assert!(map.contains_key("enable_std_output"));
        assert!(map.contains_key("term_type"));
        // quiet was not supplied, so it is omitted.
        assert!(!map.contains_key("quiet"));
    }

    #[test]
    fn cli_source_collect_with_paths() {
        let cli = parse(&[
            "mps",
            "-p",
            "/tmp/priv",
            "-k",
            "/tmp/pub",
            "-c",
            "/tmp/config.toml",
            "-t",
            "/tmp/trace.log",
            "--term-type",
            "tmux-256color",
        ]);
        let map = cli.collect().expect("collect should succeed");
        assert!(map.contains_key("private_key_path"));
        assert!(map.contains_key("public_key_path"));
        assert!(map.contains_key("config_path"));
        assert!(map.contains_key("tracing_path"));
        assert!(map.contains_key("term_type"));
        // Verify the term_type value was collected correctly
        let term_type = map.get("term_type").expect("term_type should be in map");
        assert_eq!(
            term_type.clone().into_string().ok(),
            Some("tmux-256color".to_string())
        );
    }

    #[test]
    fn cli_min_protocol_version_absent_by_default() {
        let cli = parse(&["mps"]);
        assert!(cli.min_protocol_version().is_none());
        let map = cli.collect().expect("collect should succeed");
        assert!(!map.contains_key("min_protocol_version"));
    }

    #[test]
    fn cli_min_protocol_version_collected() {
        let cli = parse(&["mps", "--min-protocol-version", "3"]);
        assert_eq!(cli.min_protocol_version(), Some(3));
        let map = cli.collect().expect("collect should succeed");
        let value = map
            .get("min_protocol_version")
            .expect("min_protocol_version should be in map");
        assert_eq!(value.clone().into_uint().ok(), Some(3));
    }

    #[test]
    fn cli_source_collect_optional_delays() {
        let cli = parse(&["mps", "--warmup-delay-ms", "50", "--pacing-delay-us", "250"]);
        let map = cli.collect().expect("collect should succeed");
        let warmup = map
            .get("warmup_delay_ms")
            .expect("warmup_delay_ms should be in map");
        assert_eq!(warmup.clone().into_uint().ok(), Some(50));
        let pacing = map
            .get("pacing_delay_us")
            .expect("pacing_delay_us should be in map");
        assert_eq!(pacing.clone().into_uint().ok(), Some(250));
    }

    #[test]
    fn cli_source_collect_algo_table() {
        let cli = parse(&[
            "mps",
            "--kex-algos",
            "x25519-sha256, ml-kem-768-sha256",
            "--aead-algos",
            "aes256-gcm-siv",
            "--mac-algos",
            "hmac-sha512",
            "--kdf-algos",
            "hkdf-sha256",
        ]);
        let map = cli.collect().expect("collect should succeed");
        let value = map
            .get("preferred_algorithms")
            .expect("preferred_algorithms should be in map");
        let table = value
            .clone()
            .into_table()
            .expect("preferred_algorithms should be a table");
        assert!(table.contains_key("kex"));
        assert!(table.contains_key("aead"));
        assert!(table.contains_key("mac"));
        assert!(table.contains_key("kdf"));
    }

    #[test]
    fn cli_path_defaults() {
        use libmoshpit::PathDefaults as _;
        let cli = parse(&["mps"]);
        assert_eq!(cli.env_prefix(), "MOSHPITS");
        assert_eq!(cli.default_file_path(), "moshpits");
        assert_eq!(cli.default_file_name(), "moshpits");
        assert_eq!(cli.default_tracing_path(), "moshpits/logs");
        assert_eq!(cli.default_tracing_file_name(), "moshpits");
        assert!(cli.config_absolute_path().is_none());
        assert!(cli.tracing_absolute_path().is_none());
    }
}