lbry-sdk/torba/basetransaction.py

import logging
import typing
from typing import List, Iterable, Optional
from binascii import hexlify

from twisted.internet import defer

from torba.basescript import BaseInputScript, BaseOutputScript
from torba.baseaccount import BaseAccount
from torba.constants import COIN, NULL_HASH32
from torba.bcd_data_stream import BCDataStream
from torba.hash import sha256, TXRef, TXRefImmutable
from torba.util import ReadOnlyList

if typing.TYPE_CHECKING:
    from torba import baseledger


log = logging.getLogger()


class TXRefMutable(TXRef):

    __slots__ = ('tx',)

    def __init__(self, tx: 'BaseTransaction') -> None:
        super().__init__()
        self.tx = tx

    @property
    def id(self):
        if self._id is None:
            self._id = hexlify(self.hash[::-1]).decode()
        return self._id

    @property
    def hash(self):
        if self._hash is None:
            self._hash = sha256(sha256(self.tx.raw))
        return self._hash

    def reset(self):
        self._id = None
        self._hash = None


class TXORef:

    __slots__ = 'tx_ref', 'position'

    def __init__(self, tx_ref: TXRef, position: int) -> None:
        self.tx_ref = tx_ref
        self.position = position

    @property
    def id(self):
        return '{}:{}'.format(self.tx_ref.id, self.position)

    @property
    def is_null(self):
        return self.tx_ref.is_null

    @property
    def txo(self) -> Optional['BaseOutput']:
        return None


class TXORefResolvable(TXORef):

    __slots__ = ('_txo',)

    def __init__(self, txo: 'BaseOutput') -> None:
        assert txo.tx_ref is not None
        assert txo.position is not None
        super().__init__(txo.tx_ref, txo.position)
        self._txo = txo

    @property
    def txo(self):
        return self._txo


class InputOutput:

    __slots__ = 'tx_ref', 'position'

    def __init__(self, tx_ref: TXRef = None, position: int = None) -> None:
        self.tx_ref = tx_ref
        self.position = position

    @property
    def size(self) -> int:
        """ Size of this input / output in bytes. """
        stream = BCDataStream()
        self.serialize_to(stream)
        return len(stream.get_bytes())

    def get_fee(self, ledger):
        return self.size * ledger.fee_per_byte

    def serialize_to(self, stream, alternate_script=None):
        raise NotImplementedError


class BaseInput(InputOutput):

    script_class = BaseInputScript

    NULL_SIGNATURE = b'\x00'*72
    NULL_PUBLIC_KEY = b'\x00'*33

    __slots__ = 'txo_ref', 'sequence', 'coinbase', 'script'

    def __init__(self, txo_ref: TXORef, script: BaseInputScript, sequence: int = 0xFFFFFFFF,
                 tx_ref: TXRef = None, position: int = None) -> None:
        super().__init__(tx_ref, position)
        self.txo_ref = txo_ref
        self.sequence = sequence
        self.coinbase = script if txo_ref.is_null else None
        self.script = script if not txo_ref.is_null else None

    @property
    def is_coinbase(self):
        return self.coinbase is not None

    @classmethod
    def spend(cls, txo: 'BaseOutput') -> 'BaseInput':
        """ Create an input to spend the output."""
        assert txo.script.is_pay_pubkey_hash, 'Attempting to spend unsupported output.'
        script = cls.script_class.redeem_pubkey_hash(cls.NULL_SIGNATURE, cls.NULL_PUBLIC_KEY)
        return cls(txo.ref, script)

    @property
    def amount(self) -> int:
        """ Amount this input adds to the transaction. """
        if self.txo_ref.txo is None:
            raise ValueError('Cannot resolve output to get amount.')
        return self.txo_ref.txo.amount

    @property
    def is_my_account(self) -> Optional[bool]:
        """ True if the output this input spends is yours. """
        if self.txo_ref.txo is None:
            return False
        return self.txo_ref.txo.is_my_account

    @classmethod
    def deserialize_from(cls, stream):
        tx_ref = TXRefImmutable.from_hash(stream.read(32))
        position = stream.read_uint32()
        script = stream.read_string()
        sequence = stream.read_uint32()
        return cls(
            TXORef(tx_ref, position),
            cls.script_class(script) if not tx_ref.is_null else script,
            sequence
        )

    def serialize_to(self, stream, alternate_script=None):
        stream.write(self.txo_ref.tx_ref.hash)
        stream.write_uint32(self.txo_ref.position)
        if alternate_script is not None:
            stream.write_string(alternate_script)
        else:
            if self.is_coinbase:
                stream.write_string(self.coinbase)
            else:
                stream.write_string(self.script.source)
        stream.write_uint32(self.sequence)


