trezor-firmware/trezorctl

#!/usr/bin/env python3

# This file is part of the TREZOR project.
#
# Copyright (C) 2012-2017 Marek Palatinus <slush@satoshilabs.com>
# Copyright (C) 2012-2017 Pavol Rusnak <stick@satoshilabs.com>
# Copyright (C) 2016-2017 Jochen Hoenicke <hoenicke@gmail.com>
# Copyright (C) 2017      mruddy
#
# This library is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this library.  If not, see <http://www.gnu.org/licenses/>.

import base64
import binascii
import hashlib
import io
import json
import os
import sys

import click

from trezorlib import (
    btc,
    coins,
    cosi,
    debuglink,
    device,
    ethereum,
    firmware,
    lisk,
    log,
    messages as proto,
    misc,
    nem,
    protobuf,
    ripple,
    stellar,
    tools,
)
from trezorlib.client import TrezorClient
from trezorlib.transport import enumerate_devices, get_transport


class ChoiceType(click.Choice):
    def __init__(self, typemap):
        super(ChoiceType, self).__init__(typemap.keys())
        self.typemap = typemap

    def convert(self, value, param, ctx):
        value = super(ChoiceType, self).convert(value, param, ctx)
        return self.typemap[value]


CHOICE_RECOVERY_DEVICE_TYPE = ChoiceType(
    {
        "scrambled": proto.RecoveryDeviceType.ScrambledWords,
        "matrix": proto.RecoveryDeviceType.Matrix,
    }
)

CHOICE_INPUT_SCRIPT_TYPE = ChoiceType(
    {
        "address": proto.InputScriptType.SPENDADDRESS,
        "segwit": proto.InputScriptType.SPENDWITNESS,
        "p2shsegwit": proto.InputScriptType.SPENDP2SHWITNESS,
    }
)

CHOICE_OUTPUT_SCRIPT_TYPE = ChoiceType(
    {
        "address": proto.OutputScriptType.PAYTOADDRESS,
        "segwit": proto.OutputScriptType.PAYTOWITNESS,
        "p2shsegwit": proto.OutputScriptType.PAYTOP2SHWITNESS,
    }
)


def enable_logging():
    log.enable_debug_output()
    log.OMITTED_MESSAGES.add(proto.Features)


@click.group(context_settings={"max_content_width": 400})
@click.option(
    "-p",
    "--path",
    help="Select device by specific path.",
    default=os.environ.get("TREZOR_PATH"),
)
@click.option("-v", "--verbose", is_flag=True, help="Show communication messages.")
@click.option(
    "-j", "--json", "is_json", is_flag=True, help="Print result as JSON object"
)
@click.pass_context
def cli(ctx, path, verbose, is_json):
    if verbose:
        enable_logging()

    def get_device():
        try:
            device = get_transport(path, prefix_search=False)
        except Exception:
            try:
                device = get_transport(path, prefix_search=True)
            except Exception:
                click.echo("Failed to find a TREZOR device.")
                if path is not None:
                    click.echo("Using path: {}".format(path))
                sys.exit(1)
        return TrezorClient(transport=device)

    ctx.obj = get_device


@cli.resultcallback()
def print_result(res, path, verbose, is_json):
    if is_json:
        if isinstance(res, protobuf.MessageType):
            click.echo(json.dumps({res.__class__.__name__: res.__dict__}))
        else:
            click.echo(json.dumps(res, sort_keys=True, indent=4))
    else:
        if isinstance(res, list):
            for line in res:
                click.echo(line)
        elif isinstance(res, dict):
            for k, v in res.items():
                if isinstance(v, dict):
                    for kk, vv in v.items():
                        click.echo("%s.%s: %s" % (k, kk, vv))
                else:
                    click.echo("%s: %s" % (k, v))
        elif isinstance(res, protobuf.MessageType):
            click.echo(protobuf.format_message(res))
        else:
            click.echo(res)


#
# Common functions
#


@cli.command(name="list", help="List connected TREZOR devices.")
def ls():
    return enumerate_devices()


@cli.command(help="Show version of trezorctl/trezorlib.")
def version():
    from trezorlib import __version__ as VERSION

    return VERSION


#
# Basic device functions
#


@cli.command(help="Send ping message.")
@click.argument("message")
@click.option("-b", "--button-protection", is_flag=True)
@click.option("-p", "--pin-protection", is_flag=True)
@click.option("-r", "--passphrase-protection", is_flag=True)
@click.pass_obj
def ping(connect, message, button_protection, pin_protection, passphrase_protection):
    return connect().ping(
        message,
        button_protection=button_protection,
        pin_protection=pin_protection,
        passphrase_protection=passphrase_protection,
    )


@cli.command(help="Clear session (remove cached PIN, passphrase, etc.).")
@click.pass_obj
def clear_session(connect):
    return connect().clear_session()


