From 1bdbcbc6430408f4562b766bbfc3c26b5e17e200 Mon Sep 17 00:00:00 2001
From: benk10 <ben.kaufman10@gmail.com>
Date: Tue, 3 Sep 2024 20:57:30 +0400
Subject: [PATCH] Add prompt_pin and send_pin HWI command

---
 src/interface.rs | 30 ++++++++++++++++++++++++++++++
 src/lib.rs       | 39 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/src/interface.rs b/src/interface.rs
index e324bbe..4f40dc0 100644
--- a/src/interface.rs
+++ b/src/interface.rs
@@ -497,6 +497,36 @@ impl HWIClient {
         })
     }
 
+    /// Prompt PIN on a device
+    pub fn prompt_pin(&self) -> Result<(), Error> {
+        Python::with_gil(|py| {
+            let func_args = (&self.hw_client,);
+            let output = self
+                .hwilib
+                .commands
+                .getattr(py, "prompt_pin")?
+                .call1(py, func_args)?;
+            let output = self.hwilib.json_dumps.call1(py, (output,))?;
+            let status: HWIStatus = deserialize_obj!(&output.to_string())?;
+            status.into()
+        })
+    }
+
+    /// Send PIN to a device
+    pub fn send_pin(&self, pin: &str) -> Result<(), Error> {
+        Python::with_gil(|py| {
+            let func_args = (&self.hw_client, pin);
+            let output = self
+                .hwilib
+                .commands
+                .getattr(py, "send_pin")?
+                .call1(py, func_args)?;
+            let output = self.hwilib.json_dumps.call1(py, (output,))?;
+            let status: HWIStatus = deserialize_obj!(&output.to_string())?;
+            status.into()
+        })
+    }
+
     /// Get the installed version of hwilib. Returns None if hwi is not installed.
     pub fn get_version() -> Option<String> {
         Python::with_gil(|py| {
diff --git a/src/lib.rs b/src/lib.rs
index d549e92..5a75d6e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -41,12 +41,12 @@ pub mod types;
 
 #[cfg(test)]
 mod tests {
-    use crate::types::{self, HWIDeviceType, TESTNET};
+    use crate::types::{self, HWIChain, HWIDevice, HWIDeviceType, TESTNET};
     use crate::HWIClient;
     use std::collections::BTreeMap;
     use std::str::FromStr;
 
-    use bitcoin::bip32::{DerivationPath, KeySource};
+    use bitcoin::bip32::{DerivationPath, Fingerprint, KeySource};
     use bitcoin::locktime::absolute;
     use bitcoin::psbt::{Input, Output};
     use bitcoin::{secp256k1, Transaction};
@@ -435,6 +435,41 @@ mod tests {
         }
     }
 
+    #[test]
+    #[serial]
+    #[ignore = "Needs a Trezor One device for the test"]
+    fn test_prompt_pin_to_trezor_device() {
+        let client = HWIClient::find_device(
+            None,
+            Some(HWIDeviceType::Trezor),
+            None,
+            false,
+            Network::Testnet,
+        )
+        .unwrap();
+        client.prompt_pin().unwrap();
+    }
+
+    #[test]
+    #[serial]
+    #[ignore = "Needs a Trezor One device for the test"]
+    fn test_send_pin_to_trezor_device() {
+        let client = HWIClient::get_client(
+            &HWIDevice {
+                device_type: HWIDeviceType::Trezor,
+                model: "trezor_1".to_string(),
+                path: "webusb:000:1:2".to_string(),
+                needs_pin_sent: true,
+                needs_passphrase_sent: false,
+                fingerprint: Fingerprint::from_str("00000000").unwrap(),
+            },
+            false,
+            HWIChain::from(Network::Testnet),
+        )
+        .unwrap();
+        client.send_pin("123456").unwrap();
+    }
+
     #[test]
     #[serial]
     #[ignore]