2 files changed, 10 insertions, 6 deletions
diff --git a/nitrocli/src/args.rs b/nitrocli/src/args.rs
index 6f02832..246490e 100644
--- a/nitrocli/src/args.rs
+++ b/nitrocli/src/args.rs
@@ -64,6 +64,15 @@ Enum! {DeviceModel, [
   Storage => "storage"
 ]}
 
+impl From<DeviceModel> for nitrokey::Model {
+  fn from(model: DeviceModel) -> nitrokey::Model {
+    match model {
+      DeviceModel::Pro => nitrokey::Model::Pro,
+      DeviceModel::Storage => nitrokey::Model::Storage,
+    }
+  }
+}
+
 /// A top-level command for nitrocli.
 Enum! {Command, [
   Config => "config",
diff --git a/nitrocli/src/commands.rs b/nitrocli/src/commands.rs
index ed3c2c4..b37f9c5 100644
--- a/nitrocli/src/commands.rs
+++ b/nitrocli/src/commands.rs
@@ -59,12 +59,7 @@ fn get_device(ctx: &mut args::ExecCtx<'_>) -> Result<nitrokey::DeviceWrapper> {
   set_log_level(ctx);
 
   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)
-      }
-    },
+    Some(model) => nitrokey::connect_model(model.into()),
     None => nitrokey::connect(),
   }
   .map_err(|_| Error::Error("Nitrokey device not found".to_string()))