#!/usr/bin/env python3
# Copyright (c) 2016-2024 The Hush developers
# Copyright (c) 2014 The Bitcoin Core developers
# Distributed under the GPLv3 software license, see the accompanying
# file COPYING or https://www.gnu.org/licenses/gpl-3.0.en.html

from test_framework.test_framework import BitcoinTestFramework
from test_framework.authproxy import JSONRPCException
from test_framework.util import initialize_chain_clean, start_node

from binascii import a2b_hex, b2a_hex
from hashlib import sha256
from struct import pack


def check_array_result(object_array, to_match, expected):
    """
    Pass in array of JSON objects, a dictionary with key/value pairs
    to match against, and another dictionary with expected key/value
    pairs.
    """
    num_matched = 0
    for item in object_array:
        all_match = True
        for key,value in to_match.items():
            if item[key] != value:
                all_match = False
        if not all_match:
            continue
        for key,value in expected.items():
            if item[key] != value:
                raise AssertionError("%s : expected %s=%s"%(str(item), str(key), str(value)))
            num_matched = num_matched+1
    if num_matched == 0:
        raise AssertionError("No objects matched %s"%(str(to_match)))

def b2x(b):
    return b2a_hex(b).decode('ascii')

# NOTE: This does not work for signed numbers (set the high bit) or zero (use b'\0')
def encodeUNum(n):
    s = bytearray(b'\1')
    while n > 127:
        s[0] += 1
        s.append(n % 256)
        n //= 256
    s.append(n)
    return bytes(s)

def varlenEncode(n):
    if n < 0xfd:
        return pack('<B', n)
    if n <= 0xffff:
        return b'\xfd' + pack('<H', n)
    if n <= 0xffffffff:
        return b'\xfe' + pack('<L', n)
    return b'\xff' + pack('<Q', n)

def dblsha(b):
    return sha256(sha256(b).digest()).digest()

def genmrklroot(leaflist):
    cur = leaflist
    while len(cur) > 1:
        n = []
        if len(cur) & 1:
            cur.append(cur[-1])
        for i in range(0, len(cur), 2):
            n.append(dblsha(cur[i] + cur[i+1]))
        cur = n
    return cur[0]

# ---------------------------------------------------------------------------
# Sapling v4 transaction layout.
#
# This test was written against the pre-Overwinter serialization, where a tx
# began with a 4-byte nVersion immediately followed by the vin count, so the
# first input's prevout hash lived at offset 4+1. DragonX transactions are
# Sapling v4: 4-byte header (nVersion | fOverwintered) + 4-byte nVersionGroupId
# + vin count, so the prevout hash starts 4 bytes further in. Poking the old
# offset corrupts nVersionGroupId and every proposal below just comes back
# "Block decode failed" instead of exercising any consensus rule.
CB_PREVOUT_OFF = 4+4+1
# Likewise the tx no longer ends at nLockTime: nExpiryHeight (4), valueBalance
# (8) and the empty vShieldedSpend/vShieldedOutput/vJoinSplit counts (1 each)
# trail it, so nLockTime is the 4 bytes at [-19:-15].
TX_TAIL_AFTER_LOCKTIME = 4+8+1+1+1

def tx_seq_off(tx):
    """Offset of the first input's nSequence in a Sapling v4 tx."""
    scriptlen_off = CB_PREVOUT_OFF + 32 + 4     # after prevout hash + prevout.n
    return scriptlen_off + 1 + tx[scriptlen_off]

def tx_vout0_value_off(tx):
    """Offset of the first output's 8-byte value in a Sapling v4 tx."""
    return tx_seq_off(tx) + 4 + 1               # after nSequence + vout count

