1import io
2
3from hashlib import sha256
4
5from ...encoding.hash import double_sha256
6from ...encoding.bytes32 import from_bytes_32
7from ...intbytes import byte2int, indexbytes
8
9from ..SolutionChecker import SolutionChecker, ScriptError
10from pycoin.satoshi import errno
11
12from pycoin.satoshi.satoshi_struct import stream_struct
13from pycoin.satoshi.satoshi_string import stream_satoshi_string
14
15from pycoin.satoshi.flags import (
16    SIGHASH_NONE, SIGHASH_SINGLE, SIGHASH_ANYONECANPAY,
17    VERIFY_DISCOURAGE_UPGRADABLE_WITNESS_PROGRAM, VERIFY_CLEANSTACK, VERIFY_WITNESS
18)
19
20from .ScriptTools import BitcoinScriptTools
21
22
23ZERO32 = b'\0' * 32
24
25
26class SegwitChecker(SolutionChecker):
27    # you must set VM
28    # you must set ScriptTools
29
30    V0_len20_prefix = BitcoinScriptTools.compile("OP_DUP OP_HASH160")
31    V0_len20_postfix = BitcoinScriptTools.compile("OP_EQUALVERIFY OP_CHECKSIG")
32
33    OP_0 = BitcoinScriptTools.int_for_opcode("OP_0")
34    OP_1 = BitcoinScriptTools.int_for_opcode("OP_1")
35    OP_16 = BitcoinScriptTools.int_for_opcode("OP_16")
36
37    def _make_witness_sighash_f(self, tx_in_idx):
38
39        def witness_signature_for_hash_type(hash_type, sig_blobs, vm):
40            return self._signature_for_hash_type_segwit(
41                vm.script[vm.begin_code_hash:], tx_in_idx, hash_type)
42
43        return witness_signature_for_hash_type
44
45    def _puzzle_script_for_len20_segwit(self, witness_program):
46        return self.V0_len20_prefix + self.ScriptTools.compile_push_data_list(
47            [witness_program]) + self.V0_len20_postfix
48
49    def _check_witness_program_v0(self, witness_solution_stack, witness_program):
50        size = len(witness_program)
51        if size == 32:
52            if len(witness_solution_stack) == 0:
53                raise ScriptError("witness program witness empty", errno.WITNESS_PROGRAM_WITNESS_EMPTY)
54            puzzle_script = witness_solution_stack[-1]
55            if sha256(puzzle_script).digest() != witness_program:
56                raise ScriptError("witness program mismatch", errno.WITNESS_PROGRAM_MISMATCH)
57            stack = list(witness_solution_stack[:-1])
58        elif size == 20:
59            # special case for pay-to-pubkeyhash; signature + pubkey in witness
60            if len(witness_solution_stack) != 2:
61                raise ScriptError("witness program mismatch", errno.WITNESS_PROGRAM_MISMATCH)
62            puzzle_script = self._puzzle_script_for_len20_segwit(witness_program)
63            stack = list(witness_solution_stack)
64        else:
65            raise ScriptError("witness program wrong length", errno.WITNESS_PROGRAM_WRONG_LENGTH)
66        return stack, puzzle_script
67
68    def _witness_program_version(self, script):
69        size = len(script)
70        if size < 4 or size > 42:
71            return None
72        first_opcode = byte2int(script)
73        if indexbytes(script, 1) + 2 != size:
74            return None
75        if first_opcode == self.OP_0:
76            return 0
77        if self.OP_1 <= first_opcode <= self.OP_16:
78            return first_opcode - self.OP_1 + 1
79        return None
80
81    def _hash_prevouts(self, hash_type):
82        if hash_type & SIGHASH_ANYONECANPAY:
83            return ZERO32
84        f = io.BytesIO()
85        for tx_in in self.tx.txs_in:
86            f.write(tx_in.previous_hash)
87            stream_struct("L", f, tx_in.previous_index)
88        return double_sha256(f.getvalue())
89
90    def _hash_sequence(self, hash_type):
91        if (
92                (hash_type & SIGHASH_ANYONECANPAY) or
93                ((hash_type & 0x1f) == SIGHASH_SINGLE) or
94                ((hash_type & 0x1f) == SIGHASH_NONE)
95        ):
96            return ZERO32
97
98        f = io.BytesIO()
99        for tx_in in self.tx.txs_in:
100            stream_struct("L", f, tx_in.sequence)
101        return double_sha256(f.getvalue())
102
103    def _hash_outputs(self, hash_type, tx_in_idx):
104        txs_out = self.tx.txs_out
105        if hash_type & 0x1f == SIGHASH_SINGLE:
106            if tx_in_idx >= len(txs_out):
107                return ZERO32
108            txs_out = txs_out[tx_in_idx:tx_in_idx+1]
109        elif hash_type & 0x1f == SIGHASH_NONE:
110            return ZERO32
111        f = io.BytesIO()
112        for tx_out in txs_out:
113            stream_struct("QS", f, tx_out.coin_value, tx_out.script)
114        return double_sha256(f.getvalue())
115
116    def _segwit_signature_preimage(self, script, tx_in_idx, hash_type):
117        f = io.BytesIO()
118        stream_struct("L", f, self.tx.version)
119        # calculate hash prevouts
120        f.write(self._hash_prevouts(hash_type))
121        f.write(self._hash_sequence(hash_type))
122        tx_in = self.tx.txs_in[tx_in_idx]
123        f.write(tx_in.previous_hash)
124        stream_struct("L", f, tx_in.previous_index)
125        tx_out = self.tx.unspents[tx_in_idx]
126        stream_satoshi_string(f, script)
127        stream_struct("Q", f, tx_out.coin_value)
128        stream_struct("L", f, tx_in.sequence)
129        f.write(self._hash_outputs(hash_type, tx_in_idx))
130        stream_struct("L", f, self.tx.lock_time)
131        stream_struct("L", f, hash_type)
132        return f.getvalue()
133
134    def _signature_for_hash_type_segwit(self, script, tx_in_idx, hash_type):
135        return from_bytes_32(double_sha256(self._segwit_signature_preimage(script, tx_in_idx, hash_type)))
136
137    def witness_program_tuple(self, tx_context, puzzle_script, solution_stack, flags, is_p2sh):
138        if not flags & VERIFY_WITNESS:
139            return
140
141        witness_version = self._witness_program_version(puzzle_script)
142        if witness_version is None:
143            if len(tx_context.witness_solution_stack) > 0:
144                raise ScriptError("witness unexpected", errno.WITNESS_UNEXPECTED)
145        else:
146            witness_program = puzzle_script[2:]
147            if len(solution_stack) > 0:
148                err = errno.WITNESS_MALLEATED_P2SH if is_p2sh else errno.WITNESS_MALLEATED
149                raise ScriptError("script sig is not blank on segwit input", err)
150
151            for s in tx_context.witness_solution_stack:
152                if len(s) > self.VM.MAX_BLOB_LENGTH:
153                    raise ScriptError("pushing too much data onto stack", errno.PUSH_SIZE)
154
155            if witness_version == 0:
156                stack, puzzle_script = self._check_witness_program_v0(
157                    tx_context.witness_solution_stack, witness_program)
158                sighash_f = self._make_witness_sighash_f(tx_context.tx_in_idx)
159                return puzzle_script, stack, flags | VERIFY_CLEANSTACK, sighash_f
160            elif flags & VERIFY_DISCOURAGE_UPGRADABLE_WITNESS_PROGRAM:
161                raise ScriptError(
162                    "this version witness program not yet supported", errno.DISCOURAGE_UPGRADABLE_WITNESS_PROGRAM)
163