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.

5 files changed, 68 insertions, 5 deletions

diff --git a/nitrocli/CHANGELOG.md b/nitrocli/CHANGELOG.md
index e6a4480..41cfd91 100644
--- a/nitrocli/CHANGELOG.md
+++ b/nitrocli/CHANGELOG.md
@@ -1,6 +1,8 @@
 Unreleased
 ----------
 - Added the `-v`/`--verbose` option to control libnitrokey log level
+- Added the `-m`/`--model` option to restrict connections to a device
+  model
 
 
 0.2.1
diff --git a/nitrocli/doc/nitrocli.1 b/nitrocli/doc/nitrocli.1
index 7786fc1..e4dbf10 100644
--- a/nitrocli/doc/nitrocli.1
+++ b/nitrocli/doc/nitrocli.1
@@ -3,7 +3,7 @@
 nitrocli \- access Nitrokey devices
 .SH SYNOPSIS
 .B nitrocli
-\fR[\fB\-v\fR|\fB\-\-verbose\fR]
+[\fB\-m\fR|\fB\-\-model pro\fR|\fBstorage\fR] \fR[\fB\-v\fR|\fB\-\-verbose\fR]
 \fIcommand\fR
 [\fIarguments\fR]
 .SH DESCRIPTION
@@ -13,7 +13,12 @@ It can be used to access the encrypted volume, the one-time password generator,
 and the password safe.
 .SH OPTIONS
 .TP
-.B \-v, \-\-verbose
+\fB\-m\fR, \fB\-\-model pro\fR|\fBstorage\fR
+Restrict connections to the given device model.
+If this option is not set, nitrocli will connect to any connected Nitrokey Pro
+or Nitrokey Storage device.
+.TP
+\fB\-v\fR, \fB\-\-verbose\fR
 Enable additional logging and control its verbosity. Logging enabled through
 this option will appear on the standard error stream. This option can be
 supplied multiple times. A single occurrence will show additional warnings.
@@ -76,7 +81,7 @@ This command might require the user PIN (see the Configuration section).
 .TP
 \fBnitrocli otp set \fIslot name secret \
 \fR[\fB\-a\fR|\fB\-\-algorithm \fIalgorithm\fR] \
-[\fB\-d\fR|\fB\-\-digits \fI digits\fR] [\fB\-c\fR|\fB\-\-counter \fIcounter\fR] \
+[\fB\-d\fR|\fB\-\-digits \fIdigits\fR] [\fB\-c\fR|\fB\-\-counter \fIcounter\fR] \
 [\fB\-t\fR|\fB\-\-time-window \fItime window\fR] [\fB\-\-ascii\fR]
 Configure a one-time password slot.
 \fIslot\fR is the number of the slot to configure.
diff --git a/nitrocli/doc/nitrocli.1.pdf b/nitrocli/doc/nitrocli.1.pdf
index c5e8898..5127435 100644
--- a/nitrocli/doc/nitrocli.1.pdf
+++ b/nitrocli/doc/nitrocli.1.pdf
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(),