From d0b5b9d2cd9a8afb96a3adc563669af52af1a54a Mon Sep 17 00:00:00 2001
From: Greg Martin <gmartin7t2@gmail.com>
Date: Sat, 4 Nov 2023 01:03:45 +0000
Subject: [PATCH] Add `--address-type` flag to `wallet create` and `wallet
 restore`.

---
 src/subcommand/preview.rs        |  1 +
 src/subcommand/wallet.rs         | 26 +++++++++++++++++++++++---
 src/subcommand/wallet/create.rs  |  4 +++-
 src/subcommand/wallet/restore.rs |  4 +++-
 4 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/subcommand/preview.rs b/src/subcommand/preview.rs
index d25ae0fbba..c4fb887fbb 100644
--- a/src/subcommand/preview.rs
+++ b/src/subcommand/preview.rs
@@ -63,6 +63,7 @@ impl Preview {
 
     super::wallet::Wallet::Create(super::wallet::create::Create {
       passphrase: "".into(),
+      address_type: super::wallet::AddressType::Bech32m,
     })
     .run(options.clone())?;
 
diff --git a/src/subcommand/wallet.rs b/src/subcommand/wallet.rs
index a526ac7b05..7d4d1c3e1a 100644
--- a/src/subcommand/wallet.rs
+++ b/src/subcommand/wallet.rs
@@ -53,6 +53,14 @@ pub(crate) enum Wallet {
   Cardinals,
 }
 
+#[derive(clap::ValueEnum, Clone, Debug)]
+pub(crate) enum AddressType {
+  Legacy,
+  P2SHSegwit,
+  Bech32,
+  Bech32m,
+}
+
 impl Wallet {
   pub(crate) fn run(self, options: Options) -> SubcommandResult {
     match self {
@@ -80,7 +88,7 @@ fn get_change_address(client: &Client, chain: Chain) -> Result<Address> {
   )
 }
 
-pub(crate) fn initialize_wallet(options: &Options, seed: [u8; 64]) -> Result {
+pub(crate) fn initialize_wallet(options: &Options, seed: [u8; 64], address_type: AddressType) -> Result {
   let client = options.bitcoin_rpc_client_for_wallet_command(true)?;
   let network = options.chain().network();
 
@@ -93,7 +101,12 @@ pub(crate) fn initialize_wallet(options: &Options, seed: [u8; 64]) -> Result {
   let fingerprint = master_private_key.fingerprint(&secp);
 
   let derivation_path = DerivationPath::master()
-    .child(ChildNumber::Hardened { index: 86 })
+    .child(ChildNumber::Hardened { index: match address_type {
+      AddressType::Legacy => 44,
+      AddressType::P2SHSegwit => 49,
+      AddressType::Bech32 => 84,
+      AddressType::Bech32m => 86,
+    } })
     .child(ChildNumber::Hardened {
       index: u32::from(network != Network::Bitcoin),
     })
@@ -108,6 +121,7 @@ pub(crate) fn initialize_wallet(options: &Options, seed: [u8; 64]) -> Result {
       (fingerprint, derivation_path.clone()),
       derived_private_key,
       change,
+      &address_type,
     )?;
   }
 
@@ -120,6 +134,7 @@ fn derive_and_import_descriptor(
   origin: (Fingerprint, DerivationPath),
   derived_private_key: ExtendedPrivKey,
   change: bool,
+  address_type: &AddressType,
 ) -> Result {
   let secret_key = DescriptorSecretKey::XPrv(DescriptorXKey {
     origin: Some(origin),
@@ -135,7 +150,12 @@ fn derive_and_import_descriptor(
   let mut key_map = std::collections::HashMap::new();
   key_map.insert(public_key.clone(), secret_key);
 
-  let desc = Descriptor::new_tr(public_key, None)?;
+  let desc = match address_type {
+    AddressType::Legacy => Descriptor::new_pkh(public_key),
+    AddressType::P2SHSegwit => Descriptor::new_sh_wpkh(public_key),
+    AddressType::Bech32 => Descriptor::new_wpkh(public_key),
+    AddressType::Bech32m => Descriptor::new_tr(public_key, None),
+  }?;
 
   client.import_descriptors(ImportDescriptors {
     descriptor: desc.to_string_with_secret(&key_map),
diff --git a/src/subcommand/wallet/create.rs b/src/subcommand/wallet/create.rs
index 34cd762450..a95b79015f 100644
--- a/src/subcommand/wallet/create.rs
+++ b/src/subcommand/wallet/create.rs
@@ -14,6 +14,8 @@ pub(crate) struct Create {
     help = "Use <PASSPHRASE> to derive wallet seed."
   )]
   pub(crate) passphrase: String,
+  #[arg(long, value_enum, default_value="bech32m")]
+  pub(crate) address_type: AddressType,
 }
 
 impl Create {
@@ -23,7 +25,7 @@ impl Create {
 
     let mnemonic = Mnemonic::from_entropy(&entropy)?;
 
-    initialize_wallet(&options, mnemonic.to_seed(self.passphrase.clone()))?;
+    initialize_wallet(&options, mnemonic.to_seed(self.passphrase.clone()), self.address_type)?;
 
     Ok(Box::new(Output {
       mnemonic,
diff --git a/src/subcommand/wallet/restore.rs b/src/subcommand/wallet/restore.rs
index 0d2f596ad0..7478f27a30 100644
--- a/src/subcommand/wallet/restore.rs
+++ b/src/subcommand/wallet/restore.rs
@@ -10,11 +10,13 @@ pub(crate) struct Restore {
     help = "Use <PASSPHRASE> when deriving wallet"
   )]
   pub(crate) passphrase: String,
+  #[arg(long, value_enum, default_value="bech32m")]
+  pub(crate) address_type: AddressType,
 }
 
 impl Restore {
   pub(crate) fn run(self, options: Options) -> SubcommandResult {
-    initialize_wallet(&options, self.mnemonic.to_seed(self.passphrase))?;
+    initialize_wallet(&options, self.mnemonic.to_seed(self.passphrase), self.address_type)?;
     Ok(Box::new(Empty {}))
   }
 }