From b59d91db06278338bbef3f7569994ff50fe8c31d Mon Sep 17 00:00:00 2001
From: Halimao <1065621723@qq.com>
Date: Tue, 28 Nov 2023 17:00:20 +0800
Subject: [PATCH] feat: support display private key discreetly

---
 Cargo.lock                  | 38 ++++++++++++++++++++++++++++++++++++-
 Cargo.toml                  |  3 +++
 leo/cli/commands/account.rs | 27 ++++++++++++++++++++++++--
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4704d76bfa..9a0a28264e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1448,6 +1448,7 @@ dependencies = [
  "serde_json",
  "snarkvm",
  "sys-info",
+ "termion",
  "test_dir",
  "toml 0.8.8",
  "tracing",
@@ -1549,6 +1550,17 @@ dependencies = [
  "redox_syscall",
 ]
 
+[[package]]
+name = "libredox"
+version = "0.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3af92c55d7d839293953fcd0fda5ecfe93297cfde6ffbdec13b41d99c0ba6607"
+dependencies = [
+ "bitflags 2.4.1",
+ "libc",
+ "redox_syscall",
+]
+
 [[package]]
 name = "libz-sys"
 version = "1.1.12"
@@ -1757,6 +1769,12 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3"
 
+[[package]]
+name = "numtoa"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b8f8bdf33df195859076e54ab11ee78a1b208382d3a26ec40d142ffc1ecc49ef"
+
 [[package]]
 name = "object"
 version = "0.32.1"
@@ -2101,6 +2119,12 @@ dependencies = [
  "bitflags 1.3.2",
 ]
 
+[[package]]
+name = "redox_termios"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "20145670ba436b55d91fc92d25e71160fbfbdd57831631c8d7d36377a476f1cb"
+
 [[package]]
 name = "redox_users"
 version = "0.4.4"
@@ -2108,7 +2132,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a18479200779601e498ada4e8c1e1f50e3ee19deb0259c25825a98b5603b2cb4"
 dependencies = [
  "getrandom",
- "libredox",
+ "libredox 0.0.1",
  "thiserror",
 ]
 
@@ -3522,6 +3546,18 @@ dependencies = [
  "winapi-util",
 ]
 
+[[package]]
+name = "termion"
+version = "2.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c4648c7def6f2043b2568617b9f9b75eae88ca185dbc1f1fda30e95a85d49d7d"
+dependencies = [
+ "libc",
+ "libredox 0.0.2",
+ "numtoa",
+ "redox_termios",
+]
+
 [[package]]
 name = "termtree"
 version = "0.4.1"
diff --git a/Cargo.toml b/Cargo.toml
index 1b72c1fa2d..c594e2bbd7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -157,6 +157,9 @@ features = [ "fmt" ]
 [dependencies.zip]
 version = "^0.6"
 
+[dependencies.termion]
+version = "2.0.3"
+
 [target."cfg(windows)".dependencies.ansi_term]
 version = "0.12.1"
 
diff --git a/leo/cli/commands/account.rs b/leo/cli/commands/account.rs
index eea0ed6883..0e196c0d6f 100644
--- a/leo/cli/commands/account.rs
+++ b/leo/cli/commands/account.rs
@@ -20,6 +20,7 @@ use snarkvm::prelude::{Address, PrivateKey, ViewKey};
 
 use rand::SeedableRng;
 use rand_chacha::ChaChaRng;
+use std::io::{Read, Write};
 
 /// Commands to manage Aleo accounts.
 #[derive(Parser, Debug)]
@@ -98,9 +99,12 @@ fn print_keys(private_key: PrivateKey<CurrentNetwork>) -> Result<()> {
     let view_key = ViewKey::try_from(&private_key)?;
     let address = Address::<CurrentNetwork>::try_from(&view_key)?;
 
+    display_string_discreetly(
+        &private_key.to_string(), 
+        "### Do not share or lose this private key! Press any key to complete. ###",
+    )?;
     println!(
-        "\n {:>12}  {private_key}\n {:>12}  {view_key}\n {:>12}  {address}\n",
-        "Private Key".cyan().bold(),
+        "\n {:>12}  {view_key}\n {:>12}  {address}\n",
         "View Key".cyan().bold(),
         "Address".cyan().bold(),
     );
@@ -115,3 +119,22 @@ fn write_to_env_file(private_key: PrivateKey<CurrentNetwork>, ctx: &Context) ->
     tracing::info!("✅ Private Key written to {}", program_dir.join(".env").display());
     Ok(())
 }
+
+/// Print the string to an alternate screen, so that the string won't been printed to the terminal.
+fn display_string_discreetly(
+    discreet_string: &str,
+    continue_message: &str,
+) -> Result<()> {
+    use termion::screen::IntoAlternateScreen;
+    let mut screen = std::io::stdout().into_alternate_screen().unwrap();
+    writeln!(screen, "{discreet_string}").unwrap();
+    screen.flush().unwrap();
+    println!("\n{continue_message}");
+    wait_for_keypress();
+    Ok(())
+}
+
+fn wait_for_keypress() {
+    let mut single_key = [0u8];
+    std::io::stdin().read_exact(&mut single_key).unwrap();
+}
\ No newline at end of file