class BaseOutputEffectiveAmountEstimator:

    __slots__ = 'txo', 'txi', 'fee', 'effective_amount'

    def __init__(self, ledger: 'baseledger.BaseLedger', txo: 'BaseOutput') -> None:
        self.txo = txo
        self.txi = ledger.transaction_class.input_class.spend(txo)
        self.fee: int = self.txi.get_fee(ledger)
        self.effective_amount: int = txo.amount - self.fee

    def __lt__(self, other):
        return self.effective_amount < other.effective_amount


class BaseOutput(InputOutput):

    script_class = BaseOutputScript
    estimator_class = BaseOutputEffectiveAmountEstimator

    __slots__ = 'amount', 'script', 'is_change', 'is_my_account'

    def __init__(self, amount: int, script: BaseOutputScript,
                 tx_ref: TXRef = None, position: int = None,
                 is_change: Optional[bool] = None, is_my_account: Optional[bool] = None
                 ) -> None:
        super().__init__(tx_ref, position)
        self.amount = amount
        self.script = script
        self.is_change = is_change
        self.is_my_account = is_my_account

    def update_annotations(self, annotated):
        if annotated is None:
            self.is_change = False
            self.is_my_account = False
        else:
            self.is_change = annotated.is_change
            self.is_my_account = annotated.is_my_account

    @property
    def ref(self):
        return TXORefResolvable(self)

    @property
    def id(self):
        return self.ref.id

    def get_address(self, ledger):
        return ledger.hash160_to_address(
            self.script.values['pubkey_hash']
        )

    def get_estimator(self, ledger):
        return self.estimator_class(ledger, self)

    @classmethod
    def pay_pubkey_hash(cls, amount, pubkey_hash):
        return cls(amount, cls.script_class.pay_pubkey_hash(pubkey_hash))

    @classmethod
    def deserialize_from(cls, stream):
        return cls(
            amount=stream.read_uint64(),
            script=cls.script_class(stream.read_string())
        )

    def serialize_to(self, stream, alternate_script=None):
        stream.write_uint64(self.amount)
        stream.write_string(self.script.source)


