From 704a2ac14d987f10e5e79c03903bbb472c35aa3d Mon Sep 17 00:00:00 2001
From: Ariel Gabizon <ariel.gabizon@gmail.com>
Date: Wed, 13 Jun 2018 15:45:05 +0200
Subject: [PATCH] Implement ZIP 143 test vector generation

No support for JoinSplits yet.

Co-authored-by: Jack Grigg <jack@z.cash>
---
 transaction.py | 117 ++++++++++++++++++++++++++++++++++++
 zip_0143.py    | 157 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 274 insertions(+)
 create mode 100644 transaction.py
 create mode 100644 zip_0143.py

diff --git a/transaction.py b/transaction.py
new file mode 100644
index 0000000..8cadb4a
--- /dev/null
+++ b/transaction.py
@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+import struct
+
+MAX_MONEY = 21000000 * 100000000
+TX_EXPIRY_HEIGHT_THRESHOLD = 500000000
+
+OVERWINTER_VERSION_GROUP_ID = 0x03C48270
+OVERWINTER_TX_VERSION = 3
+
+
+RAND_OPCODES = [
+    0x00, # OP_FALSE,
+    0x51, # OP_1,
+    0x52, # OP_2,
+    0x53, # OP_3,
+    0xac, # OP_CHECKSIG,
+    0x63, # OP_IF,
+    0x65, # OP_VERIF,
+    0x6a, # OP_RETURN,
+]
+
+class Script(object):
+    def __init__(self, rand):
+        self._script = bytes([
+            rand.a(RAND_OPCODES) for i in range(rand.u8() % 10)
+        ])
+
+    def raw(self):
+        return self._script
+
+    def __bytes__(self):
+        return struct.pack('b', len(self._script)) + self._script
+
+
+class OutPoint(object):
+    def __init__(self, rand):
+        self.txid = rand.b(32)
+        self.n = rand.u32()
+
+    def __bytes__(self):
+        return self.txid + struct.pack('<I', self.n)
+
+
+class TxIn(object):
+    def __init__(self, rand):
+        self.prevout = OutPoint(rand)
+        self.scriptSig = Script(rand)
+        self.nSequence = rand.u32()
+
+    def __bytes__(self):
+        return (
+            bytes(self.prevout) +
+            bytes(self.scriptSig) +
+            struct.pack('<I', self.nSequence)
+        )
+
+
+class TxOut(object):
+    def __init__(self, rand):
+        self.nValue = rand.u64() % (MAX_MONEY + 1)
+        self.scriptPubKey = Script(rand)
+
+    def __bytes__(self):
+        return struct.pack('<Q', self.nValue) + bytes(self.scriptPubKey)
+
+
+class Transaction(object):
+    def __init__(self, rand, version):
+        if version == OVERWINTER_TX_VERSION:
+            self.fOverwintered = True
+            self.nVersionGroupId = OVERWINTER_VERSION_GROUP_ID
+            self.nVersion = OVERWINTER_TX_VERSION
+        else:
+            self.fOverwintered = False
+            self.nVersion = rand.u32() & ((1 << 31) - 1)
+
+        self.vin = []
+        for i in range(rand.u8() % 3):
+            self.vin.append(TxIn(rand))
+
+        self.vout = []
+        for i in range(rand.u8() % 3):
+            self.vout.append(TxOut(rand))
+
+        self.nLockTime = rand.u32()
+        self.nExpiryHeight = rand.u32() % TX_EXPIRY_HEIGHT_THRESHOLD
+
+    def header(self):
+        return self.nVersion | (1 << 31 if self.fOverwintered else 0)
+
+    def __bytes__(self):
+        ret = b''
+        ret += struct.pack('<I', self.header())
+        if self.fOverwintered:
+            ret += struct.pack('<I', self.nVersionGroupId)
+
+        isOverwinterV3 = \
+            self.fOverwintered and \
+            self.nVersionGroupId == OVERWINTER_VERSION_GROUP_ID and \
+            self.nVersion == OVERWINTER_TX_VERSION
+
+        ret += struct.pack('b', len(self.vin))
+        for x in self.vin:
+            ret += bytes(x)
+
+        ret += struct.pack('b', len(self.vout))
+        for x in self.vout:
+            ret += bytes(x)
+
+        ret += struct.pack('<I', self.nLockTime)
+        if isOverwinterV3:
+            ret += struct.pack('<I', self.nExpiryHeight)
+
+        if self.nVersion >= 2:
+            ret += struct.pack('b', 0)
+
+        return ret
diff --git a/zip_0143.py b/zip_0143.py
new file mode 100644
index 0000000..b2b279c
--- /dev/null
+++ b/zip_0143.py
@@ -0,0 +1,157 @@
+#!/usr/bin/env python3
+from pyblake2 import blake2b
+import struct
+
+from transaction import (
+    MAX_MONEY,
+    OVERWINTER_TX_VERSION,
+    Script,
+    Transaction,
+)
+from tv_output import render_args, render_tv
+from tv_rand import Rand
+
+
+SIGHASH_ALL = 1
+SIGHASH_NONE = 2
+SIGHASH_SINGLE = 3
+SIGHASH_ANYONECANPAY = 0x80
+
+NOT_AN_INPUT = -1 # For portability of the test vectors
+
+def getHashPrevouts(tx):
+    digest = blake2b(digest_size=32, person=b'ZcashPrevoutHash')
+    for x in tx.vin:
+        digest.update(bytes(x.prevout))
+    return digest.digest()
+
+def getHashSequence(tx):
+    digest = blake2b(digest_size=32, person=b'ZcashSequencHash')
+    for x in tx.vin:
+        digest.update(struct.pack('<I', x.nSequence))
+    return digest.digest()
+
+def getHashOutputs(tx):
+    digest = blake2b(digest_size=32, person=b'ZcashOutputsHash')
+    for x in tx.vout:
+        digest.update(bytes(x))
+    return digest.digest()
+
+
+# Currently assumes the nHashType is SIGHASHALL
+# and that there are no joinSplits
+def signature_hash(scriptCode, tx, nIn, nHashType, amount, consensusBranchId):
+    hashPrevouts = b'\x00'*32
+    hashSequence = b'\x00'*32
+    hashOutputs = b'\x00'*32
+    hashJoinSplits = b'\x00'*32
+
+    if not (nHashType & SIGHASH_ANYONECANPAY):
+        hashPrevouts = getHashPrevouts(tx)
+
+    if (not (nHashType & SIGHASH_ANYONECANPAY)) and \
+        (nHashType & 0x1f) != SIGHASH_SINGLE and \
+        (nHashType & 0x1f) != SIGHASH_NONE:
+        hashSequence = getHashSequence(tx)
+
+    if (nHashType & 0x1f) != SIGHASH_SINGLE and \
+        (nHashType & 0x1f) != SIGHASH_NONE:
+        hashOutputs = getHashOutputs(tx)
+    elif (nHashType & 0x1f) == SIGHASH_SINGLE and \
+        0 <= nIn and nIn < len(tx.vout):
+        digest = blake2b(digest_size=32, person=b'ZcashOutputsHash')
+        digest.update(bytes(tx.vout[nIn]))
+        hashOutputs = digest.digest()
+
+    digest = blake2b(
+        digest_size=32,
+        person=b'ZcashSigHash' + struct.pack('<I', consensusBranchId),
+    )
+
+    digest.update(struct.pack('<I', tx.header()))
+    digest.update(struct.pack('<I', tx.nVersionGroupId))
+    digest.update(hashPrevouts)
+    digest.update(hashSequence)
+    digest.update(hashOutputs)
+    digest.update(hashJoinSplits)
+    digest.update(struct.pack('<I', tx.nLockTime))
+    digest.update(struct.pack('<I', tx.nExpiryHeight))
+    digest.update(struct.pack('<I', nHashType))
+
+    if nIn != NOT_AN_INPUT:
+        digest.update(bytes(tx.vin[nIn].prevout))
+        digest.update(bytes(scriptCode))
+        digest.update(struct.pack('<Q', amount))
+        digest.update(struct.pack('<I', tx.vin[nIn].nSequence))
+
+    return digest.digest()
+
+
+def main():
+    args = render_args()
+
+    from random import Random
+    rng = Random(0xabad533d)
+    def randbytes(l):
+        ret = []
+        while len(ret) < l:
+            ret.append(rng.randrange(0, 256))
+        return bytes(ret)
+    rand = Rand(randbytes)
+
+    consensusBranchId = 0x5ba81b19 # Overwinter
+
+    test_vectors = []
+    for i in range(10):
+        tx = Transaction(rand, OVERWINTER_TX_VERSION)
+        scriptCode = Script(rand)
+        nIn = rand.u8() % (len(tx.vin) + 1)
+        if nIn == len(tx.vin):
+            nIn = NOT_AN_INPUT
+        nHashType = SIGHASH_ALL if nIn == NOT_AN_INPUT else rand.a([
+            SIGHASH_ALL,
+            SIGHASH_NONE,
+            SIGHASH_SINGLE,
+            SIGHASH_ALL | SIGHASH_ANYONECANPAY,
+            SIGHASH_NONE | SIGHASH_ANYONECANPAY,
+            SIGHASH_SINGLE | SIGHASH_ANYONECANPAY,
+        ])
+        amount = rand.u64() % (MAX_MONEY + 1)
+
+        sighash = signature_hash(
+            scriptCode,
+            tx,
+            nIn,
+            nHashType,
+            amount,
+            consensusBranchId,
+        )
+
+        test_vectors.append({
+            'tx': bytes(tx),
+            'scriptCode': scriptCode.raw(),
+            'nIn': nIn,
+            'nHashType': nHashType,
+            'amount': amount,
+            'consensusBranchId': consensusBranchId,
+            'sighash': sighash,
+        })
+
+    render_tv(
+        args,
+        'zip_0143',
+        (
+            ('tx', 'Vec<u8>'),
+            ('scriptCode', 'Vec<u8>'),
+            ('nIn', 'u32'),
+            ('nHashType', 'u32'),
+            ('amount', 'u64'),
+            ('consensusBranchId', 'u32'),
+            ('sighash', '[u8; 32]'),
+        ),
+        test_vectors,
+    )
+
+
+if __name__ == '__main__':
+    main()