6 files changed, 393 insertions, 3 deletions

diff --git a/nitrocli/src/args.rs b/nitrocli/src/args.rs
index f47f1a0..2e56e9e 100644
--- a/nitrocli/src/args.rs
+++ b/nitrocli/src/args.rs
@@ -17,8 +17,14 @@
 // * along with this program.  If not, see <http://www.gnu.org/licenses/>. *
 // *************************************************************************
 
+use std::collections;
+use std::env;
 use std::ffi;
+use std::fmt;
+use std::fs;
 use std::io;
+use std::path;
+use std::process;
 use std::result;
 use std::str;
 
@@ -28,6 +34,7 @@ use crate::pinentry;
 use crate::RunCtx;
 
 type Result<T> = result::Result<T, Error>;
+type Extensions = collections::BTreeMap<String, path::PathBuf>;
 
 /// Wraps a writer and buffers its output.
 ///
@@ -141,6 +148,87 @@ Enum! {Builtin, [
   Unencrypted => ("unencrypted", unencrypted),
 ]}
 
+#[derive(Debug)]
+enum Command {
+  Builtin(Builtin),
+  Extension(String),
+}
+
+impl Command {
+  pub fn execute(
+    &self,
+    ctx: &mut ExecCtx<'_>,
+    args: Vec<String>,
+    extensions: Extensions,
+  ) -> Result<()> {
+    match self {
+      Command::Builtin(command) => command.execute(ctx, args),
+      Command::Extension(extension) => {
+        match extensions.get(extension) {
+          Some(path) => {
+            // Note that theoretically we could just exec the extension
+            // and be done. However, the problem with that approach is
+            // that it makes testing extension support much more nasty,
+            // because the test process would be overwritten in the
+            // process, requiring us to essentially fork & exec nitrocli
+            // beforehand -- which is much more involved from a cargo
+            // test context.
+            let mut cmd = process::Command::new(path);
+
+            if let Some(model) = ctx.model {
+              let _ = cmd.args(&["--model", model.as_ref()]);
+            };
+
+            let out = cmd
+              // TODO: We may want to take this path from the command
+              //       execution context.
+              .args(&["--nitrocli", &env::current_exe()?.to_string_lossy()])
+              .args(&["--verbosity", &ctx.verbosity.to_string()])
+              .args(&args[1..])
+              .output()
+              .map_err(Into::<Error>::into)?;
+            ctx.stdout.write_all(&out.stdout)?;
+            ctx.stderr.write_all(&out.stderr)?;
+            if out.status.success() {
+              Ok(())
+            } else {
+              Err(Error::ExtensionFailed(
+                extension.to_string(),
+                out.status.code(),
+              ))
+            }
+          }
+          None => Err(Error::Error(format!("Unknown command: {}", extension))),
+        }
+      }
+    }
+  }
+}
+
+impl fmt::Display for Command {
+  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+    match self {
+      Command::Builtin(cmd) => write!(f, "{}", cmd),
+      Command::Extension(ext) => write!(f, "{}", ext),
+    }
+  }
+}
+
+impl str::FromStr for Command {
+  type Err = ();
+
+  fn from_str(s: &str) -> result::Result<Self, Self::Err> {
+    Ok(match Builtin::from_str(s) {
+      Ok(cmd) => Command::Builtin(cmd),
+      // Note that at this point we cannot know whether the extension
+      // exists or not and so we always return success. However, if we
+      // fail looking up the corresponding command an error will be
+      // emitted later on.
+      Err(()) => Command::Extension(s.to_string()),
+    })
+  }
+}
+
 Enum! {ConfigCommand, [
   Get => ("get", config_get),
   Set => ("set", config_set),
@@ -913,6 +1001,60 @@ fn pws_status(ctx: &mut ExecCtx<'_>, args: Vec<String>) -> Result<()> {
   commands::pws_status(ctx, all)
 }
 
+/// Find all the available extensions. Extensions are (executable) files
+/// that have the "nitrocli-" prefix and are discoverable via the `PATH`
+/// environment variable.
+// Note that we use a BTreeMap here to have a stable ordering among
+// extensions. That makes for a nicer user experience over a HashMap as
+// they appear in the help text and random changes in position are
+// confusing.
+fn find_extensions(path: &ffi::OsStr) -> Result<Extensions> {
+  // The std::env module has several references to the PATH environment
+  // variable, indicating that this name is considered platform
+  // independent from their perspective. We do the same.
+  let dirs = env::split_paths(path);
+  let mut commands = Extensions::new();
+  let prefix = format!("{}-", crate::NITROCLI);
+
+  for dir in dirs {
+    match fs::read_dir(&path::Path::new(&dir)) {
+      Ok(entries) => {
+        for entry in entries {
+          let entry = entry?;
+          let path = entry.path();
+          if path.is_file() {
+            let file = String::from(entry.file_name().to_str().unwrap());
+            // Note that we deliberately do not check whether the file
+            // we found is executable. If it is not we will just fail
+            // later on with a permission denied error. The reasons for
+            // this behavior are two fold:
+            // 1) Checking whether a file is executable in Rust is
+            //    painful (as of 1.37 there exists the PermissionsExt
+            //    trait but it is available only for Unix based
+            //    systems).
+            // 2) It is considered a better user experience to show an
+            //    extension that we found (we list them in the help
+            //    text) even if it later turned out to be not usable
+            //    over not showing it and silently doing nothing --
+            //    mostly because anything residing in PATH should be
+            //    executable anyway and given that its name also starts
+            //    with nitrocli- we are pretty sure that's a bug on the
+            //    user's side.
+            if file.starts_with(&prefix) {
+              let mut file = file;
+              file.replace_range(..prefix.len(), "");
+              assert!(commands.insert(file, path).is_none());
+            }
+          }
+        }
+      }
+      Err(ref err) if err.kind() == io::ErrorKind::NotFound => (),
+      x => x.map(|_| ())?,
+    }
+  }
+  Ok(commands)
+}
+
 /// Parse the command-line arguments and execute the selected command.
 pub(crate) fn handle_arguments(ctx: &mut RunCtx<'_>, args: Vec<String>) -> Result<()> {
   use std::io::Write;
@@ -924,8 +1066,20 @@ pub(crate) fn handle_arguments(ctx: &mut RunCtx<'_>, args: Vec<String>) -> Resul
     fmt_enum!(DeviceModel::all_variants())
   );
   let mut verbosity = 0;
-  let mut command = Builtin::Status;
-  let cmd_help = cmd_help!(command);
+  let path = ctx.path.take().unwrap_or_else(ffi::OsString::new);
+  let extensions = find_extensions(&path)?;
+  let mut command = Command::Builtin(Builtin::Status);
+  let commands = Builtin::all_variants()
+    .iter()
+    .map(AsRef::as_ref)
+    .map(ToOwned::to_owned)
+    .chain(extensions.keys().cloned())
+    .collect::<Vec<_>>()
+    .join("|");
+  // argparse's help text formatting is pretty bad for our intents and
+  // purposes. In particular, line breaks are just ignored by its custom
+  // line wrapping algorithm.
+  let cmd_help = format!("The command to execute ({})", commands);
   let mut subargs = vec![];
   let mut parser = argparse::ArgumentParser::new();
   let _ = parser.refer(&mut version).add_option(
@@ -981,6 +1135,50 @@ pub(crate) fn handle_arguments(ctx: &mut RunCtx<'_>, args: Vec<String>) -> Resul
       no_cache: ctx.no_cache,
       verbosity,
     };
-    command.execute(&mut ctx, subargs)
+    command.execute(&mut ctx, subargs, extensions)
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  fn no_extensions_available() -> Result<()> {
+    let exts = find_extensions(&ffi::OsString::new())?;
+    assert!(exts.is_empty(), "{:?}", exts);
+    Ok(())
+  }
+
+  #[test]
+  fn discover_extensions() -> 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()])
+        .map_err(|err| Error::Error(err.to_string()))?;
+      let exts = find_extensions(&path)?;
+
+      let mut it = exts.iter();
+      // Because we control the file names and the order of directories
+      // in `PATH` we can safely assume a fixed order in which
+      // extensions should be discovered.
+      assert_eq!(it.next(), Some((&"ext1".to_string(), &ext1_path)));
+      assert_eq!(it.next(), Some((&"ext2".to_string(), &ext2_path)));
+      assert_eq!(
+        it.next(),
+        Some((&"super-1337-extensions111one".to_string(), &ext3_path))
+      );
+      assert_eq!(it.next(), None);
+    }
+    Ok(())
   }
 }
