260 lines
8.1 KiB
Python
Executable File
260 lines
8.1 KiB
Python
Executable File
#!/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'))
|
|
|
|
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}),
|
|
)
|
|
|
|
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()
|