Add the --model option to restrict the device model to connect to

This patch adds the -m/--model option that can be used to restrict the
device model to connect to.  Per default, nitrocli connects to any
available Nitrokey device.  If this new option is set, it will instead
only connect to devices of the given Nitrokey model.

We introduce a new struct DeviceModel instead of using
nitrokey::DeviceModel to make sure that the command-line options are
parsed properly.  On the long term, we should add a connect_model
function to the nitrokey crate to make the connection code easier.

2 files changed, 58 insertions, 2 deletions

diff --git a/nitrocli/src/args.rs b/nitrocli/src/args.rs
index 869734c..ad296c2 100644
--- a/nitrocli/src/args.rs
+++ b/nitrocli/src/args.rs
@@ -28,10 +28,43 @@ use crate::pinentry;
 
 type Result<T> = result::Result<T, Error>;
 
+/// The available Nitrokey models.
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum DeviceModel {
+  Pro,
+  Storage,
+}
+
+impl fmt::Display for DeviceModel {
+  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+    write!(
+      f,
+      "{}",
+      match *self {
+        DeviceModel::Pro => "pro",
+        DeviceModel::Storage => "storage",
+      }
+    )
+  }
+}
+
+impl str::FromStr for DeviceModel {
+  type Err = ();
+
+  fn from_str(s: &str) -> result::Result<Self, Self::Err> {
+    match s {
+      "pro" => Ok(DeviceModel::Pro),
+      "storage" => Ok(DeviceModel::Storage),
+      _ => Err(()),
+    }
+  }
+}
+
 /// A command execution context that captures additional data pertaining
 /// the command execution.
 #[derive(Debug)]
 pub struct ExecCtx {
+  pub model: Option<DeviceModel>,
   pub verbosity: u64,
 }
 
@@ -952,6 +985,7 @@ fn pws_status(ctx: &ExecCtx, args: Vec<String>) -> Result<()> {
 /// Parse the command-line arguments and return the selected command and
 /// the remaining arguments for the command.
 fn parse_arguments(args: Vec<String>) -> Result<(Command, ExecCtx, Vec<String>)> {
+  let mut model: Option<DeviceModel> = None;
   let mut verbosity = 0;
   let mut command = Command::Status;
   let mut subargs = vec![];
@@ -961,6 +995,11 @@ fn parse_arguments(args: Vec<String>) -> Result<(Command, ExecCtx, Vec<String>)>
     argparse::IncrBy::<u64>(1),
     "Increase the log level (can be supplied multiple times)",
   );
+  let _ = parser.refer(&mut model).add_option(
+    &["-m", "--model"],
+    argparse::StoreOption,
+    "Select the device model to connect to (pro|storage)",
+  );
   parser.set_description("Provides access to a Nitrokey device");
   let _ = parser.refer(&mut command).required().add_argument(
     "command",
@@ -978,7 +1017,7 @@ fn parse_arguments(args: Vec<String>) -> Result<(Command, ExecCtx, Vec<String>)>
 
   subargs.insert(0, format!("nitrocli {}", command));
 
-  let ctx = ExecCtx { verbosity };
+  let ctx = ExecCtx { model, verbosity };
   Ok((command, ctx, subargs))
 }
 
diff --git a/nitrocli/src/commands.rs b/nitrocli/src/commands.rs
index 27faf05..e125f17 100644
--- a/nitrocli/src/commands.rs
+++ b/nitrocli/src/commands.rs
@@ -58,13 +58,30 @@ fn set_log_level(ctx: &args::ExecCtx) {
 fn get_device(ctx: &args::ExecCtx) -> Result<nitrokey::DeviceWrapper> {
   set_log_level(ctx);
 
-  nitrokey::connect().map_err(|_| Error::Error("Nitrokey device not found".to_string()))
+  match ctx.model {
+    Some(model) => match model {
+      args::DeviceModel::Pro => nitrokey::Pro::connect().map(nitrokey::DeviceWrapper::Pro),
+      args::DeviceModel::Storage => {
+        nitrokey::Storage::connect().map(nitrokey::DeviceWrapper::Storage)
+      }
+    },
+    None => nitrokey::connect(),
+  }
+  .map_err(|_| Error::Error("Nitrokey device not found".to_string()))
 }
 
 /// Connect to a Nitrokey Storage device and return it.
 fn get_storage_device(ctx: &args::ExecCtx) -> Result<nitrokey::Storage> {
   set_log_level(ctx);
 
+  if let Some(model) = ctx.model {
+    if model != args::DeviceModel::Storage {
+      return Err(Error::Error(
+        "This command is only available on the Nitrokey Storage".to_string(),
+      ));
+    }
+  }
+
   nitrokey::Storage::connect().or_else(|_| {
     Err(Error::Error(
       "Nitrokey Storage device not found".to_string(),