Use Manager to connect to a device

1 files changed, 19 insertions, 24 deletions

diff --git a/nitrocli/src/commands.rs b/nitrocli/src/commands.rs
index 5deefa2..ffe4496 100644
--- a/nitrocli/src/commands.rs
+++ b/nitrocli/src/commands.rs
@@ -55,19 +55,27 @@ fn set_log_level(ctx: &mut args::ExecCtx<'_>) {
   nitrokey::set_log_level(log_lvl);
 }
 
-/// Connect to any Nitrokey device and return it.
-fn get_device(ctx: &mut args::ExecCtx<'_>) -> Result<nitrokey::DeviceWrapper> {
+fn with_device<F>(ctx: &mut args::ExecCtx<'_>, op: F) -> Result<()>
+where
+  F: FnOnce(&mut args::ExecCtx<'_>, nitrokey::DeviceWrapper) -> Result<()>,
+{
+  let mut manager = nitrokey::take()?;
   set_log_level(ctx);
 
-  match ctx.model {
-    Some(model) => nitrokey::connect_model(model.into()),
-    None => nitrokey::connect(),
+  let device = match ctx.model {
+    Some(model) => manager.connect_model(model.into()),
+    None => manager.connect(),
   }
-  .map_err(|_| Error::from("Nitrokey device not found"))
+  .map_err(|_| Error::from("Nitrokey device not found"))?;
+
+  op(ctx, device)
 }
 
-/// Connect to a Nitrokey Storage device and return it.
-fn get_storage_device(ctx: &mut args::ExecCtx<'_>) -> Result<nitrokey::Storage> {
+fn with_storage_device<F>(ctx: &mut args::ExecCtx<'_>, op: F) -> Result<()>
+where
+  F: FnOnce(&mut args::ExecCtx<'_>, nitrokey::Storage) -> Result<()>,
+{
+  let mut manager = nitrokey::take()?;
   set_log_level(ctx);
 
   if let Some(model) = ctx.model {
@@ -78,22 +86,9 @@ fn get_storage_device(ctx: &mut args::ExecCtx<'_>) -> Result<nitrokey::Storage>
     }
   }
 
-  nitrokey::Storage::connect().or_else(|_| Err(Error::from("Nitrokey Storage device not found")))
-}
-
-fn with_device<F>(ctx: &mut args::ExecCtx<'_>, op: F) -> Result<()>
-where
-  F: FnOnce(&mut args::ExecCtx<'_>, nitrokey::DeviceWrapper) -> Result<()>,
-{
-  let device = get_device(ctx)?;
-  op(ctx, device)
-}
-
-fn with_storage_device<F>(ctx: &mut args::ExecCtx<'_>, op: F) -> Result<()>
-where
-  F: FnOnce(&mut args::ExecCtx<'_>, nitrokey::Storage) -> Result<()>,
-{
-  let device = get_storage_device(ctx)?;
+  let device = manager
+    .connect_storage()
+    .or_else(|_| Err(Error::from("Nitrokey Storage device not found")))?;
   op(ctx, device)
 }