@cli.command(help="Get example entropy.")
@click.argument("size", type=int)
@click.pass_obj
def get_entropy(connect, size):
    return binascii.hexlify(misc.get_entropy(connect(), size))


@cli.command(help="Retrieve device features and settings.")
@click.pass_obj
def get_features(connect):
    return connect().features


#
# Device management functions
#


@cli.command(help="Change new PIN or remove existing.")
@click.option("-r", "--remove", is_flag=True)
@click.pass_obj
def change_pin(connect, remove):
    return device.change_pin(connect(), remove)


@cli.command(help="Enable passphrase.")
@click.pass_obj
def enable_passphrase(connect):
    return device.apply_settings(connect(), use_passphrase=True)


@cli.command(help="Disable passphrase.")
@click.pass_obj
def disable_passphrase(connect):
    return device.apply_settings(connect(), use_passphrase=False)


@cli.command(help="Set new device label.")
@click.option("-l", "--label")
@click.pass_obj
def set_label(connect, label):
    return device.apply_settings(connect(), label=label)


@cli.command(help="Set passphrase source.")
@click.argument("source", type=int)
@click.pass_obj
def set_passphrase_source(connect, source):
    return device.apply_settings(connect(), passphrase_source=source)


@cli.command(help="Set auto-lock delay (in seconds).")
@click.argument("delay", type=str)
@click.pass_obj
def set_auto_lock_delay(connect, delay):
    value, unit = delay[:-1], delay[-1:]
    units = {"s": 1, "m": 60, "h": 3600}
    if unit in units:
        seconds = float(value) * units[unit]
    else:
        seconds = float(delay)  # assume seconds if no unit is specified
    return device.apply_settings(connect(), auto_lock_delay_ms=int(seconds * 1000))


@cli.command(help="Set device flags.")
@click.argument("flags")
@click.pass_obj
def set_flags(connect, flags):
    flags = flags.lower()
    if flags.startswith("0b"):
        flags = int(flags, 2)
    elif flags.startswith("0x"):
        flags = int(flags, 16)
    else:
        flags = int(flags)
    return device.apply_flags(connect(), flags=flags)


