From 308c6561bc23b1d2b567ad6edbd8cb1246e7da91 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Mon, 9 Oct 2023 22:28:16 +0800
Subject: [PATCH] feat: find kid first arg

---
 README.md          |  6 +++---
 src/cmd_decrypt.rs | 10 +++-------
 src/config.rs      |  8 ++++++++
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 8c50432..b1cc7c2 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,7 @@
 # tiny-encrypt-rs
 
-**IMPORTANT**: To use tiny-encrypt, a Yubikey(https://www.yubico.com/products/) or CanoKey(https://www.canokeys.org/) is
-required, the Key NEED supports PIV or OpenPGP.
+**IMPORTANT**: To use tiny-encrypt, a Yubikey(https://www.yubico.com/products/) is
+required, the key MUST support PIV or OpenPGP.
 
 ![](https://cdn.hatter.ink/doc/7684_4DB4452911E2A25AB993429AA7FFCD65/yubikey-5-family.png)
 
@@ -37,7 +37,7 @@ Encrypt config `~/.tinyencrypt/config-rs.json`:
       "KID-1",
       "KID-2"
     ],
-    "leve2": [
+    "l2,leve2": [
       "KID-2"
     ]
   }
diff --git a/src/cmd_decrypt.rs b/src/cmd_decrypt.rs
index 4847663..059367f 100644
--- a/src/cmd_decrypt.rs
+++ b/src/cmd_decrypt.rs
@@ -270,13 +270,9 @@ fn read_slot(config: &Option<TinyEncryptConfig>, kid: &str, slot: &Option<String
         Some(slot) => Ok(slot.to_string()),
         None => {
             if let Some(config) = config {
-                if let Some(e) = config.find_by_kid(kid) {
-                    if let Some(args) = &e.args {
-                        if let Some(first_arg) = args.iter().next() {
-                            information!("Found kid: {}'s slot: {}", kid, first_arg);
-                            return Ok(first_arg.to_string());
-                        }
-                    }
+                if let Some(first_arg) = config.find_first_arg_by_kid(kid) {
+                    information!("Found kid: {}'s slot: {}", kid, first_arg);
+                    return Ok(first_arg.to_string());
                 }
             }
             print!("Input slot(eg 82, 83 ...): ");
diff --git a/src/config.rs b/src/config.rs
index e1c1e16..0e2b6db 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -72,6 +72,14 @@ impl TinyEncryptConfig {
         Ok(config)
     }
 
+    pub fn find_first_arg_by_kid(&self, kid: &str) -> Option<&String> {
+        self.find_args_by_kid(kid).map(|a| a.iter().next()).flatten()
+    }
+
+    pub fn find_args_by_kid(&self, kid: &str) -> Option<&Vec<String>> {
+        self.find_by_kid(kid).map(|e| e.args.as_ref()).flatten()
+    }
+
     pub fn find_by_kid(&self, kid: &str) -> Option<&TinyEncryptConfigEnvelop> {
         self.envelops.iter().find(|e| e.kid == kid)
     }