diff --git a/nitrocli/src/error.rs b/nitrocli/src/error.rs
index 819bed8..1f6424e 100644
--- a/nitrocli/src/error.rs
+++ b/nitrocli/src/error.rs
@@ -45,6 +45,7 @@ pub enum Error {
   NitrokeyError(Option<&'static str>, nitrokey::Error),
   Utf8Error(str::Utf8Error),
   Error(String),
+  ExtensionFailed(String, Option<i32>),
 }
 
 impl TryInto<nitrokey::Error> for Error {
@@ -99,6 +100,13 @@ impl fmt::Display for Error {
       Error::Utf8Error(_) => write!(f, "Encountered UTF-8 conversion error"),
       Error::IoError(ref e) => write!(f, "IO error: {}", e),
       Error::Error(ref e) => write!(f, "{}", e),
+      Error::ExtensionFailed(ref ext, rc) => {
+        write!(f, "Extension {} failed", ext)?;
+        if let Some(rc) = rc {
+          write!(f, " with exit code {}", rc)?;
+        }
+        Ok(())
+      }
     }
   }
 }
diff --git a/nitrocli/src/main.rs b/nitrocli/src/main.rs
index bb4b007..92ecfe6 100644
--- a/nitrocli/src/main.rs
+++ b/nitrocli/src/main.rs
@@ -101,6 +101,8 @@ pub(crate) struct RunCtx<'io> {
   pub stdout: &'io mut dyn io::Write,
   /// The `Write` object used as standard error throughout the program.
   pub stderr: &'io mut dyn io::Write,
