CLI updates

ant-cli/src/commands.rs

@@ -158,14 +158,26 @@ pub enum WalletCmd {
         /// Optional flag to not add a password.
         #[clap(long, action)]
         no_password: bool,
-        /// Optional hex-encoded private key.
-        #[clap(long)]
-        private_key: Option<String>,
         /// Optional password to encrypt the wallet with.
         #[clap(long, short)]
         password: Option<String>,
     },
 
+    /// Import an existing wallet.
+    Import {
+        /// Hex-encoded private key.
+        private_key: String,
+        /// Optional flag to not add a password.
+        #[clap(long, action)]
+        no_password: bool,
+        /// Optional password to encrypt the wallet with.
+        #[clap(long, short)]
+        password: Option<String>,
+    },
+
+    /// Print the private key of a wallet.
+    Export,
+
     /// Check the balance of the wallet.
     Balance,
 }
@@ -208,10 +220,15 @@ pub async fn handle_subcommand(opt: Opt) -> Result<()> {
         SubCmd::Wallet { command } => match command {
             WalletCmd::Create {
                 no_password,
+                password,
+            } => wallet::create(no_password, password),
+            WalletCmd::Import {
                 private_key,
+                no_password,
                 password,
-            } => wallet::create(no_password, private_key, password),
-            WalletCmd::Balance => Ok(wallet::balance().await?),
+            } => wallet::import(private_key, no_password, password),
+            WalletCmd::Export => wallet::export(),
+            WalletCmd::Balance => wallet::balance().await,
         },
     }
 }

ant-cli/src/commands/wallet.rs

@@ -6,7 +6,7 @@
 // KIND, either express or implied. Please review the Licences for the specific language governing
 // permissions and limitations relating to use of the SAFE Network Software.
 
-use crate::wallet::fs::{select_wallet, store_private_key};
+use crate::wallet::fs::{select_wallet, select_wallet_private_key, store_private_key};
 use crate::wallet::input::request_password;
 use crate::wallet::DUMMY_NETWORK;
 use autonomi::Wallet;
@@ -16,49 +16,70 @@ use prettytable::{Cell, Row, Table};
 
 const WALLET_PASSWORD_REQUIRED: bool = false;
 
-pub fn create(
+pub fn create(no_password: bool, password: Option<String>) -> Result<()> {
+    let maybe_encryption_password = maybe_request_password(no_password, password)?;
+
+    let wallet_private_key = Wallet::random_private_key();
+
+    let wallet_address = Wallet::new_from_private_key(DUMMY_NETWORK, &wallet_private_key)
+        .expect("Infallible")
+        .address()
+        .to_string();
+
+    // Save the private key file
+    let file_path = store_private_key(&wallet_private_key, maybe_encryption_password)?;
+
+    println!("Wallet address: {wallet_address}");
+    println!("Wallet private key: {wallet_private_key}");
+    println!("Stored wallet in: {file_path:?}");
+
+    Ok(())
+}
+
+pub fn import(
+    mut wallet_private_key: String,
     no_password: bool,
-    private_key: Option<String>,
     password: Option<String>,
 ) -> Result<()> {
-    if no_password && password.is_some() {
-        return Err(eyre!(
-            "Only one of `--no-password` or `--password` may be specified"
-        ));
-    }
-
-    // Set a password for encryption or not
-    let encryption_password: Option<String> = match (no_password, password) {
-        (true, _) => None,
-        (false, Some(pass)) => Some(pass.to_owned()),
-        (false, None) => request_password(WALLET_PASSWORD_REQUIRED),
-    };
-
-    let wallet_private_key = if let Some(private_key) = private_key {
-        // Validate imported key
-        Wallet::new_from_private_key(DUMMY_NETWORK, &private_key)
-            .map_err(|_| eyre!("Please provide a valid secret key in hex format"))?;
+    // Validate imported key
+    Wallet::new_from_private_key(DUMMY_NETWORK, &wallet_private_key)
+        .map_err(|_| eyre!("Please provide a valid private key in hex format"))?;
 
-        private_key
-    } else {
-        // Create a new key
-        Wallet::random_private_key()
-    };
+    let maybe_encryption_password = maybe_request_password(no_password, password)?;
 
     let wallet_address = Wallet::new_from_private_key(DUMMY_NETWORK, &wallet_private_key)
         .expect("Infallible")
         .address()
         .to_string();
 
+    // Prepend with 0x if it isn't already
+    if !wallet_private_key.starts_with("0x") {
+        wallet_private_key = format!("0x{wallet_private_key}");
+    }
+
     // Save the private key file
-    let file_path = store_private_key(&wallet_private_key, encryption_password)?;
+    let file_path = store_private_key(&wallet_private_key, maybe_encryption_password)?;
 
     println!("Wallet address: {wallet_address}");
     println!("Stored wallet in: {file_path:?}");
 
     Ok(())
 }
 
+pub fn export() -> Result<()> {
+    let wallet_private_key = select_wallet_private_key()?;
+
+    let wallet_address = Wallet::new_from_private_key(DUMMY_NETWORK, &wallet_private_key)
+        .expect("Infallible")
+        .address()
+        .to_string();
+
+    println!("Wallet address: {wallet_address}");
+    println!("Wallet private key: {wallet_private_key}");
+
+    Ok(())
+}
+
 pub async fn balance() -> Result<()> {
     let wallet = select_wallet()?;
 
@@ -83,3 +104,20 @@ pub async fn balance() -> Result<()> {
 
     Ok(())
 }
+
+fn maybe_request_password(no_password: bool, password: Option<String>) -> Result<Option<String>> {
+    if no_password && password.is_some() {
+        return Err(eyre!(
+            "Only one of `--no-password` or `--password` may be specified"
+        ));
+    }
+
+    // Set a password for encryption or not
+    let maybe_password = match (no_password, password) {
+        (true, _) => None,
+        (false, Some(pass)) => Some(pass.to_owned()),
+        (false, None) => request_password(WALLET_PASSWORD_REQUIRED),
+    };
+
+    Ok(maybe_password)
+}

ant-cli/src/wallet/encryption.rs

@@ -123,11 +123,8 @@ pub fn decrypt_private_key(encrypted_data: &str, password: &str) -> Result<Strin
             ))
         })?;
 
-    let mut private_key_bytes = [0u8; 66];
-    private_key_bytes.copy_from_slice(&decrypted_data[0..66]);
-
     // Create secret key from decrypted byte
-    Ok(String::from_utf8(private_key_bytes.to_vec()).expect("not able to convert private key"))
+    Ok(String::from_utf8(decrypted_data.to_vec()).expect("not able to convert private key"))
 }
 
 #[cfg(test)]

ant-cli/src/wallet/fs.rs

@@ -82,7 +82,10 @@ pub(crate) fn load_private_key(wallet_address: &str) -> Result<String, Error> {
     let encrypted_file_path =
         wallets_folder.join(format!("{wallet_address}{ENCRYPTED_PRIVATE_KEY_EXT}"));
 
-    let is_encrypted = encrypted_file_path.exists();
+    let is_plain = wallets_folder.join(&file_name).exists();
+
+    // Trick to favour the plain file in case they both exist
+    let is_encrypted = encrypted_file_path.exists() && !is_plain;
 
     if is_encrypted {
         file_name.push_str(ENCRYPTED_PRIVATE_KEY_EXT);