@cli.command(help="Set new homescreen.")
@click.option("-f", "--filename", default=None)
@click.pass_obj
def set_homescreen(connect, filename):
    if filename is None:
        img = b"\x00"
    elif filename.endswith(".toif"):
        img = open(filename, "rb").read()
        if img[:8] != b"TOIf\x90\x00\x90\x00":
            raise tools.CallException(
                proto.FailureType.DataError,
                "File is not a TOIF file with size of 144x144",
            )
    else:
        from PIL import Image

        im = Image.open(filename)
        if im.size != (128, 64):
            raise tools.CallException(
                proto.FailureType.DataError, "Wrong size of the image"
            )
        im = im.convert("1")
        pix = im.load()
        img = bytearray(1024)
        for j in range(64):
            for i in range(128):
                if pix[i, j]:
                    o = i + j * 128
                    img[o // 8] |= 1 << (7 - o % 8)
        img = bytes(img)
    return device.apply_settings(connect(), homescreen=img)


@cli.command(help="Set U2F counter.")
@click.argument("counter", type=int)
@click.pass_obj
def set_u2f_counter(connect, counter):
    return device.set_u2f_counter(connect(), counter)


@cli.command(help="Reset device to factory defaults and remove all private data.")
@click.option(
    "-b",
    "--bootloader",
    help="Wipe device in bootloader mode. This also erases the firmware.",
    is_flag=True,
)
@click.pass_obj
def wipe_device(connect, bootloader):
    client = connect()
    if bootloader:
        if not client.features.bootloader_mode:
            click.echo("Please switch your device to bootloader mode.")
            sys.exit(1)
        else:
            click.echo(
                "Wiping user data and firmware! Please confirm the action on your device ..."
            )
    else:
        if client.features.bootloader_mode:
            click.echo(
                "Your device is in bootloader mode. This operation would also erase firmware."
            )
            click.echo(
                'Specify "--bootloader" if that is what you want, or disconnect and reconnect device in normal mode.'
            )
            click.echo("Aborting.")
            sys.exit(1)
        else:
            click.echo("Wiping user data! Please confirm the action on your device ...")

    try:
        return device.wipe(connect())
    except tools.CallException as e:
        click.echo("Action failed: {} {}".format(*e.args))
        sys.exit(3)


@cli.command(help="Load custom configuration to the device.")
@click.option("-m", "--mnemonic")
@click.option("-e", "--expand", is_flag=True)
@click.option("-x", "--xprv")
@click.option("-p", "--pin", default="")
@click.option("-r", "--passphrase-protection", is_flag=True)
@click.option("-l", "--label", default="")
@click.option("-i", "--ignore-checksum", is_flag=True)
@click.option("-s", "--slip0014", is_flag=True)
@click.pass_obj
def load_device(
    connect,
    mnemonic,
    expand,
    xprv,
    pin,
    passphrase_protection,
    label,
    ignore_checksum,
    slip0014,
):
    if not mnemonic and not xprv and not slip0014:
        raise tools.CallException(
            proto.FailureType.DataError, "Please provide mnemonic or xprv"
        )

    client = connect()
    if mnemonic:
        return debuglink.load_device_by_mnemonic(
            client,
            mnemonic,
            pin,
            passphrase_protection,
            label,
            "english",
            ignore_checksum,
            expand,
        )
    if xprv:
        return debuglink.load_device_by_xprv(
            client, xprv, pin, passphrase_protection, label, "english"
        )
    if slip0014:
        return debuglink.load_device_by_mnemonic(
            client, " ".join(["all"] * 12), pin, passphrase_protection, "SLIP-0014"
        )


@cli.command(help="Start safe recovery workflow.")
@click.option("-w", "--words", type=click.Choice(["12", "18", "24"]), default="24")
@click.option("-e", "--expand", is_flag=True)
@click.option("-p", "--pin-protection", is_flag=True)
@click.option("-r", "--passphrase-protection", is_flag=True)
@click.option("-l", "--label")
@click.option(
    "-t", "--type", "rec_type", type=CHOICE_RECOVERY_DEVICE_TYPE, default="scrambled"
)
@click.option("-d", "--dry-run", is_flag=True)
@click.pass_obj
def recovery_device(
    connect,
    words,
    expand,
    pin_protection,
    passphrase_protection,
    label,
    rec_type,
    dry_run,
):
    return device.recover(
        connect(),
        int(words),
        passphrase_protection,
        pin_protection,
        label,
        "english",
        rec_type,
        expand,
        dry_run,
    )


@cli.command(help="Perform device setup and generate new seed.")
@click.option("-e", "--entropy", is_flag=True)
@click.option(
    "-t", "--strength", type=click.Choice(["128", "192", "256"]), default="256"
)
@click.option("-r", "--passphrase-protection", is_flag=True)
@click.option("-p", "--pin-protection", is_flag=True)
@click.option("-l", "--label")
@click.option("-u", "--u2f-counter", default=0)
@click.option("-s", "--skip-backup", is_flag=True)
@click.pass_obj
def reset_device(
    connect,
    entropy,
    strength,
    passphrase_protection,
    pin_protection,
    label,
    u2f_counter,
    skip_backup,
):
    return device.reset(
        connect(),
        entropy,
        int(strength),
        passphrase_protection,
        pin_protection,
        label,
        "english",
        u2f_counter,
        skip_backup,
    )


@cli.command(help="Perform device seed backup.")
@click.pass_obj
def backup_device(connect):
    return device.backup(connect())


#
# Firmware update
#


@cli.command(help="Upload new firmware to device (must be in bootloader mode).")
@click.option("-f", "--filename")
@click.option("-u", "--url")
@click.option("-v", "--version")
@click.option("-s", "--skip-check", is_flag=True)
@click.option("--fingerprint", help="Expected firmware fingerprint in hex")
@click.pass_obj
def firmware_update(connect, filename, url, version, skip_check, fingerprint):
    if sum(bool(x) for x in (filename, url, version)) > 1:
        click.echo("You can use only one of: firmware, url, version.")
        sys.exit(1)

    client = connect()
    if not client.features.bootloader_mode:
        click.echo("Please switch your device to bootloader mode.")
        sys.exit(1)

    firmware_version = client.features.major_version

    if filename:
        fp = open(filename, "rb").read()
    elif url:
        import requests

        click.echo("Downloading from", url)
        r = requests.get(url)
        fp = r.content
    else:
        import requests

        r = requests.get(
            "https://wallet.trezor.io/data/firmware/{}/releases.json".format(
                firmware_version
            )
        )
        releases = r.json()

        def version_func(r):
            return r["version"]

        def version_string(r):
            return ".".join(map(str, version_func(r)))

        if version:
            try:
                release = next(r for r in releases if version_string(r) == version)
            except StopIteration:
                click.echo("Version {} not found.".format(version))
                sys.exit(1)
        else:
            release = max(releases, key=version_func)
            click.echo("Fetching version: %s" % version_string(release))

        if not fingerprint:
            fingerprint = release["fingerprint"]
        url = "https://wallet.trezor.io/" + release["url"]
        click.echo("Downloading from %s" % url)
        r = requests.get(url)
        fp = r.content

    if not skip_check:
        if fp[:8] == b"54525a52" or fp[:8] == b"54525a56":
            fp = binascii.unhexlify(fp)
        if fp[:4] != b"TRZR" and fp[:4] != b"TRZV":
            click.echo("Trezor firmware header expected.")
            sys.exit(2)

        if fingerprint and firmware_version > 1:
            click.echo("Checking Trezor T fingerprint is not supported yet.")
        elif firmware_version == 1:
            calculated_fingerprint = hashlib.sha256(fp[256:]).hexdigest()
            click.echo("Firmware fingerprint: {}".format(calculated_fingerprint))
            if fingerprint and fingerprint != calculated_fingerprint:
                click.echo("Expected fingerprint: {}".format(fingerprint))
                click.echo("Fingerprints do not match, aborting.")
                sys.exit(5)

    click.echo("If asked, please confirm the action on your device ...")

    try:
        return firmware.update(client, fp=io.BytesIO(fp))
    except tools.CallException as e:
        if e.args[0] in (
            proto.FailureType.FirmwareError,
            proto.FailureType.ActionCancelled,
        ):
            click.echo("Update aborted on device.")
        else:
            click.echo("Update failed: {} {}".format(*e.args))
        sys.exit(3)


@cli.command(help="Perform a self-test.")
@click.pass_obj
def self_test(connect):
    return debuglink.self_test(connect())


#
# Basic coin functions
#


@cli.command(help="Get address for specified path.")
@click.option("-c", "--coin", default="Bitcoin")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/0'/0'/0/0"
)
@click.option("-t", "--script-type", type=CHOICE_INPUT_SCRIPT_TYPE, default="address")
@click.option("-d", "--show-display", is_flag=True)
@click.pass_obj
def get_address(connect, coin, address, script_type, show_display):
    client = connect()
    address_n = tools.parse_path(address)
    return btc.get_address(
        client, coin, address_n, show_display, script_type=script_type
    )