def template_to_bytes(tmpl, txlist):
    blkver = pack('<L', tmpl['version'])
    mrklroot = genmrklroot(list(dblsha(a) for a in txlist))
    # hashFinalSaplingRoot. The all-zeroes placeholder this test used predates
    # Sapling; a header carrying the wrong root is rejected with
    # 'bad-sapling-root-in-block', which would sink even Test 11 (valid block).
    reserved = a2b_hex(tmpl['finalsaplingroothash'])[::-1]
    timestamp = pack('<L', tmpl['curtime'])
    nonce = b'\0'*32
    soln = b'\0'
    # bytearray, not bytes: Test 9 mutates one byte of the result in place.
    blk = bytearray(blkver + a2b_hex(tmpl['previousblockhash'])[::-1] + mrklroot + reserved + timestamp + a2b_hex(tmpl['bits'])[::-1] + nonce + soln)
    blk += varlenEncode(len(txlist))
    for tx in txlist:
        blk += tx
    return blk

def template_to_hex(tmpl, txlist):
    return b2x(template_to_bytes(tmpl, txlist))

def assert_template(node, tmpl, txlist, expect):
    rsp = node.getblocktemplate({'data':template_to_hex(tmpl, txlist),'mode':'proposal'})
    if rsp != expect:
        raise AssertionError('unexpected: %s' % (rsp,))

