From 7ecee69c01cc6e9ee3c62fa97eb132cc4c01b726 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sat, 30 Sep 2023 20:01:15 +0800
Subject: [PATCH] feat: read slot

---
 src/cmd_decrypt.rs | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/cmd_decrypt.rs b/src/cmd_decrypt.rs
index 22553a5..ba54a41 100644
--- a/src/cmd_decrypt.rs
+++ b/src/cmd_decrypt.rs
@@ -1,4 +1,5 @@
 use std::fs::File;
+use std::io;
 use std::io::{Read, Write};
 use std::path::PathBuf;
 use std::str::FromStr;
@@ -119,10 +120,6 @@ fn try_decrypt_key(envelop: &TinyEncryptEnvelop, pin: &Option<String>, slot: &Op
 }
 
 fn try_decrypt_key_ecdh(envelop: &TinyEncryptEnvelop, pin: &Option<String>, slot: &Option<String>) -> XResult<Vec<u8>> {
-    let is_slot_none = slot.as_ref().map(|s| s.is_empty()).unwrap_or(true);
-    if is_slot_none {
-        return simple_error!("--slot is required for ecdh");
-    }
     let wrap_key = WrapKey::parse(&envelop.encrypted_key)?;
     if wrap_key.header.enc.as_str() != ENC_AES256_GCM_P256 {
         return simple_error!("Unsupported header enc.");
@@ -131,12 +128,12 @@ fn try_decrypt_key_ecdh(envelop: &TinyEncryptEnvelop, pin: &Option<String>, slot
     let e_pub_key_bytes = opt_result!(decode_base64_url_no_pad(e_pub_key), "Invalid envelop: {}");
     let (_, subject_public_key_info) = opt_result!( SubjectPublicKeyInfo::from_der(&e_pub_key_bytes), "Invalid envelop: {}");
 
-    let slot = slot.as_ref().unwrap();
+    let slot = read_slot(slot)?;
     let pin = read_pin(pin);
     let epk_bytes = subject_public_key_info.subject_public_key.as_ref();
 
     let mut yk = opt_result!(YubiKey::open(), "YubiKey not found: {}");
-    let retired_slot_id = opt_result!(RetiredSlotId::from_str(slot), "Slot not found: {}");
+    let retired_slot_id = opt_result!(RetiredSlotId::from_str(&slot), "Slot not found: {}");
     let slot_id = SlotId::Retired(retired_slot_id);
     opt_result!(yk.verify_pin(pin.as_bytes()), "YubiKey verify pin failed: {}");
     let decrypted_shared_secret = opt_result!(decrypt_data(
@@ -176,6 +173,23 @@ fn try_decrypt_key_pgp(envelop: &TinyEncryptEnvelop, pin: &Option<String>) -> XR
     Ok(key)
 }
 
+fn read_slot(slot: &Option<String>) -> XResult<String> {
+    match slot {
+        Some(slot) => Ok(slot.to_string()),
+        None => {
+            print!("Input slot(eg 82, 83 ...): ");
+            io::stdout().flush().ok();
+            let mut buff = String::new();
+            let _ = io::stdin().read_line(&mut buff).expect("Read line from stdin");
+            if buff.is_empty() {
+                simple_error!("Slot is required, and not inputed")
+            } else {
+                Ok(buff)
+            }
+        }
+    }
+}
+
 fn read_pin(pin: &Option<String>) -> String {
     match pin {
         Some(pin) => pin.to_string(),