2 files changed, 26 insertions, 2 deletions
diff --git a/src/device.rs b/src/device.rs
index 326df5d..4f96860 100644
--- a/src/device.rs
+++ b/src/device.rs
@@ -445,9 +445,25 @@ pub fn connect() -> Result<DeviceWrapper, CommandError> {
     }
 }
 
+fn get_connected_model() -> Option<Model> {
+    unsafe {
+        match nitrokey_sys::NK_get_device_model() {
+            nitrokey_sys::NK_device_model_NK_PRO => Some(Model::Pro),
+            nitrokey_sys::NK_device_model_NK_STORAGE => Some(Model::Storage),
+            _ => None,
+        }
+    }
+}
+
+fn create_device_wrapper(model: Model) -> DeviceWrapper {
+    match model {
+        Model::Pro => DeviceWrapper::Pro(Pro {}),
+        Model::Storage => DeviceWrapper::Storage(Storage {}),
+    }
+}
+
 fn get_connected_device() -> Option<DeviceWrapper> {
-    // TODO: check connected device
-    Some(DeviceWrapper::Pro(Pro {}))
+    get_connected_model().map(create_device_wrapper)
 }
 
 fn connect_model(model: Model) -> bool {
diff --git a/src/tests/device.rs b/src/tests/device.rs
index 60ca33a..c2c5336 100644
--- a/src/tests/device.rs
+++ b/src/tests/device.rs
@@ -33,6 +33,10 @@ fn connect_pro() {
     assert!(::connect().is_ok());
     assert!(::Pro::connect().is_ok());
     assert!(::Storage::connect().is_err());
+    match ::connect().unwrap() {
+        ::DeviceWrapper::Pro(_) => assert!(true),
+        ::DeviceWrapper::Storage(_) => assert!(false),
+    };
 }
 
 #[test]
@@ -41,6 +45,10 @@ fn connect_storage() {
     assert!(::connect().is_ok());
     assert!(::Pro::connect().is_err());
     assert!(::Storage::connect().is_ok());
+    match ::connect().unwrap() {
+        ::DeviceWrapper::Pro(_) => assert!(false),
+        ::DeviceWrapper::Storage(_) => assert!(true),
+    };
 }
 
 fn assert_empty_serial_number() {