@cli.command(help="Get public node of given path.")
@click.option("-c", "--coin", default="Bitcoin")
@click.option("-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/0'/0'")
@click.option("-e", "--curve")
@click.option("-d", "--show-display", is_flag=True)
@click.pass_obj
def get_public_node(connect, coin, address, curve, show_display):
    client = connect()
    address_n = tools.parse_path(address)
    result = btc.get_public_node(
        client,
        address_n,
        ecdsa_curve_name=curve,
        show_display=show_display,
        coin_name=coin,
    )
    return {
        "node": {
            "depth": result.node.depth,
            "fingerprint": "%08x" % result.node.fingerprint,
            "child_num": result.node.child_num,
            "chain_code": binascii.hexlify(result.node.chain_code),
            "public_key": binascii.hexlify(result.node.public_key),
        },
        "xpub": result.xpub,
    }


#
# Signing options
#


@cli.command(help="Sign transaction.")
@click.option("-c", "--coin", default="Bitcoin")
# @click.option('-n', '--address', required=True, help="BIP-32 path, e.g. m/44'/0'/0'/0/0")
# @click.option('-t', '--script-type', type=CHOICE_INPUT_SCRIPT_TYPE, default='address')
# @click.option('-o', '--output', required=True, help='Transaction output')
# @click.option('-f', '--fee', required=True, help='Transaction fee (sat/B)')
@click.pass_obj
def sign_tx(connect, coin):
    client = connect()
    if coin in coins.tx_api:
        txapi = coins.tx_api[coin]
    else:
        click.echo('Coin "%s" is not recognized.' % coin, err=True)
        click.echo(
            "Supported coin types: %s" % ", ".join(coins.tx_api.keys()), err=True
        )
        sys.exit(1)

    client.set_tx_api(txapi)

    def default_script_type(address_n):
        script_type = "address"

        if address_n is None:
            pass
        elif address_n[0] == tools.H_(49):
            script_type = "p2shsegwit"

        return script_type

    def outpoint(s):
        txid, vout = s.split(":")
        return binascii.unhexlify(txid), int(vout)

    inputs = []
    while True:
        click.echo()
        prev = click.prompt(
            "Previous output to spend (txid:vout)", type=outpoint, default=""
        )
        if not prev:
            break
        prev_hash, prev_index = prev
        address_n = click.prompt("BIP-32 path to derive the key", type=tools.parse_path)
        amount = click.prompt("Input amount (satoshis)", type=int, default=0)
        sequence = click.prompt(
            "Sequence Number to use (RBF opt-in enabled by default)",
            type=int,
            default=0xfffffffd,
        )
        script_type = click.prompt(
            "Input type",
            type=CHOICE_INPUT_SCRIPT_TYPE,
            default=default_script_type(address_n),
        )
        script_type = (
            script_type
            if isinstance(script_type, int)
            else CHOICE_INPUT_SCRIPT_TYPE.typemap[script_type]
        )
        if txapi.bip115:
            prev_output = txapi.get_tx(
                binascii.hexlify(prev_hash).decode("utf-8")
            ).bin_outputs[prev_index]
            prev_blockhash = prev_output.block_hash
            prev_blockheight = prev_output.block_height
        inputs.append(
            proto.TxInputType(
                address_n=address_n,
                prev_hash=prev_hash,
                prev_index=prev_index,
                amount=amount,
                script_type=script_type,
                sequence=sequence,
                prev_block_hash_bip115=prev_blockhash,
                prev_block_height_bip115=prev_blockheight,
            )
        )

    if txapi.bip115:
        current_block_height = txapi.current_height()
        block_height = (
            current_block_height - 300
        )  # Zencash recommendation for the better protection
        block_hash = txapi.get_block_hash(block_height)

    outputs = []
    while True:
        click.echo()
        address = click.prompt("Output address (for non-change output)", default="")
        if address:
            address_n = None
        else:
            address = None
            address_n = click.prompt(
                "BIP-32 path (for change output)", type=tools.parse_path, default=""
            )
            if not address_n:
                break
        amount = click.prompt("Amount to spend (satoshis)", type=int)
        script_type = click.prompt(
            "Output type",
            type=CHOICE_OUTPUT_SCRIPT_TYPE,
            default=default_script_type(address_n),
        )
        script_type = (
            script_type
            if isinstance(script_type, int)
            else CHOICE_OUTPUT_SCRIPT_TYPE.typemap[script_type]
        )
        outputs.append(
            proto.TxOutputType(
                address_n=address_n,
                address=address,
                amount=amount,
                script_type=script_type,
                block_hash_bip115=block_hash[::-1],  # Blockhash passed in reverse order
                block_height_bip115=block_height,
            )
        )

    tx_version = click.prompt("Transaction version", type=int, default=2)
    tx_locktime = click.prompt("Transaction locktime", type=int, default=0)

    _, serialized_tx = btc.sign_tx(
        client, coin, inputs, outputs, tx_version, tx_locktime
    )

    client.close()

    click.echo()
    click.echo("Signed Transaction:")
    click.echo(binascii.hexlify(serialized_tx))
    click.echo()
    click.echo("Use the following form to broadcast it to the network:")
    click.echo(txapi.pushtx_url)


