From 3d951ab54daea9c362bdae2e0945701955c412e5 Mon Sep 17 00:00:00 2001
From: Robin Krahl <robin.krahl@ireas.org>
Date: Sun, 30 Dec 2018 00:04:23 +0100
Subject: Add function Device::get_model

This patch adds the function Device::get_model that returns the model of
the connected Nitrokey stick.
---
 CHANGELOG.md    |  1 +
 src/device.rs   | 54 +++++++++++++++++++++++++++++++++++++++++++++++++-----
 src/lib.rs      |  2 +-
 tests/device.rs | 17 +++++++++++++++++
 4 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fd1a307..a30e531 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@
 - Update to Rust edition 2018.
 - Remove the `test-no-device` feature.
 - Update the rand dependency to version 0.6.
+- Add function `Device::get_model` that returns the connected model.
 
 # v0.2.1 (2018-12-10)
 
diff --git a/src/device.rs b/src/device.rs
index 1455a57..1f9ae3d 100644
--- a/src/device.rs
+++ b/src/device.rs
@@ -1,3 +1,5 @@
+use std::fmt;
+
 use libc;
 use nitrokey_sys;
 
@@ -8,14 +10,27 @@ use crate::pws::GetPasswordSafe;
 use crate::util::{get_command_result, get_cstring, get_last_error, result_from_string, CommandError};
 
 /// Available Nitrokey models.
-#[derive(Debug, PartialEq)]
-enum Model {
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum Model {
     /// The Nitrokey Storage.
     Storage,
     /// The Nitrokey Pro.
     Pro,
 }
 
+impl fmt::Display for Model {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "{}",
+            match *self {
+                Model::Pro => "Pro",
+                Model::Storage => "Storage",
+            }
+        )
+    }
+}
+
 /// A wrapper for a Nitrokey device of unknown type.
 ///
 /// Use the function [`connect`][] to obtain a wrapped instance.  The wrapper implements all traits
@@ -211,6 +226,20 @@ pub struct StorageStatus {
 /// This trait provides the commands that can be executed without authentication and that are
 /// present on all supported Nitrokey devices.
 pub trait Device: Authenticate + GetPasswordSafe + GenerateOtp {
+    /// Returns the model of the connected Nitrokey device.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # use nitrokey::CommandError;
+    ///
+    /// # fn try_main() -> Result<(), CommandError> {
+    /// let device = nitrokey::connect()?;
+    /// println!("Connected to a Nitrokey {}", device.get_model());
+    /// #    Ok(())
+    /// # }
+    fn get_model(&self) -> Model;
+
     /// Returns the serial number of the Nitrokey device.  The serial number is the string
     /// representation of a hex number.
     ///
@@ -563,7 +592,14 @@ impl GenerateOtp for DeviceWrapper {
     }
 }
 
-impl Device for DeviceWrapper {}
+impl Device for DeviceWrapper {
+    fn get_model(&self) -> Model {
+        match *self {
+            DeviceWrapper::Pro(_) => Model::Pro,
+            DeviceWrapper::Storage(_) => Model::Storage,
+        }
+    }
+}
 
 impl Pro {
     /// Connects to a Nitrokey Pro.
@@ -597,7 +633,11 @@ impl Drop for Pro {
     }
 }
 
-impl Device for Pro {}
+impl Device for Pro {
+    fn get_model(&self) -> Model {
+        Model::Pro
+    }
+}
 
 impl GenerateOtp for Pro {}
 
@@ -743,7 +783,11 @@ impl Drop for Storage {
     }
 }
 
-impl Device for Storage {}
+impl Device for Storage {
+    fn get_model(&self) -> Model {
+        Model::Pro
+    }
+}
 
 impl GenerateOtp for Storage {}
 
diff --git a/src/lib.rs b/src/lib.rs
index d0e4357..37f839a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -102,7 +102,7 @@ use nitrokey_sys;
 
 pub use crate::auth::{Admin, Authenticate, User};
 pub use crate::config::Config;
-pub use crate::device::{connect, Device, DeviceWrapper, Pro, Storage, StorageStatus, VolumeStatus};
+pub use crate::device::{connect, Device, DeviceWrapper, Model, Pro, Storage, StorageStatus, VolumeStatus};
 pub use crate::otp::{ConfigureOtp, GenerateOtp, OtpMode, OtpSlotData};
 pub use crate::pws::{GetPasswordSafe, PasswordSafe, SLOT_COUNT};
 pub use crate::util::{CommandError, LogLevel};
diff --git a/tests/device.rs b/tests/device.rs
index 2c8cfa6..a6098b7 100644
--- a/tests/device.rs
+++ b/tests/device.rs
@@ -81,6 +81,23 @@ fn disconnect() {
     assert_empty_serial_number();
 }
 
+fn require_model(model: nitrokey::Model) {
+    assert_eq!(model, nitrokey::connect().unwrap().get_model());
+    assert_eq!(model, Target::connect().unwrap().get_model());
+}
+
+#[test]
+#[cfg_attr(not(feature = "test-pro"), ignore)]
+fn get_model_pro() {
+    require_model(nitrokey::Model::Pro);
+}
+
+#[test]
+#[cfg_attr(not(feature = "test-storage"), ignore)]
+fn get_model_storage() {
+    require_model(nitrokey::Model::Storage);
+}
+
 #[test]
 #[cfg_attr(not(any(feature = "test-pro", feature = "test-storage")), ignore)]
 fn get_serial_number() {
-- 
cgit v1.2.1