class GetBlockTemplateProposalTest(BitcoinTestFramework):
    '''
    Test block proposals with getblocktemplate.
    '''

    def setup_chain(self):
        print("Initializing test directory "+self.options.tmpdir)
        initialize_chain_clean(self.options.tmpdir, 1)

    def setup_network(self, split=False):
        # -daaforkheight=0: ContextualCheckBlockHeader only enforces nBits for a
        # smart chain above daaForkHeight, which defaults to
        # ASSETCHAINS_RANDOMX_VALIDATION+62000 (millions of blocks) so that a
        # fresh sync accepts DragonX's historical bad-nBits window. Without this
        # flag Test 8 (bad bits) is unreachable at regtest heights and the
        # daemon happily accepts a block with arbitrary nBits.
        self.nodes = [ start_node(0, self.options.tmpdir, ['-daaforkheight=0']) ]
        self.is_network_split = False

    def run_test(self):
        node = self.nodes[0]
        node.generate(1) # Mine a block to leave initial block download
        tmpl = node.getblocktemplate()
        if 'coinbasetxn' not in tmpl:
            rawcoinbase = encodeUNum(tmpl['height'])
            rawcoinbase += b'\x01-'
            hexcoinbase = b2x(rawcoinbase)
            hexoutval = b2x(pack('<Q', tmpl['coinbasevalue']))
            tmpl['coinbasetxn'] = {'data': '01000000' + '01' + '0000000000000000000000000000000000000000000000000000000000000000ffffffff' + ('%02x' % (len(rawcoinbase),)) + hexcoinbase + 'fffffffe' + '01' + hexoutval + '00' + '00000000'}
        txlist = list(bytearray(a2b_hex(a['data'])) for a in (tmpl['coinbasetxn'],) + tuple(tmpl['transactions']))

        # Test 0: Capability advertised
        assert('proposal' in tmpl['capabilities'])

        # NOTE: This test currently FAILS (regtest mode doesn't enforce block height in coinbase)
        ## Test 1: Bad height in coinbase
        #txlist[0][4+1+36+1+1] += 1
        #assert_template(node, tmpl, txlist, 'FIXME')
        #txlist[0][4+1+36+1+1] -= 1

        # Test 2: Bad input hash for gen tx
        txlist[0][CB_PREVOUT_OFF] += 1
        assert_template(node, tmpl, txlist, 'bad-cb-missing')
        txlist[0][CB_PREVOUT_OFF] -= 1

        # Test 3: Truncated final tx
        lastbyte = txlist[-1].pop()
        try:
            assert_template(node, tmpl, txlist, 'n/a')
        except JSONRPCException:
            pass  # Expected
        txlist[-1].append(lastbyte)

        # Test 4: Add an invalid tx to the end (duplicate of gen tx)
        txlist.append(txlist[0])
        assert_template(node, tmpl, txlist, 'bad-txns-duplicate')
        txlist.pop()

        # Test 5: Add an invalid tx to the end (non-duplicate)
        txlist.append(bytearray(txlist[0]))
        txlist[-1][CB_PREVOUT_OFF] = 0xff
        # DragonX is a fully private chain (ASSETCHAINS_PRIVATE): a non-coinbase
        # tx paying a positive amount to a t-addr is killed in CheckTransaction
        # with 'bad-txns-acprivacy-chain' before ConnectBlock ever looks up its
        # inputs. Zero the output value -- CheckTransaction exempts zero-value
        # vouts -- so the tx survives to the missing-input check this case is
        # actually about.
        _val_off = tx_vout0_value_off(txlist[-1])
        txlist[-1][_val_off:_val_off+8] = b'\0'*8
        assert_template(node, tmpl, txlist, 'bad-txns-inputs-missingorspent')
        txlist.pop()

        # Test 6: Future tx lock time
        # The server-supplied coinbase carries nSequence=0xffffffff, and
        # IsFinalTx() short-circuits on all-final inputs, so a future nLockTime
        # alone leaves the tx final and the block valid. (DragonX's IsFinalTx
        # also whitelists 0xfffffffe below the Hush hardfork height, which
        # regtest is.) Make the input genuinely non-final so the future
        # nLockTime is the thing under test.
        seq_off = tx_seq_off(txlist[0])
        realseq = txlist[0][seq_off:seq_off+4]
        txlist[0][seq_off:seq_off+4] = b'\0\0\0\0'
        reallocktime = txlist[0][-TX_TAIL_AFTER_LOCKTIME-4:-TX_TAIL_AFTER_LOCKTIME]
        txlist[0][-TX_TAIL_AFTER_LOCKTIME-4:-TX_TAIL_AFTER_LOCKTIME] = b'\xff\xff\xff\xff'
        assert_template(node, tmpl, txlist, 'bad-txns-nonfinal')
        txlist[0][-TX_TAIL_AFTER_LOCKTIME-4:-TX_TAIL_AFTER_LOCKTIME] = reallocktime
        txlist[0][seq_off:seq_off+4] = realseq

        # Test 7: Bad tx count
        txlist.append(b'')
        try:
            assert_template(node, tmpl, txlist, 'n/a')
        except JSONRPCException:
            pass  # Expected
        txlist.pop()

        # Test 8: Bad bits
        realbits = tmpl['bits']
        tmpl['bits'] = '1c0000ff'  # impossible in the real world
        assert_template(node, tmpl, txlist, 'bad-diffbits')
        tmpl['bits'] = realbits

        # Test 9: Bad merkle root
        rawtmpl = template_to_bytes(tmpl, txlist)
        rawtmpl[4+32] = (rawtmpl[4+32] + 1) % 0x100
        rsp = node.getblocktemplate({'data':b2x(rawtmpl),'mode':'proposal'})
        if rsp != 'bad-txnmrklroot':
            raise AssertionError('unexpected: %s' % (rsp,))

        # Test 10: Bad timestamps
        realtime = tmpl['curtime']
        tmpl['curtime'] = 0x7fffffff
        assert_template(node, tmpl, txlist, 'time-too-new')
        tmpl['curtime'] = 0
        # DragonX split Bitcoin's 'time-too-old' into 'time-too-old-median'
        # (block time <= prev MedianTimePast) and 'time-too-old-prevblock'.
        assert_template(node, tmpl, txlist, 'time-too-old-median')
        tmpl['curtime'] = realtime

        # Test 11: Valid block
        assert_template(node, tmpl, txlist, None)

        # Test 12: Orphan block
        tmpl['previousblockhash'] = 'ff00' * 16
        assert_template(node, tmpl, txlist, 'inconclusive-not-best-prevblk')

if __name__ == '__main__':
    GetBlockTemplateProposalTest().main()
