trezor-firmware/trezorlib/btc.py

from . import messages as proto
from .tools import expect, CallException, normalize_nfc, session


@expect(proto.PublicKey)
def get_public_node(client, n, ecdsa_curve_name=None, show_display=False, coin_name=None):
    return client.call(proto.GetPublicKey(address_n=n, ecdsa_curve_name=ecdsa_curve_name, show_display=show_display, coin_name=coin_name))


@expect(proto.Address, field="address")
def get_address(client, coin_name, n, show_display=False, multisig=None, script_type=proto.InputScriptType.SPENDADDRESS):
    if multisig:
        return client.call(proto.GetAddress(address_n=n, coin_name=coin_name, show_display=show_display, multisig=multisig, script_type=script_type))
    else:
        return client.call(proto.GetAddress(address_n=n, coin_name=coin_name, show_display=show_display, script_type=script_type))


@expect(proto.MessageSignature)
def sign_message(client, coin_name, n, message, script_type=proto.InputScriptType.SPENDADDRESS):
    message = normalize_nfc(message)
    return client.call(proto.SignMessage(coin_name=coin_name, address_n=n, message=message, script_type=script_type))


def verify_message(client, coin_name, address, signature, message):
    message = normalize_nfc(message)
    try:
        resp = client.call(proto.VerifyMessage(address=address, signature=signature, message=message, coin_name=coin_name))
    except CallException as e:
        resp = e
    return isinstance(resp, proto.Success)


@session
def sign_tx(client, coin_name, inputs, outputs, version=None, lock_time=None, expiry=None, overwintered=None, debug_processor=None):
    # start = time.time()
    txes = client._prepare_sign_tx(inputs, outputs)

    # Prepare and send initial message
    tx = proto.SignTx()
    tx.inputs_count = len(inputs)
    tx.outputs_count = len(outputs)
    tx.coin_name = coin_name
    if version is not None:
        tx.version = version
    if lock_time is not None:
        tx.lock_time = lock_time
    if expiry is not None:
        tx.expiry = expiry
    if overwintered is not None:
        tx.overwintered = overwintered
    res = client.call(tx)

    # Prepare structure for signatures
    signatures = [None] * len(inputs)
    serialized_tx = b''

    counter = 0
    while True:
        counter += 1

        if isinstance(res, proto.Failure):
            raise CallException("Signing failed")

        if not isinstance(res, proto.TxRequest):
            raise CallException("Unexpected message")

        # If there's some part of signed transaction, let's add it
        if res.serialized and res.serialized.serialized_tx:
            # log("RECEIVED PART OF SERIALIZED TX (%d BYTES)" % len(res.serialized.serialized_tx))
            serialized_tx += res.serialized.serialized_tx

        if res.serialized and res.serialized.signature_index is not None:
            if signatures[res.serialized.signature_index] is not None:
                raise ValueError("Signature for index %d already filled" % res.serialized.signature_index)
            signatures[res.serialized.signature_index] = res.serialized.signature

        if res.request_type == proto.RequestType.TXFINISHED:
            # Device didn't ask for more information, finish workflow
            break

        # Device asked for one more information, let's process it.
        if not res.details.tx_hash:
            current_tx = txes[None]
        else:
            current_tx = txes[bytes(res.details.tx_hash)]

        if res.request_type == proto.RequestType.TXMETA:
            msg = proto.TransactionType()
            msg.version = current_tx.version
            msg.lock_time = current_tx.lock_time
            msg.inputs_cnt = len(current_tx.inputs)
            if res.details.tx_hash:
                msg.outputs_cnt = len(current_tx.bin_outputs)
            else:
                msg.outputs_cnt = len(current_tx.outputs)
            msg.extra_data_len = len(current_tx.extra_data) if current_tx.extra_data else 0
            res = client.call(proto.TxAck(tx=msg))
            continue

        elif res.request_type == proto.RequestType.TXINPUT:
            msg = proto.TransactionType()
            msg.inputs = [current_tx.inputs[res.details.request_index]]
            if debug_processor is not None:
                # msg needs to be deep copied so when it's modified
                # the other messages stay intact
                from copy import deepcopy
                msg = deepcopy(msg)
                # If debug_processor function is provided,
                # pass thru it the request and prepared response.
                # This is useful for tests, see test_msg_signtx
                msg = debug_processor(res, msg)

            res = client.call(proto.TxAck(tx=msg))
            continue

        elif res.request_type == proto.RequestType.TXOUTPUT:
            msg = proto.TransactionType()
            if res.details.tx_hash:
                msg.bin_outputs = [current_tx.bin_outputs[res.details.request_index]]
            else:
                msg.outputs = [current_tx.outputs[res.details.request_index]]

            if debug_processor is not None:
                # msg needs to be deep copied so when it's modified
                # the other messages stay intact
                from copy import deepcopy
                msg = deepcopy(msg)
                # If debug_processor function is provided,
                # pass thru it the request and prepared response.
                # This is useful for tests, see test_msg_signtx
                msg = debug_processor(res, msg)

            res = client.call(proto.TxAck(tx=msg))
            continue

        elif res.request_type == proto.RequestType.TXEXTRADATA:
            o, l = res.details.extra_data_offset, res.details.extra_data_len
            msg = proto.TransactionType()
            msg.extra_data = current_tx.extra_data[o:o + l]
            res = client.call(proto.TxAck(tx=msg))
            continue

    if None in signatures:
        raise RuntimeError("Some signatures are missing!")

    # log("SIGNED IN %.03f SECONDS, CALLED %d MESSAGES, %d BYTES" %
    #    (time.time() - start, counter, len(serialized_tx)))

    return (signatures, serialized_tx)