From 7bc7d9448512f49d78e0b46ca87eb3556282f60e Mon Sep 17 00:00:00 2001
From: Lennart <lennart@1200wd.com>
Date: Thu, 9 Nov 2023 23:16:43 +0100
Subject: [PATCH] [ADD] Use multiple witnesses in multisig wallet

---
 bitcoinlib/wallets.py | 17 ++++++++++-------
 tests/test_wallets.py |  6 ++++--
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/bitcoinlib/wallets.py b/bitcoinlib/wallets.py
index ced04dfd..9d8358d4 100644
--- a/bitcoinlib/wallets.py
+++ b/bitcoinlib/wallets.py
@@ -1672,7 +1672,8 @@ def import_key(self, key, account_id=0, name='', network=None, purpose=84, key_t
                     return w.import_master_key(hdkey)
             raise WalletError("Unknown key: Can only import a private key for a known public key in multisig wallets")
 
-    def _new_key_multisig(self, public_keys, name, account_id, change, cosigner_id, network, address_index):
+    def _new_key_multisig(self, public_keys, name, account_id, change, cosigner_id, network, address_index,
+                          witness_type):
         if self.sort_keys:
             public_keys.sort(key=lambda pubk: pubk.key_public)
         public_key_list = [pubk.key_public for pubk in public_keys]
@@ -1685,9 +1686,9 @@ def _new_key_multisig(self, public_keys, name, account_id, change, cosigner_id,
         redeemscript = Script(script_types=['multisig'], keys=public_key_list,
                               sigs_required=self.multisig_n_required).serialize()
         script_type = 'p2sh'
-        if self.witness_type == 'p2sh-segwit':
+        if witness_type == 'p2sh-segwit':
             script_type = 'p2sh_p2wsh'
-        address = Address(redeemscript, encoding=self.encoding, script_type=script_type, network=network)
+        address = Address(redeemscript, script_type=script_type, network=network, witness_type=witness_type)
         already_found_key = self._session.query(DbKey).filter_by(wallet_id=self.wallet_id,
                                                                  address=address.address).first()
         if already_found_key:
@@ -1701,7 +1702,7 @@ def _new_key_multisig(self, public_keys, name, account_id, change, cosigner_id,
             name=name[:80], wallet_id=self.wallet_id, purpose=self.purpose, account_id=account_id,
             depth=depth, change=change, address_index=address_index, parent_id=0, is_private=False, path=path,
             public=address.hash_bytes, wif='multisig-%s' % address, address=address.address, cosigner_id=cosigner_id,
-            key_type='multisig', witness_type=self.witness_type, network_name=network)
+            key_type='multisig', witness_type=witness_type, network_name=network)
         self._session.add(multisig_key)
         self._commit()
         for child_id in public_key_ids:
@@ -2152,7 +2153,7 @@ def key_for_path(self, path, level_offset=None, name=None, account_id=None, cosi
             level_offset_key = level_offset - self.main_key.depth
         witness_type = witness_type if witness_type else self.witness_type
         if ((not self.main_key or not self.main_key.is_private or self.main_key.depth != 0) and
-                self.witness_type != witness_type):
+                self.witness_type != witness_type) and not self.multisig:
             raise WalletError("This wallet has no private key, cannot use multiple witness types")
         key_path = self.key_path
         purpose = self.purpose
@@ -2172,9 +2173,11 @@ def key_for_path(self, path, level_offset=None, name=None, account_id=None, cosi
                     wk = wlt.main_key
                 else:
                     wk = wlt.key_for_path(path, level_offset=level_offset, account_id=account_id, name=name,
-                                          cosigner_id=cosigner_id, network=network, recreate=recreate)
+                                          cosigner_id=cosigner_id, network=network, recreate=recreate,
+                                          witness_type=witness_type)
                 public_keys.append(wk)
-            return self._new_key_multisig(public_keys, name, account_id, change, cosigner_id, network, address_index)
+            return self._new_key_multisig(public_keys, name, account_id, change, cosigner_id, network, address_index,
+                                          witness_type)
 
         # Check for closest ancestor in wallet
         wpath = fullpath
diff --git a/tests/test_wallets.py b/tests/test_wallets.py
index c666d4c3..bbb1b519 100644
--- a/tests/test_wallets.py
+++ b/tests/test_wallets.py
@@ -2570,8 +2570,10 @@ def test_wallet_segwit_multiple_networks_accounts(self):
         self.assertListEqual(sorted([nw.name for nw in wallet.networks()]), networks_expected)
         self.assertListEqual([k.path for k in wallet.keys_accounts(network='litecoin')],
                              ["m/48'/2'/0'/2'", "m/48'/2'/1'/2'"])
-        self.assertEqual(wallet.keys(network='litecoin')[0].address, "MQNA8FYrN2fvD7SSYny3Ccvpapvsu9cVJH")
-        self.assertEqual(wallet.keys(network='bitcoin')[0].address, "3L6XFzC6RPeXSFpZS8v4S86v4gsNmKFnFT")
+        self.assertEqual(wallet.keys(network='litecoin')[0].address,
+                         "ltc1qcu83hcdwy46dv85vmwnnfgzeu8d9arfy7kfyct52dxqcyvq2q6ds5kq2ah")
+        self.assertEqual(wallet.keys(network='bitcoin')[0].address,
+                         "bc1qd0f952amxkmqc9e60u4g8w4r5a3cx22lt5vqeeyljllacxq8ezusclkwa0")
 
 
 @parameterized_class(*params)