drk2: finished up wallet subcommand functionality

bin/drk2/Cargo.toml

@@ -10,13 +10,14 @@ edition = "2021"
 
 [dependencies]
 # Darkfi
-darkfi = {path = "../../", features = ["async-daemonize", "rpc"]}
+darkfi = {path = "../../", features = ["async-daemonize", "bs58", "rpc"]}
 darkfi_money_contract = {path = "../../src/contract/money", features = ["no-entrypoint", "client"]}
 darkfi_dao_contract = {path = "../../src/contract/dao", features = ["no-entrypoint", "client"]}
 darkfi-sdk = {path = "../../src/sdk", features = ["async"]}
 darkfi-serial = {path = "../../src/serial"}
 
 # Misc
+bs58 = "0.5.0"
 log = "0.4.20"
 prettytable-rs = "0.10.0"
 rand = "0.8.5"

bin/drk2/src/main.rs

@@ -16,7 +16,7 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use std::{fs, process::exit, sync::Arc, time::Instant};
+use std::{fs, io::stdin, process::exit, sync::Arc, time::Instant};
 
 use prettytable::{format, row, Table};
 use smol::stream::StreamExt;
@@ -29,6 +29,8 @@ use darkfi::{
     util::{parse::encode_base10, path::expand_path},
     Result,
 };
+use darkfi_sdk::pasta::pallas;
+use darkfi_serial::{deserialize, serialize};
 
 /// Error codes
 mod error;
@@ -191,8 +193,8 @@ impl Drk {
         let req = JsonRequest::new("ping", JsonValue::Array(vec![]));
         let rep = self.rpc_client.oneshot_request(req).await?;
         let latency = latency.elapsed();
-        eprintln!("Got reply: {:?}", rep);
-        eprintln!("Latency: {:?}", latency);
+        eprintln!("Got reply: {rep:?}");
+        eprintln!("Latency: {latency:?}");
         Ok(())
     }
 }
@@ -284,9 +286,9 @@ async fn realmain(args: Args, ex: Arc<smol::Executor<'static>>) -> Result<()> {
                 }
 
                 if table.is_empty() {
-                    println!("No unspent balances found");
+                    eprintln!("No unspent balances found");
                 } else {
-                    println!("{table}");
+                    eprintln!("{table}");
                 }
 
                 return Ok(())
@@ -301,7 +303,7 @@ async fn realmain(args: Args, ex: Arc<smol::Executor<'static>>) -> Result<()> {
                     }
                 };
 
-                println!("{address}");
+                eprintln!("{address}");
 
                 return Ok(())
             }
@@ -322,9 +324,9 @@ async fn realmain(args: Args, ex: Arc<smol::Executor<'static>>) -> Result<()> {
                 }
 
                 if table.is_empty() {
-                    println!("No addresses found");
+                    eprintln!("No addresses found");
                 } else {
-                    println!("{table}");
+                    eprintln!("{table}");
                 }
 
                 return Ok(())
@@ -338,9 +340,109 @@ async fn realmain(args: Args, ex: Arc<smol::Executor<'static>>) -> Result<()> {
                 return Ok(())
             }
 
-            // TODO
+            if secrets {
+                let v = drk.get_money_secrets().await?;
 
-            Ok(())
+                for i in v {
+                    eprintln!("{i}");
+                }
+
+                return Ok(())
+            }
+
+            if import_secrets {
+                let mut secrets = vec![];
+                let lines = stdin().lines();
+                for (i, line) in lines.enumerate() {
+                    if let Ok(line) = line {
+                        let bytes = bs58::decode(&line.trim()).into_vec()?;
+                        let Ok(secret) = deserialize(&bytes) else {
+                            eprintln!("Warning: Failed to deserialize secret on line {}", i);
+                            continue
+                        };
+                        secrets.push(secret);
+                    }
+                }
+
+                let pubkeys = match drk.import_money_secrets(secrets).await {
+                    Ok(p) => p,
+                    Err(e) => {
+                        eprintln!("Failed to import secret keys into wallet: {e:?}");
+                        exit(2);
+                    }
+                };
+
+                for key in pubkeys {
+                    eprintln!("{key}");
+                }
+
+                return Ok(())
+            }
+
+            if tree {
+                let tree = drk.get_money_tree().await?;
+
+                eprintln!("{tree:#?}");
+
+                return Ok(())
+            }
+
+            if coins {
+                let coins = drk.get_coins(true).await?;
+
+                let aliases_map = drk.get_aliases_mapped_by_token().await?;
+
+                if coins.is_empty() {
+                    return Ok(())
+                }
+
+                let mut table = Table::new();
+                table.set_format(*format::consts::FORMAT_NO_BORDER_LINE_SEPARATOR);
+                table.set_titles(row![
+                    "Coin",
+                    "Spent",
+                    "Token ID",
+                    "Aliases",
+                    "Value",
+                    "Spend Hook",
+                    "User Data"
+                ]);
+                let zero = pallas::Base::zero();
+                for coin in coins {
+                    let aliases = match aliases_map.get(&coin.0.note.token_id.to_string()) {
+                        Some(a) => a,
+                        None => "-",
+                    };
+
+                    let spend_hook = if coin.0.note.spend_hook != zero {
+                        bs58::encode(&serialize(&coin.0.note.spend_hook)).into_string().to_string()
+                    } else {
+                        String::from("-")
+                    };
+
+                    let user_data = if coin.0.note.user_data != zero {
+                        bs58::encode(&serialize(&coin.0.note.user_data)).into_string().to_string()
+                    } else {
+                        String::from("-")
+                    };
+
+                    table.add_row(row![
+                        bs58::encode(&serialize(&coin.0.coin.inner())).into_string().to_string(),
+                        coin.1,
+                        coin.0.note.token_id,
+                        aliases,
+                        format!("{} ({})", coin.0.note.value, encode_base10(coin.0.note.value, 8)),
+                        spend_hook,
+                        user_data
+                    ]);
+                }
+
+                println!("{table}");
+
+                return Ok(())
+            }
+
+            unreachable!()
         }
     }
 }