#
# Message functions
#


@cli.command(help="Sign message using address of given path.")
@click.option("-c", "--coin", default="Bitcoin")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/0'/0'/0/0"
)
@click.option(
    "-t",
    "--script-type",
    type=click.Choice(["address", "segwit", "p2shsegwit"]),
    default="address",
)
@click.argument("message")
@click.pass_obj
def sign_message(connect, coin, address, message, script_type):
    client = connect()
    address_n = tools.parse_path(address)
    typemap = {
        "address": proto.InputScriptType.SPENDADDRESS,
        "segwit": proto.InputScriptType.SPENDWITNESS,
        "p2shsegwit": proto.InputScriptType.SPENDP2SHWITNESS,
    }
    script_type = typemap[script_type]
    res = btc.sign_message(client, coin, address_n, message, script_type)
    return {
        "message": message,
        "address": res.address,
        "signature": base64.b64encode(res.signature),
    }


@cli.command(help="Verify message.")
@click.option("-c", "--coin", default="Bitcoin")
@click.argument("address")
@click.argument("signature")
@click.argument("message")
@click.pass_obj
def verify_message(connect, coin, address, signature, message):
    signature = base64.b64decode(signature)
    return btc.verify_message(connect(), coin, address, signature, message)


@cli.command(help="Sign message with Ethereum address.")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/60'/0'/0/0"
)
@click.argument("message")
@click.pass_obj
def ethereum_sign_message(connect, address, message):
    client = connect()
    address_n = tools.parse_path(address)
    ret = ethereum.sign_message(client, address_n, message)
    output = {
        "message": message,
        "address": "0x%s" % binascii.hexlify(ret.address).decode(),
        "signature": "0x%s" % binascii.hexlify(ret.signature).decode(),
    }
    return output


def ethereum_decode_hex(value):
    if value.startswith("0x") or value.startswith("0X"):
        return binascii.unhexlify(value[2:])
    else:
        return binascii.unhexlify(value)


@cli.command(help="Verify message signed with Ethereum address.")
@click.argument("address")
@click.argument("signature")
@click.argument("message")
@click.pass_obj
def ethereum_verify_message(connect, address, signature, message):
    address = ethereum_decode_hex(address)
    signature = ethereum_decode_hex(signature)
    return ethereum.verify_message(connect(), address, signature, message)


@cli.command(help="Encrypt value by given key and path.")
@click.option("-n", "--address", required=True, help="BIP-32 path, e.g. m/10016'/0")
@click.argument("key")
@click.argument("value")
@click.pass_obj
def encrypt_keyvalue(connect, address, key, value):
    client = connect()
    address_n = tools.parse_path(address)
    res = misc.encrypt_keyvalue(client, address_n, key, value.encode())
    return binascii.hexlify(res)


