python-trezor/cmd.py

#!/usr/bin/python
import os
import binascii
import argparse
import json
import threading

from trezorlib.client import TrezorClient, pin_func
from trezorlib.debuglink import DebugLink
from trezorlib.protobuf_json import pb2json
from trezorlib.pinmatrix import PinMatrixWidget

def parse_args(commands):
    parser = argparse.ArgumentParser(description='Commandline tool for Trezor devices.')
    parser.add_argument('-t', '--transport', dest='transport',  choices=['usb', 'serial', 'pipe', 'socket'], default='usb', help="Transport used for talking with the device")
    parser.add_argument('-p', '--path', dest='path', default='', help="Path used by the transport (usually serial port)")
    parser.add_argument('-dt', '--debuglink-transport', dest='debuglink_transport', choices=['usb', 'serial', 'pipe', 'socket'], default='socket', help="Debuglink transport")
    parser.add_argument('-dp', '--debuglink-path', dest='debuglink_path', default='127.0.0.1:2000', help="Path used by the transport (usually serial port)")         
    parser.add_argument('-j', '--json', dest='json', action='store_true', help="Prints result as json object")
    parser.add_argument('-d', '--debug', dest='debug', action='store_true', help='Enable low-level debugging')

    cmdparser = parser.add_subparsers(title='Available commands')
    
    for cmd in commands._list_commands():
        func = object.__getattribute__(commands, cmd)
        try:
            help = func.help
        except AttributeError:
            help = ''
            
        try:
            arguments = func.arguments
        except AttributeError:
            arguments = ((('params',), {'nargs': '*'}),)
        
        item = cmdparser.add_parser(cmd, help=func.help)
        for arg in arguments:
            item.add_argument(*arg[0], **arg[1])
            
        item.set_defaults(func=func)
        item.set_defaults(cmd=cmd)
    
    return parser.parse_args()

def get_transport(transport_string, path):
    if transport_string == 'usb':
        from trezorlib.transport_hid import HidTransport

        if path == '':
            try:
                path = list_usb()[0]
            except IndexError:
                raise Exception("No Trezor found on USB")

        return HidTransport(path)
 
    if transport_string == 'serial':
        from trezorlib.transport_serial import SerialTransport
        return SerialTransport(path)

    if transport_string == 'pipe':
        from trezorlib.transport_pipe import PipeTransport
        return PipeTransport(path, is_device=False)
    
    if transport_string == 'socket':
        from trezorlib.transport_socket import SocketTransportClient
        return SocketTransportClient(path)
    
    if transport_string == 'fake':
        from trezorlib.transport_fake import FakeTransport
        return FakeTransport(path)
    
    raise NotImplemented("Unknown transport")

