Просмотр исходного кода

drk: Add secret key wallet import functionality.

parazyd 3 лет назад
Родитель
Сommit
682136f823
2 измененных файлов с 92 добавлено и 2 удалено
  1. 51 2
      bin/drk/src/main.rs
  2. 41 0
      bin/drk/src/rpc_wallet.rs

+ 51 - 2
bin/drk/src/main.rs

@@ -104,6 +104,10 @@ enum Subcmd {
         /// Print all the secret keys from the wallet
         secrets: bool,
 
+        #[arg(long)]
+        /// Import secret keys from stdin into the wallet, separated by newlines
+        import_secrets: bool,
+
         #[arg(long)]
         /// Print the Merkle tree in the wallet
         tree: bool,
@@ -229,8 +233,25 @@ async fn main() -> Result<()> {
             Ok(())
         }
 
-        Subcmd::Wallet { initialize, keygen, balance, address, secrets, tree, coins } => {
-            if !initialize && !keygen && !balance && !address && !secrets && !tree && !coins {
+        Subcmd::Wallet {
+            initialize,
+            keygen,
+            balance,
+            address,
+            secrets,
+            import_secrets,
+            tree,
+            coins,
+        } => {
+            if !initialize &&
+                !keygen &&
+                !balance &&
+                !address &&
+                !secrets &&
+                !tree &&
+                !coins &&
+                !import_secrets
+            {
                 eprintln!("Error: You must use at least one flag for this subcommand");
                 eprintln!("Run with \"wallet -h\" to see the subcommand usage.");
                 exit(2);
@@ -281,6 +302,34 @@ async fn main() -> Result<()> {
                 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 = drk
+                    .wallet_import_secrets(secrets)
+                    .await
+                    .with_context(|| "Failed to import secret keys into wallet")?;
+
+                drk.rpc_client.close().await?;
+
+                for key in pubkeys {
+                    println!("{}", key);
+                }
+
+                return Ok(())
+            }
+
             if tree {
                 let v = drk.wallet_tree().await.with_context(|| "Failed to fetch Merkle tree")?;
                 drk.rpc_client.close().await?;

+ 41 - 0
bin/drk/src/rpc_wallet.rs

@@ -360,6 +360,47 @@ impl Drk {
         Ok(secrets)
     }
 
+    /// Import given secret keys into the wallet. The query uses INSERT, so if the key already
+    /// exists, it will simply be skipped.
+    pub async fn wallet_import_secrets(&self, secrets: Vec<SecretKey>) -> Result<Vec<PublicKey>> {
+        let mut ret = vec![];
+
+        for secret in secrets {
+            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,
+            );
+
+            let params = json!([
+                query,
+                QueryType::Integer as u8,
+                is_default,
+                QueryType::Blob as u8,
+                public,
+                QueryType::Blob as u8,
+                secret,
+            ]);
+
+            let req = JsonRequest::new("wallet.exec_sql", params);
+            let rep = self.rpc_client.request(req).await?;
+
+            if rep != true {
+                // Something weird happened?
+                eprintln!("Got unexpected reply from darkfid: {}", rep);
+            }
+        }
+
+        Ok(ret)
+    }
+
     /// Get the Merkle tree from the wallet
     pub async fn wallet_tree(&self) -> Result<BridgeTree<MerkleNode, MERKLE_DEPTH>> {
         let query = format!("SELECT * FROM {}", MONEY_TREE_TABLE);