1 files changed, 18 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 {