diff --git a/src/commands.rs b/src/commands.rs index 08befeac..5c892295 100644 --- a/src/commands.rs +++ b/src/commands.rs @@ -8,6 +8,7 @@ use std::convert::TryFrom as _; use std::env; use std::ffi; use std::fmt; +use std::fs; use std::io; use std::mem; use std::ops; @@ -33,6 +34,8 @@ use crate::output; use crate::pinentry; use crate::Context; +const NITROCLI_EXT_PREFIX: &str = "nitrocli-"; + /// Set `libnitrokey`'s log level based on the execution context's verbosity. fn set_log_level(ctx: &mut Context<'_>) { let log_lvl = match ctx.config.verbosity { @@ -1096,15 +1099,52 @@ pub fn pws_status(ctx: &mut Context<'_>, all: bool) -> anyhow::Result<()> { }) } +/// Find and list all available extensions. +/// +/// The logic used in this function should use the same criteria as +/// `resolve_extension`. +pub(crate) fn discover_extensions(path_var: &ffi::OsStr) -> anyhow::Result> { + let dirs = env::split_paths(path_var); + let mut commands = Vec::new(); + + for dir in dirs { + match fs::read_dir(&dir) { + Ok(entries) => { + for entry in entries { + let entry = entry?; + let path = entry.path(); + if path.is_file() { + let name = entry.file_name(); + let file = name.to_string_lossy(); + if file.starts_with(NITROCLI_EXT_PREFIX) { + let mut file = file.into_owned(); + file.replace_range(..NITROCLI_EXT_PREFIX.len(), ""); + commands.push(file); + } + } + } + } + Err(ref err) if err.kind() == io::ErrorKind::NotFound => (), + x => x + .map(|_| ()) + .with_context(|| format!("Failed to iterate entries of directory {}", dir.display()))?, + } + } + Ok(commands) +} + /// Resolve an extension provided by name to an actual path. /// /// Extensions are (executable) files that have the "nitrocli-" prefix /// and are discoverable via the `PATH` environment variable. +/// +/// The logic used in this function should use the same criteria as +/// `discover_extensions`. pub(crate) fn resolve_extension( path_var: &ffi::OsStr, ext_name: &ffi::OsStr, ) -> anyhow::Result { - let mut bin_name = ffi::OsString::from("nitrocli-"); + let mut bin_name = ffi::OsString::from(NITROCLI_EXT_PREFIX); bin_name.push(ext_name); for dir in env::split_paths(path_var) { diff --git a/src/main.rs b/src/main.rs index e7a7d2f5..c8b73dc1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -69,6 +69,10 @@ use std::io; use std::process; use std::str; +use structopt::clap::ErrorKind; +use structopt::clap::SubCommand; +use structopt::StructOpt; + const NITROCLI_BINARY: &str = "NITROCLI_BINARY"; const NITROCLI_MODEL: &str = "NITROCLI_MODEL"; const NITROCLI_USB_PATH: &str = "NITROCLI_USB_PATH"; @@ -105,15 +109,44 @@ impl fmt::Display for DirectExitError { impl error::Error for DirectExitError {} /// Parse the command-line arguments and execute the selected command. -fn handle_arguments(ctx: &mut Context<'_>, args: Vec) -> anyhow::Result<()> { - use structopt::StructOpt; - - match args::Args::from_iter_safe(args.iter()) { +fn handle_arguments(ctx: &mut Context<'_>, argv: Vec) -> anyhow::Result<()> { + match args::Args::from_iter_safe(argv.iter()) { Ok(args) => { ctx.config.update(&args); args.cmd.execute(ctx) } - Err(err) => { + Err(mut err) => { + if err.kind == ErrorKind::HelpDisplayed { + // For the convenience of the user we'd like to list the + // available extensions in the help text. At the same time, we + // don't want to unconditionally iterate through PATH (which may + // contain directories with loads of files that need scanning) + // for every command invoked. So we do that listing only if a + // help text is actually displayed. + let path = ctx.path.clone().unwrap_or_else(ffi::OsString::new); + if let Ok(extensions) = commands::discover_extensions(&path) { + let mut clap = args::Args::clap(); + for name in extensions { + // Because of clap's brain dead API, we see no other way + // but to leak the string we created here. That's okay, + // though, because we exit in a moment anyway. + let about = Box::leak(format!("Run the {} extension", name).into_boxed_str()); + clap = clap.subcommand( + SubCommand::with_name(&name) + // Use some magic number here that causes all + // extensions to be listed after all other + // subcommands. + .display_order(1000) + .about(about as &'static str), + ); + } + // At this point we are *pretty* sure that repeated invocation + // will result in another error. So should be fine to unwrap + // here. + err = clap.get_matches_from_safe(argv.iter()).unwrap_err(); + } + } + if err.use_stderr() { Err(err.into()) } else { diff --git a/src/tests/extensions.rs b/src/tests/extensions.rs index a295949e..d546e8a5 100644 --- a/src/tests/extensions.rs +++ b/src/tests/extensions.rs @@ -8,6 +8,36 @@ use std::fs; use super::*; +#[test] +fn no_extensions_to_discover() -> anyhow::Result<()> { + let exts = crate::commands::discover_extensions(&ffi::OsString::new())?; + assert!(exts.is_empty(), "{:?}", exts); + Ok(()) +} + +#[test] +fn extension_discovery() -> anyhow::Result<()> { + let dir1 = tempfile::tempdir()?; + let dir2 = tempfile::tempdir()?; + + { + let ext1_path = dir1.path().join("nitrocli-ext1"); + let ext2_path = dir1.path().join("nitrocli-ext2"); + let ext3_path = dir2.path().join("nitrocli-super-1337-extensions111one"); + let _ext1 = fs::File::create(&ext1_path)?; + let _ext2 = fs::File::create(&ext2_path)?; + let _ext3 = fs::File::create(&ext3_path)?; + + let path = env::join_paths(&[dir1.path(), dir2.path()])?; + let mut exts = crate::commands::discover_extensions(&path)?; + // We can't assume a fixed ordering of extensions, because that is + // platform/file system dependent. So sort here to fix it. + exts.sort(); + assert_eq!(exts, vec!["ext1", "ext2", "super-1337-extensions111one"]); + } + Ok(()) +} + #[test] fn resolve_extensions() -> anyhow::Result<()> { let dir1 = tempfile::tempdir()?;