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(),