@cli.command(help="Decrypt value by given key and path.")
@click.option("-n", "--address", required=True, help="BIP-32 path, e.g. m/10016'/0")
@click.argument("key")
@click.argument("value")
@click.pass_obj
def decrypt_keyvalue(connect, address, key, value):
    client = connect()
    address_n = tools.parse_path(address)
    return misc.decrypt_keyvalue(client, address_n, key, binascii.unhexlify(value))


# @cli.command(help='Encrypt message.')
# @click.option('-c', '--coin', default='Bitcoin')
# @click.option('-d', '--display-only', is_flag=True)
# @click.option('-n', '--address', required=True, help="BIP-32 path, e.g. m/44'/0'/0'/0/0")
# @click.argument('pubkey')
# @click.argument('message')
# @click.pass_obj
# def encrypt_message(connect, coin, display_only, address, pubkey, message):
#     client = connect()
#     pubkey = binascii.unhexlify(pubkey)
#     address_n = tools.parse_path(address)
#     res = client.encrypt_message(pubkey, message, display_only, coin, address_n)
#     return {
#         'nonce': binascii.hexlify(res.nonce),
#         'message': binascii.hexlify(res.message),
#         'hmac': binascii.hexlify(res.hmac),
#         'payload': base64.b64encode(res.nonce + res.message + res.hmac),
#     }


# @cli.command(help='Decrypt message.')
# @click.option('-n', '--address', required=True, help="BIP-32 path, e.g. m/44'/0'/0'/0/0")
# @click.argument('payload')
# @click.pass_obj
# def decrypt_message(connect, address, payload):
#     client = connect()
#     address_n = tools.parse_path(address)
#     payload = base64.b64decode(payload)
#     nonce, message, msg_hmac = payload[:33], payload[33:-8], payload[-8:]
#     return client.decrypt_message(address_n, nonce, message, msg_hmac)


#
# Ethereum functions
#


@cli.command(help="Get Ethereum address in hex encoding.")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/60'/0'/0/0"
)
@click.option("-d", "--show-display", is_flag=True)
@click.pass_obj
def ethereum_get_address(connect, address, show_display):
    client = connect()
    address_n = tools.parse_path(address)
    address = ethereum.get_address(client, address_n, show_display)
    return "0x%s" % binascii.hexlify(address).decode()


@cli.command(
    help='Sign (and optionally publish) Ethereum transaction. Use TO as destination address or set TO to "" for contract creation.'
)
@click.option(
    "-a",
    "--host",
    default="localhost:8545",
    help="RPC port of ethereum node for automatic gas/nonce estimation and publishing",
)
@click.option("-c", "--chain-id", type=int, help="EIP-155 chain id (replay protection)")
@click.option(
    "-n",
    "--address",
    required=True,
    help="BIP-32 path to source address, e.g., m/44'/60'/0'/0/0",
)
@click.option(
    "-v", "--value", default="0", help='Ether amount to transfer, e.g. "100 milliether"'
)
@click.option(
    "-g", "--gas-limit", type=int, help="Gas limit - Required for offline signing"
)
@click.option(
    "-t",
    "--gas-price",
    help='Gas price, e.g. "20 nanoether" - Required for offline signing',
)
@click.option(
    "-i", "--nonce", type=int, help="Transaction counter - Required for offline signing"
)
@click.option("-d", "--data", default="", help="Data as hex string, e.g. 0x12345678")
@click.option("-p", "--publish", is_flag=True, help="Publish transaction via RPC")
@click.option("-x", "--tx-type", type=int, help="TX type (used only for Wanchain)")
@click.argument("to")
@click.pass_obj
def ethereum_sign_tx(
    connect,
    host,
    chain_id,
    address,
    value,
    gas_limit,
    gas_price,
    nonce,
    data,
    publish,
    to,
    tx_type,
):
    from ethjsonrpc import EthJsonRpc
    import rlp

    # fmt: off
    ether_units = {
        'wei':          1,
        'kwei':         1000,
        'babbage':      1000,
        'femtoether':   1000,
        'mwei':         1000000,
        'lovelace':     1000000,
        'picoether':    1000000,
        'gwei':         1000000000,
        'shannon':      1000000000,
        'nanoether':    1000000000,
        'nano':         1000000000,
        'szabo':        1000000000000,
        'microether':   1000000000000,
        'micro':        1000000000000,
        'finney':       1000000000000000,
        'milliether':   1000000000000000,
        'milli':        1000000000000000,
        'ether':        1000000000000000000,
        'eth':          1000000000000000000,
    }
    # fmt: on

    if " " in value:
        value, unit = value.split(" ", 1)
        if unit.lower() not in ether_units:
            raise tools.CallException(
                proto.Failure.DataError, "Unrecognized ether unit %r" % unit
            )
        value = int(value) * ether_units[unit.lower()]
    else:
        value = int(value)

    if gas_price is not None:
        if " " in gas_price:
            gas_price, unit = gas_price.split(" ", 1)
            if unit.lower() not in ether_units:
                raise tools.CallException(
                    proto.Failure.DataError, "Unrecognized gas price unit %r" % unit
                )
            gas_price = int(gas_price) * ether_units[unit.lower()]
        else:
            gas_price = int(gas_price)

    if gas_limit is not None:
        gas_limit = int(gas_limit)

    to_address = ethereum_decode_hex(to)

    client = connect()
    address_n = tools.parse_path(address)
    address = "0x%s" % (
        binascii.hexlify(ethereum.get_address(client, address_n)).decode()
    )

    if gas_price is None or gas_limit is None or nonce is None or publish:
        host, port = host.split(":")
        eth = EthJsonRpc(host, int(port))

    if not data:
        data = ""
    data = ethereum_decode_hex(data)

    if gas_price is None:
        gas_price = eth.eth_gasPrice()

    if gas_limit is None:
        gas_limit = eth.eth_estimateGas(
            to_address=to,
            from_address=address,
            value=("0x%x" % value),
            data="0x%s" % (binascii.hexlify(data).decode()),
        )

    if nonce is None:
        nonce = eth.eth_getTransactionCount(address)

    sig = ethereum.sign_tx(
        client,
        n=address_n,
        tx_type=tx_type,
        nonce=nonce,
        gas_price=gas_price,
        gas_limit=gas_limit,
        to=to_address,
        value=value,
        data=data,
        chain_id=chain_id,
    )

    if tx_type is None:
        transaction = rlp.encode(
            (nonce, gas_price, gas_limit, to_address, value, data) + sig
        )
    else:
        transaction = rlp.encode(
            (tx_type, nonce, gas_price, gas_limit, to_address, value, data) + sig
        )
    tx_hex = "0x%s" % binascii.hexlify(transaction).decode()

    if publish:
        tx_hash = eth.eth_sendRawTransaction(tx_hex)
        return "Transaction published with ID: %s" % tx_hash
    else:
        return "Signed raw transaction: %s" % tx_hex