+  /// The content of the `PATH` environment variable.
+  pub path: Option<ffi::OsString>,
   /// The admin PIN, if provided through an environment variable.
   pub admin_pin: Option<ffi::OsString>,
   /// The user PIN, if provided through an environment variable.
@@ -129,6 +131,11 @@ fn run<'ctx, 'io: 'ctx>(ctx: &'ctx mut RunCtx<'io>, args: Vec<String>) -> i32 {
         // argparse printed an error message
         _ => 1,
       },
+      Error::ExtensionFailed(_, rc) => {
+        // We let the extension itself deal with error reporting, we
+        // just mirror its exit code (if any).
+        rc.unwrap_or(1)
+      }
       _ => {
         let _ = eprintln!(ctx, "{}", err);
         1
@@ -146,6 +153,7 @@ fn main() {
   let ctx = &mut RunCtx {
     stdout: &mut stdout,
     stderr: &mut stderr,
+    path: env::var_os("PATH"),
     admin_pin: env::var_os(NITROCLI_ADMIN_PIN),
     user_pin: env::var_os(NITROCLI_USER_PIN),
     new_admin_pin: env::var_os(NITROCLI_NEW_ADMIN_PIN),
diff --git a/nitrocli/src/tests/extension_model_test.py b/nitrocli/src/tests/extension_model_test.py
new file mode 100644
index 0000000..651c8e7
--- /dev/null
+++ b/nitrocli/src/tests/extension_model_test.py
@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+
+from argparse import (
+  ArgumentParser,
+)
+from enum import (
+  Enum,
+)
+from sys import (
+  argv,
+  exit,
+)
+
+
+class Action(Enum):
+  """An action to perform."""
+  NITROCLI = "nitrocli"
+  MODEL = "model"
+  VERBOSITY = "verbosity"
+
+  @classmethod
+  def all(cls):
+    """Return the list of all the enum members' values."""
+    return [x.value for x in cls.__members__.values()]
+
+
+def main(args):
+  """The extension's main function."""
+  parser = ArgumentParser()
+  parser.add_argument(choices=Action.all(), dest="what")
+  parser.add_argument("--nitrocli", action="store", default=None)
+  parser.add_argument("--model", action="store", default=None)
+  # We deliberately store the argument to this option as a string
+  # because we can differentiate between None and a valid value, in
+  # order to verify that it really is supplied.
+  parser.add_argument("--verbosity", action="store", default=None)
+
+  namespace = parser.parse_args(args[1:])
+  if namespace.what == Action.NITROCLI.value:
+    print(namespace.nitrocli)
+  elif namespace.what == Action.MODEL.value:
+    print(namespace.model)
+  elif namespace.what == Action.VERBOSITY.value:
+    print(namespace.verbosity)
+  else:
+    return 1
+
+  return 0
+
+
+if __name__ == "__main__":
+  exit(main(argv))
diff --git a/nitrocli/src/tests/mod.rs b/nitrocli/src/tests/mod.rs
index dc15117..ce3f861 100644
--- a/nitrocli/src/tests/mod.rs
+++ b/nitrocli/src/tests/mod.rs
@@ -85,6 +85,15 @@ impl Builder {
     self
   }
 
+  /// Set the `PATH` used for looking up extensions.
+  fn path<P>(mut self, path: P) -> Self
+  where
+    P: Into<ffi::OsString>,
+  {
+    self.0.path = Some(path.into());
+    self
+  }
+
   /// Set the password to use for certain operations.
   fn password<P>(mut self, password: P) -> Self
   where
@@ -102,6 +111,7 @@ impl Builder {
 
 struct Nitrocli {
   model: Option<nitrokey::Model>,
+  path: Option<ffi::OsString>,
   admin_pin: Option<ffi::OsString>,
   user_pin: Option<ffi::OsString>,
   new_admin_pin: Option<ffi::OsString>,
@@ -113,6 +123,7 @@ impl Nitrocli {
   pub fn new() -> Self {
     Self {
       model: None,
+      path: None,
       admin_pin: Some(nitrokey::DEFAULT_ADMIN_PIN.into()),
       user_pin: Some(nitrokey::DEFAULT_USER_PIN.into()),
       new_admin_pin: None,
@@ -166,6 +177,7 @@ impl Nitrocli {
     let ctx = &mut crate::RunCtx {
       stdout: &mut stdout,
       stderr: &mut stderr,
+      path: self.path.clone(),
       admin_pin: self.admin_pin.clone(),
       user_pin: self.user_pin.clone(),
       new_admin_pin: self.new_admin_pin.clone(),
diff --git a/nitrocli/src/tests/run.rs b/nitrocli/src/tests/run.rs
index c59c660..2212577 100644
--- a/nitrocli/src/tests/run.rs
+++ b/nitrocli/src/tests/run.rs
@@ -17,6 +17,11 @@
 // * along with this program.  If not, see <http://www.gnu.org/licenses/>. *
 // *************************************************************************
 
+use std::fs;
+use std::io::Write;
+use std::os::unix::fs::OpenOptionsExt;
+use std::path;
+
 use super::*;
 
 #[test]
@@ -101,3 +106,110 @@ fn version_option() {
   test(&re, "--version");
   test(&re, "-V");
 }
+
+#[test]
+fn extension() -> crate::Result<()> {
+  let ext_dir = tempfile::tempdir()?;
+  {
+    let mut ext = fs::OpenOptions::new()
+      .create(true)
+      .mode(0o755)
+      .write(true)
+      .open(ext_dir.path().join("nitrocli-ext"))?;
+
+    ext.write_all(
+      br#"#!/usr/bin/env python
+print("success")
+"#,
+    )?;
+  }
+
+  let path = ext_dir.path().as_os_str().to_os_string();
+  let out = Nitrocli::make().path(path).build().handle(&["ext"])?;
+  assert_eq!(out, "success\n");
+  Ok(())
+}
+
+#[test]
+fn extension_failure() -> crate::Result<()> {
+  let ext_dir = tempfile::tempdir()?;
+  {
+    let mut ext = fs::OpenOptions::new()
+      .create(true)
+      .mode(0o755)
+      .write(true)
+      .open(ext_dir.path().join("nitrocli-ext"))?;
+
+    ext.write_all(
+      br#"#!/usr/bin/env python
+import sys
+sys.exit(42);
+"#,
+    )?;
+  }
+
+  let path = ext_dir.path().as_os_str().to_os_string();
+  let err = Nitrocli::make()
+    .path(path)
+    .build()
+    .handle(&["ext"])
+    .unwrap_err();
+
+  match err {
+    crate::Error::ExtensionFailed(ext, rc) => {
+      assert_eq!(&ext, "ext");
+      assert_eq!(rc, Some(42));
+    }
+    _ => panic!("Unexpected error variant found: {:?}", err),
+  };
+  Ok(())
+}
+
+#[test_device]
+fn extension_arguments(model: nitrokey::Model) -> crate::Result<()> {
+  fn test<F>(model: nitrokey::Model, what: &str, args: &[&str], check: F) -> crate::Result<()>
+  where
+    F: FnOnce(&str) -> bool,
+  {
+    let ext_dir = tempfile::tempdir()?;
+    {
+      let mut ext = fs::OpenOptions::new()
+        .create(true)
+        .mode(0o755)
+        .write(true)
+        .open(ext_dir.path().join("nitrocli-ext"))?;
+
+      ext.write_all(include_bytes!("extension_model_test.py"))?;
+    }
+
+    let mut args = args.to_vec();
+    args.append(&mut vec!["ext", what]);
+
+    let path = ext_dir.path().as_os_str().to_os_string();
+    let out = Nitrocli::make()
+      .model(model)
+      .path(path)
+      .build()
+      .handle(&args)?;
+
+    assert!(check(&out), out);
+    Ok(())
+  }
+
+  test(model, "model", &[], |out| {
+    out == model.to_string().to_lowercase() + "\n"
+  })?;
+  test(model, "nitrocli", &[], |out| {
+    path::Path::new(out)
+      .file_stem()
+      .unwrap()
+      .to_str()
+      .unwrap()
+      .trim()
+      .contains("nitrocli")
+  })?;
+  test(model, "verbosity", &[], |out| out == "0\n")?;
+  test(model, "verbosity", &["-v"], |out| out == "1\n")?;
+  test(model, "verbosity", &["-v", "--verbose"], |out| out == "2\n")?;
+  Ok(())
+}