trezor: Determine segwit based on scripts instead of provided info

achow101 · achow101 · commit a47f7e583e61 · 2020-06-08T19:14:52.000-04:00
diff --git a/hwilib/devices/trezor.py b/hwilib/devices/trezor.py
@@ -45,6 +45,10 @@
 from ..serializations import (
     CTxOut,
     ExtendedKey,
+    is_p2pkh,
+    is_p2sh,
+    is_p2wsh,
+    is_witness,
     ser_uint256,
 )
 from .. import bech32
@@ -200,26 +204,46 @@ def sign_tx(self, tx):
 
                 # Detrermine spend type
                 scriptcode = b''
+                utxo = None
                 if psbt_in.witness_utxo:
                     utxo = psbt_in.witness_utxo
-                    # Check if the output is p2sh
-                    if psbt_in.witness_utxo.is_p2sh():
+                if psbt_in.non_witness_utxo:
+                    if txin.prevout.hash != psbt_in.non_witness_utxo.sha256:
+                        raise BadArgumentError('Input {} has a non_witness_utxo with the wrong hash'.format(input_num))
+                    utxo = psbt_in.non_witness_utxo.vout[txin.prevout.n]
+                if utxo is None:
+                    continue
+                scriptcode = utxo.scriptPubKey
+
+                # Check if P2SH
+                p2sh = False
+                if is_p2sh(scriptcode):
+                    # Look up redeemscript
+                    if len(psbt_in.redeem_script) == 0:
+                        continue
+                    scriptcode = psbt_in.redeem_script
+                    p2sh = True
+
+                # Check segwit
+                is_wit, _, _ = is_witness(scriptcode)
+
+                if is_wit:
+                    if p2sh:
                         txinputtype.script_type = proto.InputScriptType.SPENDP2SHWITNESS
                     else:
                         txinputtype.script_type = proto.InputScriptType.SPENDWITNESS
-                    scriptcode = psbt_in.witness_utxo.scriptPubKey
-                    txinputtype.amount = psbt_in.witness_utxo.nValue
-                elif psbt_in.non_witness_utxo:
-                    utxo = psbt_in.non_witness_utxo.vout[txin.prevout.n]
+                else:
                     txinputtype.script_type = proto.InputScriptType.SPENDADDRESS
-                    scriptcode = utxo.scriptPubKey
-                    txinputtype.amount = psbt_in.non_witness_utxo.vout[txin.prevout.n].nValue
-
-                # Set the script
-                if psbt_in.witness_script:
+                txinputtype.amount = utxo.nValue
+
+                # Check if P2WSH
+                p2wsh = False
+                if is_p2wsh(scriptcode):
+                    # Look up witnessscript
+                    if len(psbt_in.witness_script) == 0:
+                        continue
                     scriptcode = psbt_in.witness_script
-                elif psbt_in.redeem_script:
-                    scriptcode = psbt_in.redeem_script
+                    p2wsh = True
 
                 def ignore_input():
                     txinputtype.address_n = [0x80000000 | 84, 0x80000000 | (1 if self.is_testnet else 0)]
@@ -240,11 +264,11 @@ def ignore_input():
                             # Cannot sign bare multisig, ignore it
                             ignore_input()
                             continue
-                elif not is_ms and psbt_in.non_witness_utxo and not utxo.is_p2pkh:
+                elif not is_ms and not is_wit and not is_p2pkh(scriptcode):
                     # Cannot sign unknown spk, ignore it
                     ignore_input()
                     continue
-                elif not is_ms and psbt_in.witness_utxo and psbt_in.witness_script:
+                elif not is_ms and is_wit and p2wsh:
                     # Cannot sign unknown witness script, ignore it
                     ignore_input()
                     continue
@@ -273,7 +297,8 @@ def ignore_input():
                     ignore_input()
                     continue
                 elif not found and found_in_sigs: # All of our keys are in partial_sigs, ignore whatever signature is produced for this input
-                    to_ignore.append(input_num)
+                    ignore_input()
+                    continue
 
                 # append to inputs
                 inputs.append(txinputtype)