class Commands(object):
    def __init__(self, client):
        self.client = client
        
    @classmethod
    def _list_commands(cls):
        return [ x for x in dir(cls) if not x.startswith('_') ]
  
    def list(self, args):
        # Fake method for advertising 'list' command
        pass
 
    def get_address(self, args):
        return self.client.get_address(args.n)
    
    def get_entropy(self, args):
        return binascii.hexlify(self.client.get_entropy(args.size))

    def get_features(self, args):
        return pb2json(self.client.features)

    def ping(self, args):
        return self.client.ping(args.msg)

    def get_master_public_key(self, args):
        return self.client.get_master_public_key()
    
    def get_serial_number(self, args):
        return binascii.hexlify(self.client.get_serial_number())

    def set_label(self, args):
        return self.client.apply_settings(label=args.label)

    def set_coin(self, args):
        return self.client.apply_settings(coin_shortcut=args.coin_shortcut)

    def load_device(self, args):
        seed = ' '.join(args.seed)

        return self.client.load_device(seed, args.pin) 

    def firmware_update(self, args):
        if not args.file:
            raise Exception("Must provide firmware filename")
        fp = open(args.file, 'r')
        if fp.read(4) != 'TRZR':
            raise Exception("Trezor firmware header expected")

        fp.seek(0)
        return self.client.firmware_update(fp=open(args.file, 'r'), force=args.force)

    list.help = 'List connected Trezor USB devices'
    ping.help = 'Send ping message'
    get_address.help = 'Get bitcoin address in base58 encoding'
    get_entropy.help = 'Get example entropy'
    get_features.help = 'Retrieve device features and settings'
    get_serial_number.help = 'Get device\'s unique identifier'
    get_master_public_key.help = 'Get master public key'
    set_label.help = 'Set new wallet label'
    set_coin.help = 'Switch device to another crypto currency'
    load_device.help = 'Load custom configuration to the device'
    firmware_update.help = 'Upload new firmware to device (must be in bootloader mode)'

    get_address.arguments = (
        (('n',), {'metavar': 'N', 'type': int, 'nargs': '+'}),
    )
    
    get_entropy.arguments = (
        (('size',), {'type': int}),
    )

    get_features.arguments = ()

    ping.arguments = (
        (('msg',), {'type': str}),
    )
    
    set_label.arguments = (
        (('label',), {'type': str}),
    )

    set_coin.arguments = (
        (('coin_shortcut',), {'type': str}),
    )

    load_device.arguments = (
        (('-s', '--seed'), {'type': str, 'nargs': '+'}),
        (('-n', '--pin'), {'type': str, 'default': ''}),
    )

    firmware_update.arguments = (
        (('-f', '--file'), {'type': str}),
        (('-o', '--force'), {'type': bool, 'default': False})
    )

def list_usb():
    from trezorlib.transport_hid import HidTransport
    devices = HidTransport.enumerate()
    return devices

class PinMatrixThread(threading.Thread):
    '''
        Hacked PinMatrixWidget into command line tool :-).
    '''
    def __init__(self, input_text, message):
        super(PinMatrixThread, self).__init__()
        self.input_text = input_text
        self.message = message
        self.pin_value = ''

    def run(self):
        import sys
        from PyQt4.Qt import QApplication, QWidget, QVBoxLayout
        from PyQt4.QtGui import QPushButton, QLabel
        from PyQt4.QtCore import QObject, SIGNAL

        a = QApplication(sys.argv)
        matrix = PinMatrixWidget()

        def clicked():
            self.pin_value = str(matrix.get_value())
            a.closeAllWindows()

        ok = QPushButton('OK')
        QObject.connect(ok, SIGNAL('clicked()'), clicked)

        vbox = QVBoxLayout()
        vbox.addWidget(QLabel(self.input_text + self.message))
        vbox.addWidget(matrix)
        vbox.addWidget(ok)

        w = QWidget()
        w.setLayout(vbox)
        w.move(100, 100)
        w.show()

        a.exec_()

def qt_pin_func(input_text, message=None):
    '''
        This is a hack to display Qt window in non-qt application.
        Qt window just asks for PIN and closes itself, which trigger join().
    '''
    if os.getenv('DISPLAY'):
        # Let's hope that system is configured properly and this won't crash
        t = PinMatrixThread(input_text, message)
        t.start()
        t.join()
        return t.pin_value
    else:
        # Most likely no X is running,
        # let's fallback to default pin_func implementation
        return pin_func(input_text, message)

def main():
    args = parse_args(Commands)

    if args.cmd == 'list':
        devices = list_usb()
        if args.json:
            print json.dumps(devices)
        else:
            for dev in devices:
                print dev
        return

    transport = get_transport(args.transport, args.path)
    if args.debug:
        debuglink_transport = get_transport(args.debuglink_transport, args.debuglink_path)
        debuglink = DebugLink(debuglink_transport)    
    else:
        debuglink = None
        
    client = TrezorClient(transport, pin_func=qt_pin_func, debuglink=debuglink)
    client.setup_debuglink(button=True, pin_correct=True)
    cmds = Commands(client)
    
    res = args.func(cmds, args)
    
    if args.json:
        print json.dumps(res, sort_keys=True, indent=4)
    else:
        print res
    
if __name__ == '__main__':
    main()