build_psbt.py: provide witness vs redeem or both scripts

electrum/plugins/coldcard/build_psbt.py

@@ -13,13 +13,14 @@ from electrum.logging import get_logger
 from electrum.wallet import Standard_Wallet, Multisig_Wallet, Wallet
 from electrum.keystore import xpubkey_to_pubkey, Xpub
 from electrum.util import bfh, bh2u
-from electrum.crypto import hash_160
+from electrum.crypto import hash_160, sha256
 from electrum.bitcoin import DecodeBase58Check
 
 from .basic_psbt import (
         PSBT_GLOBAL_UNSIGNED_TX, PSBT_GLOBAL_XPUB, PSBT_IN_NON_WITNESS_UTXO, PSBT_IN_WITNESS_UTXO,
         PSBT_IN_SIGHASH_TYPE, PSBT_IN_REDEEM_SCRIPT, PSBT_IN_WITNESS_SCRIPT, PSBT_IN_PARTIAL_SIG,
-        PSBT_IN_BIP32_DERIVATION, PSBT_OUT_BIP32_DERIVATION, PSBT_OUT_REDEEM_SCRIPT)
+        PSBT_IN_BIP32_DERIVATION, PSBT_OUT_BIP32_DERIVATION,
+        PSBT_OUT_REDEEM_SCRIPT, PSBT_OUT_WITNESS_SCRIPT)
 from .basic_psbt import BasicPSBT
 
 from electrum.logging import get_logger
@@ -66,6 +67,7 @@ def unpacked_xfp_path(xfp, text_path):
     # - binary LE32 values, first one is the fingerprint
     # - but as ints, not bytes yet
     rv = [xfp]
+
     for x in text_path.split('/'):
         if x == 'm': continue
         if x.endswith("'"):
@@ -73,6 +75,7 @@ def unpacked_xfp_path(xfp, text_path):
         else:
             x = int(x)
         rv.append(x)
+
     return rv
 
 def xfp_for_keystore(ks):
@@ -88,7 +91,7 @@ def xfp_for_keystore(ks):
 def packed_xfp_path_for_keystore(ks, int_path=[]):
     # Return XFP + common prefix path for keystore, as binary ready for PSBT
     derv = getattr(ks, 'derivation', 'm')
-    return packed_xfp_path(xfp_for_keystore(ks), derv[2:], int_path=int_path)
+    return packed_xfp_path(xfp_for_keystore(ks), derv[2:] or 'm', int_path=int_path)
 
 # Serialization/deserialization tools
 def ser_compact_size(l):
@@ -273,8 +276,7 @@ def build_psbt(tx: Transaction, wallet: Wallet):
             if txin['type'] == 'p2wpkh-p2sh':
                 assert len(pubkeys) == 1, 'can be only one redeem script per input'
                 pa = hash_160(k)
-                assert len(pa) == 20
-                write_kv(PSBT_IN_REDEEM_SCRIPT, b'\x00\x14'+pa)
+                write_kv(PSBT_OUT_REDEEM_SCRIPT, b'\x00\x14'+pa)
 
             # optional? insert (partial) signatures that we already have
             if sigs and sigs[pk_pos]:
@@ -295,22 +297,32 @@ def build_psbt(tx: Transaction, wallet: Wallet):
                 # it is a change output (based on our standard derivation path)
                 pubkeys = [bfh(i) for i in wallet.get_public_keys(o.address)]
 
-                # always need a redeem script for multisig
+                # Add redeem/witness script?
                 if type(wallet) is Multisig_Wallet:
-                    scr = multisig_script([bh2u(i) for i in sorted(pubkeys)], wallet.m)
-                    write_kv(PSBT_OUT_REDEEM_SCRIPT, bfh(scr))
+                    # always need a redeem script for multisig cases
+                    scr = bfh(multisig_script([bh2u(i) for i in sorted(pubkeys)], wallet.m))
+
+                    if output_info.script_type == 'p2wsh-p2sh':
+                        write_kv(PSBT_OUT_WITNESS_SCRIPT, scr)
+                        write_kv(PSBT_OUT_REDEEM_SCRIPT, b'\x00\x20' + sha256(scr))
+                    elif output_info.script_type == 'p2wsh':
+                        write_kv(PSBT_OUT_WITNESS_SCRIPT, scr)
+                    elif output_info.script_type == 'p2sh':
+                        write_kv(PSBT_OUT_REDEEM_SCRIPT, scr)
+                    else:
+                        raise ValueError(output_info.script_type)
 
-                # document change output's bip32 derivation(s)
+                elif output_info.script_type == 'p2wpkh-p2sh':
+                    # need a redeem script when P2SH is used to wrap p2wpkh
+                    assert len(pubkeys) == 1
+                    pa = hash_160(pubkeys[0])
+                    write_kv(PSBT_OUT_REDEEM_SCRIPT, b'\x00\x14' + pa)
+
+                # Document change output's bip32 derivation(s)
                 for pubkey in pubkeys:
                     sk = subkeys[pubkey]
                     write_kv(PSBT_OUT_BIP32_DERIVATION, sk, pubkey)
 
-                    if output_info.script_type == 'p2wpkh-p2sh':
-                        assert len(pa) == 20
-                        assert len(pubkeys) == 1
-                        pa = hash_160(pubkey)
-                        write_kv(PSBT_OUT_REDEEM_SCRIPT, b'\x00\x14' + pa)
-
         out_fd.write(b'\x00')
 
     # capture for later use