#
# NEM functions
#


@cli.command(help="Get NEM address for specified path.")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/0'/43'/0/0"
)
@click.option("-N", "--network", type=int, default=0x68)
@click.option("-d", "--show-display", is_flag=True)
@click.pass_obj
def nem_get_address(connect, address, network, show_display):
    client = connect()
    address_n = tools.parse_path(address)
    return nem.get_address(client, address_n, network, show_display)


@cli.command(help="Sign (and optionally broadcast) NEM transaction.")
@click.option("-n", "--address", help="BIP-32 path to signing key")
@click.option(
    "-f",
    "--file",
    type=click.File("r"),
    default="-",
    help="Transaction in NIS (RequestPrepareAnnounce) format",
)
@click.option("-b", "--broadcast", help="NIS to announce transaction to")
@click.pass_obj
def nem_sign_tx(connect, address, file, broadcast):
    client = connect()
    address_n = tools.parse_path(address)
    transaction = nem.sign_tx(client, address_n, json.load(file))

    payload = {
        "data": binascii.hexlify(transaction.data).decode(),
        "signature": binascii.hexlify(transaction.signature).decode(),
    }

    if broadcast:
        import requests

        return requests.post(
            "{}/transaction/announce".format(broadcast), json=payload
        ).json()
    else:
        return payload


#
# Lisk functions
#


@cli.command(help="Get Lisk address for specified path.")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/134'/0'/0'"
)
@click.option("-d", "--show-display", is_flag=True)
@click.pass_obj
def lisk_get_address(connect, address, show_display):
    client = connect()
    address_n = tools.parse_path(address)
    return lisk.get_address(client, address_n, show_display)


@cli.command(help="Get Lisk public key for specified path.")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/134'/0'/0'"
)
@click.option("-d", "--show-display", is_flag=True)
@click.pass_obj
def lisk_get_public_key(connect, address, show_display):
    client = connect()
    address_n = tools.parse_path(address)
    res = lisk.get_public_key(client, address_n, show_display)
    output = {"public_key": binascii.hexlify(res.public_key).decode()}
    return output


@cli.command(help="Sign Lisk transaction.")
@click.option(
    "-n",
    "--address",
    required=True,
    help="BIP-32 path to signing key, e.g. m/44'/134'/0'/0'",
)
@click.option(
    "-f", "--file", type=click.File("r"), default="-", help="Transaction in JSON format"
)
# @click.option('-b', '--broadcast', help='Broadcast Lisk transaction')
@click.pass_obj
def lisk_sign_tx(connect, address, file):
    client = connect()
    address_n = tools.parse_path(address)
    transaction = lisk.sign_tx(client, address_n, json.load(file))

    payload = {"signature": binascii.hexlify(transaction.signature).decode()}

    return payload