class BaseTransaction:

    input_class = BaseInput
    output_class = BaseOutput

    def __init__(self, raw=None, version: int = 1, locktime: int = 0, is_verified: bool = False,
                 height: int = -1, position: int = -1) -> None:
        self._raw = raw
        self.ref = TXRefMutable(self)
        self.version = version
        self.locktime = locktime
        self._inputs: List[BaseInput] = []
        self._outputs: List[BaseOutput] = []
        self.is_verified = is_verified
        self.height = height
        self.position = position
        if raw is not None:
            self._deserialize()

    @property
    def id(self):
        return self.ref.id

    @property
    def hash(self):
        return self.ref.hash

    @property
    def raw(self):
        if self._raw is None:
            self._raw = self._serialize()
        return self._raw

    def _reset(self):
        self._raw = None
        self.ref.reset()

    @property
    def inputs(self) -> ReadOnlyList[BaseInput]:
        return ReadOnlyList(self._inputs)

    @property
    def outputs(self) -> ReadOnlyList[BaseOutput]:
        return ReadOnlyList(self._outputs)

    def _add(self, new_ios: Iterable[InputOutput], existing_ios: List) -> 'BaseTransaction':
        for txio in new_ios:
            txio.tx_ref = self.ref
            txio.position = len(existing_ios)
            existing_ios.append(txio)
        self._reset()
        return self

    def add_inputs(self, inputs: Iterable[BaseInput]) -> 'BaseTransaction':
        return self._add(inputs, self._inputs)

    def add_outputs(self, outputs: Iterable[BaseOutput]) -> 'BaseTransaction':
        return self._add(outputs, self._outputs)

    @property
    def size(self) -> int:
        """ Size in bytes of the entire transaction. """
        return len(self.raw)

    @property
    def base_size(self) -> int:
        """ Size of transaction without inputs or outputs in bytes. """
        return (
            self.size
            - sum(txi.size for txi in self._inputs)
            - sum(txo.size for txo in self._outputs)
        )

    @property
    def input_sum(self):
        return sum(i.amount for i in self.inputs if i.txo_ref.txo is not None)

    @property
    def output_sum(self):
        return sum(o.amount for o in self.outputs)

    @property
    def net_account_balance(self) -> int:
        balance = 0
        for txi in self.inputs:
            if txi.txo_ref.txo is None:
                continue
            if txi.is_my_account is None:
                raise ValueError(
                    "Cannot access net_account_balance if inputs/outputs do not "
                    "have is_my_account set properly."
                )
            elif txi.is_my_account:
                balance -= txi.amount
        for txo in self.outputs:
            if txo.is_my_account is None:
                raise ValueError(
                    "Cannot access net_account_balance if inputs/outputs do not "
                    "have is_my_account set properly."
                )
            elif txo.is_my_account:
                balance += txo.amount
        return balance

    @property
    def fee(self) -> int:
        return self.input_sum - self.output_sum

    def get_base_fee(self, ledger) -> int:
        """ Fee for base tx excluding inputs and outputs. """
        return self.base_size * ledger.fee_per_byte

    def get_effective_input_sum(self, ledger) -> int:
        """ Sum of input values *minus* the cost involved to spend them. """
        return sum(txi.amount - txi.get_fee(ledger) for txi in self._inputs)

    def get_total_output_sum(self, ledger) -> int:
        """ Sum of output values *plus* the cost involved to spend them. """
        return sum(txo.amount + txo.get_fee(ledger) for txo in self._outputs)

    def _serialize(self, with_inputs: bool = True) -> bytes:
        stream = BCDataStream()
        stream.write_uint32(self.version)
        if with_inputs:
            stream.write_compact_size(len(self._inputs))
            for txin in self._inputs:
                txin.serialize_to(stream)
        stream.write_compact_size(len(self._outputs))
        for txout in self._outputs:
            txout.serialize_to(stream)
        stream.write_uint32(self.locktime)
        return stream.get_bytes()

    def _serialize_for_signature(self, signing_input: int) -> bytes:
        stream = BCDataStream()
        stream.write_uint32(self.version)
        stream.write_compact_size(len(self._inputs))
        for i, txin in enumerate(self._inputs):
            if signing_input == i:
                assert txin.txo_ref.txo is not None
                txin.serialize_to(stream, txin.txo_ref.txo.script.source)
            else:
                txin.serialize_to(stream, b'')
        stream.write_compact_size(len(self._outputs))
        for txout in self._outputs:
            txout.serialize_to(stream)
        stream.write_uint32(self.locktime)
        stream.write_uint32(self.signature_hash_type(1))  # signature hash type: SIGHASH_ALL
        return stream.get_bytes()

    def _deserialize(self):
        if self._raw is not None:
            stream = BCDataStream(self._raw)
            self.version = stream.read_uint32()
            input_count = stream.read_compact_size()
            self.add_inputs([
                self.input_class.deserialize_from(stream) for _ in range(input_count)
            ])
            output_count = stream.read_compact_size()
            self.add_outputs([
                self.output_class.deserialize_from(stream) for _ in range(output_count)
            ])
            self.locktime = stream.read_uint32()

    @classmethod
    def ensure_all_have_same_ledger(cls, funding_accounts: Iterable[BaseAccount],
                                    change_account: BaseAccount = None) -> 'baseledger.BaseLedger':
        ledger = None
        for account in funding_accounts:
            if ledger is None:
                ledger = account.ledger
            if ledger != account.ledger:
                raise ValueError(
                    'All funding accounts used to create a transaction must be on the same ledger.'
                )
        if change_account is not None and change_account.ledger != ledger:
            raise ValueError('Change account must use same ledger as funding accounts.')
        if ledger is None:
            raise ValueError('No ledger found.')
        return ledger

    @classmethod
    @defer.inlineCallbacks
    def create(cls, inputs: Iterable[BaseInput], outputs: Iterable[BaseOutput],
               funding_accounts: Iterable[BaseAccount], change_account: BaseAccount):
        """ Find optimal set of inputs when only outputs are provided; add change
            outputs if only inputs are provided or if inputs are greater than outputs. """

        tx = cls() \
            .add_inputs(inputs) \
            .add_outputs(outputs)

        ledger = cls.ensure_all_have_same_ledger(funding_accounts, change_account)

        # value of the outputs plus associated fees
        cost = (
            tx.get_base_fee(ledger) +
            tx.get_total_output_sum(ledger)
        )
        # value of the inputs less the cost to spend those inputs
        payment = tx.get_effective_input_sum(ledger)

        try:

            if payment < cost:
                deficit = cost - payment
                spendables = yield ledger.get_spendable_utxos(deficit, funding_accounts)
                if not spendables:
                    raise ValueError('Not enough funds to cover this transaction.')
                payment += sum(s.effective_amount for s in spendables)
                tx.add_inputs(s.txi for s in spendables)

            if payment > cost:
                cost_of_change = (
                    tx.get_base_fee(ledger) +
                    cls.output_class.pay_pubkey_hash(COIN, NULL_HASH32).get_fee(ledger)
                )
                change = payment - cost
                if change > cost_of_change:
                    change_address = yield change_account.change.get_or_create_usable_address()
                    change_hash160 = change_account.ledger.address_to_hash160(change_address)
                    change_amount = change - cost_of_change
                    change_output = cls.output_class.pay_pubkey_hash(change_amount, change_hash160)
                    change_output.is_change = True
                    tx.add_outputs([cls.output_class.pay_pubkey_hash(change_amount, change_hash160)])

            yield tx.sign(funding_accounts)

        except Exception as e:
            log.exception('Failed to synchronize transaction:')
            yield ledger.release_outputs(tx.outputs)
            raise e

        defer.returnValue(tx)

    @staticmethod
    def signature_hash_type(hash_type):
        return hash_type

    @defer.inlineCallbacks
    def sign(self, funding_accounts: Iterable[BaseAccount]) -> defer.Deferred:
        ledger = self.ensure_all_have_same_ledger(funding_accounts)
        for i, txi in enumerate(self._inputs):
            assert txi.script is not None
            assert txi.txo_ref.txo is not None
            txo_script = txi.txo_ref.txo.script
            if txo_script.is_pay_pubkey_hash:
                address = ledger.hash160_to_address(txo_script.values['pubkey_hash'])
                private_key = yield ledger.get_private_key_for_address(address)
                tx = self._serialize_for_signature(i)
                txi.script.values['signature'] = \
                    private_key.sign(tx) + bytes((self.signature_hash_type(1),))
                txi.script.values['pubkey'] = private_key.public_key.pubkey_bytes
                txi.script.generate()
            else:
                raise NotImplementedError("Don't know how to spend this output.")
        self._reset()