bin/drk2/src/money.rs

@@ -58,7 +58,7 @@ impl Drk {
         // new tree and push it into the table.
         // For now, on success, we don't care what's returned, but in the future
         // we should actually check it.
-        if self.wallet.query_single(MONEY_TREE_TABLE, &[MONEY_TREE_COL_TREE], &[]).await.is_err() {
+        if self.get_money_tree().await.is_err() {
             eprintln!("Initializing Money Merkle tree");
             let mut tree = MerkleTree::new(100);
             tree.append(MerkleNode::from(pallas::Base::ZERO));
@@ -107,7 +107,7 @@ impl Drk {
             .await?;
 
         eprintln!("New address:");
-        println!("{}", keypair.public);
+        eprintln!("{}", keypair.public);
 
         Ok(())
     }
@@ -198,6 +198,74 @@ impl Drk {
         Ok(vec)
     }
 
+    /// Fetch all secret keys from the wallet
+    pub async fn get_money_secrets(&self) -> Result<Vec<SecretKey>> {
+        let rows =
+            match self.wallet.query_multiple(MONEY_KEYS_TABLE, &[MONEY_KEYS_COL_SECRET], &[]).await
+            {
+                Ok(r) => r,
+                Err(e) => {
+                    return Err(Error::RusqliteError(format!(
+                        "[get_money_secrets] Secret keys retrieval failed: {e:?}"
+                    )))
+                }
+            };
+
+        let mut secrets = Vec::with_capacity(rows.len());
+
+        // Let's scan through the rows and see if we got anything.
+        for row in rows {
+            let Value::Blob(ref key_bytes) = row[0] else {
+                return Err(Error::ParseFailed(
+                    "[get_money_secrets] Secret key bytes parsing failed",
+                ))
+            };
+            let secret_key: SecretKey = deserialize(key_bytes)?;
+            secrets.push(secret_key);
+        }
+
+        Ok(secrets)
+    }
+
+    /// Import given secret keys into the wallet.
+    /// If the key already exists, it will be skipped.
+    /// Returns the respective PublicKey objects for the imported keys.
+    pub async fn import_money_secrets(&self, secrets: Vec<SecretKey>) -> Result<Vec<PublicKey>> {
+        let existing_secrets = self.get_money_secrets().await?;
+
+        let mut ret = Vec::with_capacity(secrets.len());
+
+        for secret in secrets {
+            // Check if secret already exists
+            if existing_secrets.contains(&secret) {
+                eprintln!("Existing key found: {secret}");
+                continue
+            }
+
+            ret.push(PublicKey::from_secret(secret));
+            let is_default = 0;
+            let public = serialize(&PublicKey::from_secret(secret));
+            let secret = serialize(&secret);
+
+            let query = format!(
+                "INSERT INTO {} ({}, {}, {}) VALUES (?1, ?2, ?3);",
+                MONEY_KEYS_TABLE,
+                MONEY_KEYS_COL_IS_DEFAULT,
+                MONEY_KEYS_COL_PUBLIC,
+                MONEY_KEYS_COL_SECRET
+            );
+            if let Err(e) =
+                self.wallet.exec_sql(&query, rusqlite::params![is_default, public, secret]).await
+            {
+                return Err(Error::RusqliteError(format!(
+                    "[import_money_secrets] Inserting new address failed: {e:?}"
+                )))
+            }
+        }
+
+        Ok(ret)
+    }
+
     /// Fetch known unspent balances from the wallet and return them as a hashmap.
     pub async fn money_balance(&self) -> Result<HashMap<String, u64>> {
         let mut coins = self.get_coins(false).await?;
@@ -402,6 +470,25 @@ impl Drk {
         self.wallet.exec_sql(&query, rusqlite::params![serialize(tree)]).await
     }
 
+    /// Fetch the Money Merkle tree from the wallet
+    pub async fn get_money_tree(&self) -> Result<MerkleTree> {
+        let row =
+            match self.wallet.query_single(MONEY_TREE_TABLE, &[MONEY_TREE_COL_TREE], &[]).await {
+                Ok(r) => r,
+                Err(e) => {
+                    return Err(Error::RusqliteError(format!(
+                        "[get_money_tree] Tree retrieval failed: {e:?}"
+                    )))
+                }
+            };
+
+        let Value::Blob(ref tree_bytes) = row[0] else {
+            return Err(Error::ParseFailed("[get_money_tree] Tree bytes parsing failed"))
+        };
+        let tree = deserialize(tree_bytes)?;
+        Ok(tree)
+    }
+
     /// Get the last scanned slot from the wallet
     pub async fn last_scanned_slot(&self) -> WalletDbResult<u64> {
         let ret = self