ledger: Determine segwit based on scripts and allow both UTXOs

Determine whether the input is segwit based on the script. Also handle
when both non_witness_utxo and witness_utxo are provided.

hwilib/devices/ledger.py

@@ -26,6 +26,10 @@
     ExtendedKey,
     hash256,
     hash160,
+    is_p2sh,
+    is_p2wpkh,
+    is_p2wsh,
+    is_witness,
     CTransaction,
 )
 import logging
@@ -207,59 +211,54 @@ def sign_tx(self, tx):
             seq.reverse()
             seq_hex = ''.join('{:02x}'.format(x) for x in seq)
 
+            scriptcode = b""
+            utxo = None
+            if psbt_in.witness_utxo:
+                utxo = psbt_in.witness_utxo
             if psbt_in.non_witness_utxo:
-                segwit_inputs.append({"value": txin.prevout.serialize() + struct.pack("<Q", psbt_in.non_witness_utxo.vout[txin.prevout.n].nValue), "witness": True, "sequence": seq_hex})
+                if txin.prevout.hash != psbt_in.non_witness_utxo.sha256:
+                    raise BadArgumentError('Input {} has a non_witness_utxo with the wrong hash'.format(i_num))
+                utxo = psbt_in.non_witness_utxo.vout[txin.prevout.n]
+            if utxo is None:
+                raise Exception("PSBT is missing input utxo information, cannot sign")
+            scriptcode = utxo.scriptPubKey
+
+            if is_p2sh(scriptcode):
+                if len(psbt_in.redeem_script) == 0:
+                    continue
+                scriptcode = psbt_in.redeem_script
+
+            is_wit, _, _ = is_witness(scriptcode)
+
+            segwit_inputs.append({"value": txin.prevout.serialize() + struct.pack("<Q", utxo.nValue), "witness": True, "sequence": seq_hex})
+            if is_wit:
+                if is_p2wsh(scriptcode):
+                    if len(psbt_in.witness_script) == 0:
+                        continue
+                    scriptcode = psbt_in.witness_script
+                elif is_p2wpkh(scriptcode):
+                    _, _, wit_prog = is_witness(scriptcode)
+                    scriptcode = b"\x76\xa9\x14" + wit_prog + b"\x88\xac"
+                else:
+                    continue
+                has_segwit = True
+            else:
                 # We only need legacy inputs in the case where all inputs are legacy, we check
                 # later
                 ledger_prevtx = bitcoinTransaction(psbt_in.non_witness_utxo.serialize())
                 legacy_inputs.append(self.app.getTrustedInput(ledger_prevtx, txin.prevout.n))
                 legacy_inputs[-1]["sequence"] = seq_hex
                 has_legacy = True
-            else:
-                segwit_inputs.append({"value": txin.prevout.serialize() + struct.pack("<Q", psbt_in.witness_utxo.nValue), "witness": True, "sequence": seq_hex})
-                has_segwit = True
 
             pubkeys = []
             signature_attempts = []
 
-            scriptCode = b""
-            witness_program = b""
-            if psbt_in.witness_utxo is not None and psbt_in.witness_utxo.is_p2sh():
-                redeemscript = psbt_in.redeem_script
-                witness_program += redeemscript
-            elif psbt_in.non_witness_utxo is not None and psbt_in.non_witness_utxo.vout[txin.prevout.n].is_p2sh():
-                redeemscript = psbt_in.redeem_script
-            elif psbt_in.witness_utxo is not None:
-                witness_program += psbt_in.witness_utxo.scriptPubKey
-            elif psbt_in.non_witness_utxo is not None:
-                # No-op
-                redeemscript = b""
-                witness_program = b""
-            else:
-                raise Exception("PSBT is missing input utxo information, cannot sign")
-
-            # Check if witness_program is script hash
-            if len(witness_program) == 34 and witness_program[0] == 0x00 and witness_program[1] == 0x20:
-                # look up witnessscript and set as scriptCode
-                witnessscript = psbt_in.witness_script
-                scriptCode += witnessscript
-            elif len(witness_program) > 0:
-                # p2wpkh
-                scriptCode += b"\x76\xa9\x14"
-                scriptCode += witness_program[2:]
-                scriptCode += b"\x88\xac"
-            elif len(witness_program) == 0:
-                if len(redeemscript) > 0:
-                    scriptCode = redeemscript
-                else:
-                    scriptCode = psbt_in.non_witness_utxo.vout[txin.prevout.n].scriptPubKey
-
             # Save scriptcode for later signing
-            script_codes[i_num] = scriptCode
+            script_codes[i_num] = scriptcode
 
             # Find which pubkeys could sign this input (should be all?)
             for pubkey in psbt_in.hd_keypaths.keys():
-                if hash160(pubkey) in scriptCode or pubkey in scriptCode:
+                if hash160(pubkey) in scriptcode or pubkey in scriptcode:
                     pubkeys.append(pubkey)
 
             # Figure out which keys in inputs are from our wallet
@@ -287,9 +286,6 @@ def sign_tx(self, tx):
 
             # For each input we control do segwit signature
             for i in range(len(segwit_inputs)):
-                # Don't try to sign legacy inputs
-                if tx.inputs[i].non_witness_utxo is not None:
-                    continue
                 for signature_attempt in all_signature_attempts[i]:
                     self.app.startUntrustedTransaction(False, 0, [segwit_inputs[i]], script_codes[i], c_tx.nVersion)
                     tx.inputs[i].partial_sigs[signature_attempt[1]] = self.app.untrustedHashSign(signature_attempt[0], "", c_tx.nLockTime, 0x01)