3 files changed, 92 insertions, 5 deletions

diff --git a/nitrocli/src/main.rs b/nitrocli/src/main.rs
index 92aa79e..d75fe77 100644
--- a/nitrocli/src/main.rs
+++ b/nitrocli/src/main.rs
@@ -230,13 +230,37 @@ fn status() -> Result<()> {
 
 /// Open the encrypted volume on the nitrokey.
 fn open() -> Result<()> {
+  type Response = nitrokey::Response<nitrokey::StorageResponse>;
+
   return nitrokey_do(&|handle| {
-    let passphrase = pinentry::inquire_passphrase()?;
-    let payload = nitrokey::EnableEncryptedVolumeCommand::new(&passphrase);
-    let report = nitrokey::Report::from(payload);
+    let mut retry = 3;
+    loop {
+      let passphrase = pinentry::inquire_passphrase()?;
+      let payload = nitrokey::EnableEncryptedVolumeCommand::new(&passphrase);
+      let report = nitrokey::Report::from(payload);
 
-    transmit::<_, nitrokey::EmptyPayload>(handle, &report)?;
-    return Ok(());
+      let report = transmit::<_, nitrokey::EmptyPayload>(handle, &report)?;
+      let response = AsRef::<Response>::as_ref(&report.data);
+      let status = response.data.storage_status;
+
+      if status == nitrokey::StorageStatus::WrongPassword {
+        pinentry::clear_passphrase()?;
+        retry -= 1;
+
+        if retry > 0 {
+          println!("Wrong password, please reenter");
+          continue;
+        }
+        let error = "Opening encrypted volume failed: Wrong password";
+        return Err(Error::Error(error.to_string()));
+      }
+      if status != nitrokey::StorageStatus::Okay && status != nitrokey::StorageStatus::Idle {
+        let status = format!("{:?}", status);
+        let error = format!("Opening encrypted volume failed: {}", status);
+        return Err(Error::Error(error));
+      }
+      return Ok(());
+    }
   });
 }
 
diff --git a/nitrocli/src/nitrokey.rs b/nitrocli/src/nitrokey.rs
index d1d6c72..00a681c 100644
--- a/nitrocli/src/nitrokey.rs
+++ b/nitrocli/src/nitrokey.rs
@@ -241,6 +241,9 @@ pub enum CommandStatus {
 
 
 #[allow(dead_code)]
+#[derive(Copy)]
+#[derive(Clone)]
+#[derive(Debug)]
 #[derive(PartialEq)]
 #[repr(u8)]
 pub enum StorageStatus {
@@ -273,6 +276,17 @@ impl<P> AsRef<[u8]> for Response<P> {
 
 
 #[repr(packed)]
+pub struct StorageResponse {
+  pub padding1: [u8; 13],
+  pub command_counter: u8,
+  pub last_storage_command: Command,
+  pub storage_status: StorageStatus,
+  pub progress: u8,
+  pub padding2: [u8; 2],
+}
+
+
+#[repr(packed)]
 pub struct DeviceStatusResponse {
   pub padding0: [u8; 22],
   pub magic: u16,
diff --git a/nitrocli/src/pinentry.rs b/nitrocli/src/pinentry.rs
index 3fb533a..8de788f 100644
--- a/nitrocli/src/pinentry.rs
+++ b/nitrocli/src/pinentry.rs
@@ -69,6 +69,29 @@ pub fn inquire_passphrase() -> Result<Vec<u8>, Error> {
 }
 
 
+fn parse_pinentry_response(response: Vec<u8>) -> Result<(), Error> {
+  let string = String::from_utf8(response)?;
+  let lines: Vec<&str> = string.lines().collect();
+
+  if lines.len() == 1 && lines[0] == "OK" {
+    // We got the only valid answer we accept.
+    return Ok(());
+  }
+  return Err(Error::Error("Unexpected response: ".to_string() + &string));
+}
+
+
+/// Clear the cached passphrase.
+pub fn clear_passphrase() -> Result<(), Error> {
+  let command = "CLEAR_PASSPHRASE ".to_string() + CACHE_ID;
+  let output = process::Command::new("gpg-connect-agent").arg(command)
+    .arg("/bye")
+    .output()?;
+
+  return parse_pinentry_response(output.stdout);
+}
+
+
 #[cfg(test)]
 mod tests {
   use super::*;
@@ -109,4 +132,30 @@ mod tests {
       panic!("Unexpected result");
     }
   }
+
+  #[test]
+  fn parse_pinentry_response_ok() {
+    let response = "OK\n".to_string().into_bytes();
+    assert!(parse_pinentry_response(response).is_ok())
+  }
+
+  #[test]
+  fn parse_pinentry_response_ok_no_newline() {
+    let response = "OK".to_string().into_bytes();
+    assert!(parse_pinentry_response(response).is_ok())
+  }
+
+  #[test]
+  fn parse_pinentry_response_unexpected() {
+    let response = "ERR 42";
+    let expected = "Unexpected response: ".to_string() + response;
+
+    let error = parse_pinentry_response(response.to_string().into_bytes());
+
+    if let Error::Error(ref e) = error.err().unwrap() {
+      assert_eq!(e, &expected);
+    } else {
+      panic!("Unexpected result");
+    }
+  }
 }