@cli.command(help="Sign message with Lisk address.")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/134'/0'/0'"
)
@click.argument("message")
@click.pass_obj
def lisk_sign_message(connect, address, message):
    client = connect()
    address_n = client.expand_path(address)
    res = lisk.sign_message(client, address_n, message)
    output = {
        "message": message,
        "public_key": binascii.hexlify(res.public_key).decode(),
        "signature": binascii.hexlify(res.signature).decode(),
    }
    return output


@cli.command(help="Verify message signed with Lisk address.")
@click.argument("pubkey")
@click.argument("signature")
@click.argument("message")
@click.pass_obj
def lisk_verify_message(connect, pubkey, signature, message):
    signature = bytes.fromhex(signature)
    pubkey = bytes.fromhex(pubkey)
    return lisk.verify_message(connect(), pubkey, signature, message)


#
# CoSi functions
#


@cli.command(help="Ask device to commit to CoSi signing.")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/0'/0'/0/0"
)
@click.argument("data")
@click.pass_obj
def cosi_commit(connect, address, data):
    client = connect()
    address_n = tools.parse_path(address)
    return cosi.commit(client, address_n, binascii.unhexlify(data))


@cli.command(help="Ask device to sign using CoSi.")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path, e.g. m/44'/0'/0'/0/0"
)
@click.argument("data")
@click.argument("global_commitment")
@click.argument("global_pubkey")
@click.pass_obj
def cosi_sign(connect, address, data, global_commitment, global_pubkey):
    client = connect()
    address_n = tools.parse_path(address)
    return cosi.sign(
        client,
        address_n,
        binascii.unhexlify(data),
        binascii.unhexlify(global_commitment),
        binascii.unhexlify(global_pubkey),
    )


#
# Stellar functions
#
@cli.command(help="Get Stellar public address")
@click.option(
    "-n",
    "--address",
    required=False,
    help="BIP32 path. Always use hardened paths and the m/44'/148'/ prefix",
    default=stellar.DEFAULT_BIP32_PATH,
)
@click.option("-d", "--show-display", is_flag=True)
@click.pass_obj
def stellar_get_address(connect, address, show_display):
    client = connect()
    address_n = tools.parse_path(address)
    return stellar.get_address(client, address_n, show_display)


@cli.command(help="Get Stellar public key")
@click.option(
    "-n",
    "--address",
    required=False,
    help="BIP32 path. Always use hardened paths and the m/44'/148'/ prefix",
    default=stellar.DEFAULT_BIP32_PATH,
)
@click.option("-d", "--show-display", is_flag=True)
@click.pass_obj
def stellar_get_public_key(connect, address, show_display):
    client = connect()
    address_n = tools.parse_path(address)
    return binascii.hexlify(stellar.get_public_key(client, address_n, show_display))


@cli.command(help="Sign a base64-encoded transaction envelope")
@click.option(
    "-n",
    "--address",
    required=False,
    help="BIP32 path. Always use hardened paths and the m/44'/148'/ prefix",
    default=stellar.DEFAULT_BIP32_PATH,
)
@click.option(
    "-n",
    "--network-passphrase",
    default=stellar.DEFAULT_NETWORK_PASSPHRASE,
    required=False,
    help="Network passphrase (blank for public network). Testnet is: 'Test SDF Network ; September 2015'",
)
@click.argument("b64envelope")
@click.pass_obj
def stellar_sign_transaction(connect, b64envelope, address, network_passphrase):
    client = connect()
    address_n = tools.parse_path(address)
    tx, operations = stellar.parse_transaction_bytes(base64.b64decode(b64envelope))
    resp = stellar.sign_tx(client, tx, operations, address_n, network_passphrase)

    return base64.b64encode(resp.signature)


#
# Ripple functions
#
@cli.command(help="Get Ripple address")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path to key, e.g. m/44'/144'/0'/0/0"
)
@click.option("-d", "--show-display", is_flag=True)
@click.pass_obj
def ripple_get_address(connect, address, show_display):
    client = connect()
    address_n = tools.parse_path(address)
    return ripple.get_address(client, address_n, show_display)


@cli.command(help="Sign Ripple transaction")
@click.option(
    "-n", "--address", required=True, help="BIP-32 path to key, e.g. m/44'/144'/0'/0/0"
)
@click.option(
    "-f", "--file", type=click.File("r"), default="-", help="Transaction in JSON format"
)
@click.pass_obj
def ripple_sign_tx(connect, address, file):
    client = connect()
    address_n = tools.parse_path(address)
    msg = ripple.create_sign_tx_msg(json.load(file))

    result = ripple.sign_tx(client, address_n, msg)
    click.echo("Signature:")
    click.echo(binascii.hexlify(result.signature))
    click.echo()
    click.echo("Serialized tx including the signature:")
    click.echo(binascii.hexlify(result.serialized_tx))


#
# Main
#


if __name__ == "__main__":
    cli()  # pylint: disable=E1120