import os import time import binascii import hashlib import tools import messages_pb2 as proto import types_pb2 as types from api_blockchain import BlockchainApi # monkeypatching: text formatting of protobuf messages tools.monkeypatch_google_protobuf_text_format() def show_message(message): print "MESSAGE FROM DEVICE:", message def show_input(input_text, message=None): if message: print "QUESTION FROM DEVICE:", message return raw_input(input_text) def pin_func(input_text, message=None): return show_input(input_text, message) def passphrase_func(input_text): return show_input(input_text) def word_func(): return raw_input("Enter one word of mnemonic: ") class CallException(Exception): def __init__(self, code, message): super(CallException, self).__init__() self.args = [code, message] class PinException(CallException): pass PRIME_DERIVATION_FLAG = 0x80000000 class TrezorClient(object): def __init__(self, transport, debuglink=None, message_func=show_message, input_func=show_input, pin_func=pin_func, passphrase_func=passphrase_func, word_func=word_func, blockchain_api=None, debug=False): self.transport = transport self.debuglink = debuglink self.message_func = message_func self.input_func = input_func self.pin_func = pin_func self.passphrase_func = passphrase_func self.word_func = word_func self.debug = debug if blockchain_api: self.blockchain = blockchain_api else: self.blockchain = BlockchainApi() self.setup_debuglink() self.init_device() def _get_local_entropy(self): return os.urandom(32) def _convert_prime(self, n): # Convert minus signs to uint32 with flag return [ int(abs(x) | PRIME_DERIVATION_FLAG) if x < 0 else x for x in n ] def expand_path(self, n): # Convert string of bip32 path to list of uint32 integers with prime flags # 0/-1/1' -> [0, 0x80000001, 0x80000001] if not n: return [] n = n.split('/') path = [] for x in n: prime = False if x.endswith("'"): x = x.replace('\'', '') prime = True if x.startswith('-'): prime = True x = abs(int(x)) if prime: x |= PRIME_DERIVATION_FLAG path.append(x) return path def init_device(self): self.features = self.call(proto.Initialize(), proto.Features) def close(self): self.transport.close() if self.debuglink: self.debuglink.transport.close() def get_public_node(self, n): return self.call(proto.GetPublicKey(address_n=n), proto.PublicKey).node def get_address(self, coin_name, n): n = self._convert_prime(n) return self.call(proto.GetAddress(address_n=n, coin_name=coin_name), proto.Address).address def get_entropy(self, size): return self.call(proto.GetEntropy(size=size), proto.Entropy).entropy def ping(self, msg, button_protection=False, pin_protection=False, passphrase_protection=False): msg = proto.Ping(message=msg, button_protection=button_protection, pin_protection=pin_protection, passphrase_protection=passphrase_protection) return self.call(msg, proto.Success).message def get_device_id(self): return self.features.device_id def apply_settings(self, label=None, language=None): settings = proto.ApplySettings() if label != None: settings.label = label if language: settings.language = language out = self.call(settings, proto.Success).message self.init_device() # Reload Features return out def change_pin(self, remove=False): ret = self.call(proto.ChangePin(remove=remove)) self.init_device() # Re-read features return ret def _pprint(self, msg): ser = msg.SerializeToString() return "<%s> (%d bytes):\n%s" % (msg.__class__.__name__, len(ser), msg) def setup_debuglink(self, button=None, pin_correct=False): self.debug_button = button self.debug_pin = pin_correct def _get_buttonrequest_value(self, code): return [ k for k, v in types.ButtonRequestType.items() if v == code][0] def call(self, msg, expected=None, expected_buttonrequests=None): # TODO split this into normal and debug mode if self.debug: print '----------------------' print "Sending", self._pprint(msg) try: self.transport.session_begin() self.transport.write(msg) resp = self.transport.read_blocking() if isinstance(resp, proto.ButtonRequest): if expected_buttonrequests != None: try: exp = expected_buttonrequests.pop(0) if resp.code != exp: raise CallException(types.Failure_Other, "Expected %s, got %s" % \ (self._get_buttonrequest_value(exp), self._get_buttonrequest_value(resp.code))) except IndexError: raise CallException(types.Failure_Other, "Got %s, but no ButtonRequest has been expected" % \ self._get_buttonrequest_value(resp.code)) print "ButtonRequest code:", self._get_buttonrequest_value(resp.code) if self.debuglink and self.debug_button: print "Pressing button", self.debug_button self.debuglink.press_button(self.debug_button) return self.call(proto.ButtonAck(), expected_buttonrequests=expected_buttonrequests) if isinstance(resp, proto.PinMatrixRequest): if self.debuglink: if self.debug_pin == 1: pin = self.debuglink.read_pin_encoded() msg2 = proto.PinMatrixAck(pin=pin) elif self.debug_pin == -1: msg2 = proto.Cancel() else: msg2 = proto.PinMatrixAck(pin='444444222222') else: pin = self.pin_func("PIN required: ", resp.message) msg2 = proto.PinMatrixAck(pin=pin) return self.call(msg2, expected=expected, expected_buttonrequests=expected_buttonrequests) if isinstance(resp, proto.PassphraseRequest): passphrase = self.passphrase_func("Passphrase required: ") msg2 = proto.PassphraseAck(passphrase=passphrase) return self.call(msg2, expected=expected, expected_buttonrequests=expected_buttonrequests) finally: self.transport.session_end() if isinstance(resp, proto.Failure): self.message_func(resp.message) if resp.code in (types.Failure_PinInvalid, types.Failure_PinCancelled, types.Failure_PinExpected): raise PinException(resp.code, resp.message) raise CallException(resp.code, resp.message) if self.debug: print "Received", self._pprint(resp) if expected and not isinstance(resp, expected): raise CallException("Expected %s message, got %s message" % (expected.DESCRIPTOR.name, resp.DESCRIPTOR.name)) if expected_buttonrequests != None and len(expected_buttonrequests): raise CallException(types.Failure_Other, "Following ButtonRequests were not in use: %s" % \ [ self._get_buttonrequest_value(x) for x in expected_buttonrequests]) return resp def sign_message(self, n, message): n = self._convert_prime(n) return self.call(proto.SignMessage(address_n=n, message=message)) def verify_message(self, address, signature, message): try: resp = self.call(proto.VerifyMessage(address=address, signature=signature, message=message)) if isinstance(resp, proto.Success): return True except CallException: pass return False def estimate_tx_size(self, coin_name, inputs, outputs): msg = proto.EstimateTxSize() msg.coin_name = coin_name msg.inputs_count = len(inputs) msg.outputs_count = len(outputs) res = self.call(msg) return res.tx_size def _prepare_simple_sign_tx(self, coin_name, inputs, outputs): msg = proto.SimpleSignTx() msg.coin_name = coin_name msg.inputs.extend(inputs) msg.outputs.extend(outputs) known_hashes = [] for inp in inputs: if inp.prev_hash in known_hashes: continue tx = msg.transactions.add() tx.CopyFrom(self.blockchain.get_tx(binascii.hexlify(inp.prev_hash))) known_hashes.append(inp.prev_hash) return msg def simple_sign_tx(self, coin_name, inputs, outputs): msg = self._prepare_simple_sign_tx(coin_name, inputs, outputs) return self.call(msg) def sign_tx(self, coin_name, inputs, outputs): # Temporary solution, until streaming is implemented in the firmware return self.simple_sign_tx(coin_name, inputs, outputs) def _sign_tx(self, coin_name, inputs, outputs): ''' inputs: list of TxInput outputs: list of TxOutput proto.TxInput(index=0, address_n=0, amount=0, prev_hash='', prev_index=0, #script_sig= ) proto.TxOutput(index=0, address='1Bitkey', #address_n=[], amount=100000000, script_type=proto.PAYTOADDRESS, #script_args= ) ''' start = time.time() try: self.transport.session_begin() # Prepare and send initial message tx = proto.SignTx() tx.inputs_count = len(inputs) tx.outputs_count = len(outputs) res = self.call(tx) # Prepare structure for signatures signatures = [None]*len(inputs) serialized_tx = '' 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_tx: print "!!! RECEIVED PART OF SERIALIED TX (%d BYTES)" % len(res.serialized_tx) serialized_tx += res.serialized_tx if res.signed_index >= 0 and res.signature: print "!!! SIGNED INPUT", res.signed_index signatures[res.signed_index] = res.signature if res.request_index < 0: # Device didn't ask for more information, finish workflow break # Device asked for one more information, let's process it. if res.request_type == types.TXOUTPUT: res = self.call(outputs[res.request_index]) continue elif res.request_type == types.TXINPUT: print "REQUESTING", res.request_index res = self.call(inputs[res.request_index]) continue finally: self.transport.session_end() print "SIGNED IN %.03f SECONDS, CALLED %d MESSAGES, %d BYTES" % \ (time.time() - start, counter, len(serialized_tx)) return (signatures, serialized_tx) def wipe_device(self): ret = self.call(proto.WipeDevice()) self.init_device() return ret def recovery_device(self, word_count, passphrase_protection, pin_protection, label, language): if word_count not in (12, 18, 24): raise Exception("Invalid word count. Use 12/18/24") res = self.call(proto.RecoveryDevice(word_count=int(word_count), passphrase_protection=bool(passphrase_protection), pin_protection=bool(pin_protection), label=label, language=language, enforce_wordlist=True)) while isinstance(res, proto.WordRequest): word = self.word_func() res = self.call(proto.WordAck(word=word)) if not isinstance(res, proto.Success): raise Exception("Recovery device failed") self.init_device() return True def reset_device(self, display_random, strength, passphrase_protection, pin_protection, label, language): if self.features.initialized: raise Exception("Device is initialized already. Call wipe_device() and try again.") # Begin with device reset workflow msg = proto.ResetDevice(display_random=display_random, strength=strength, language=language, passphrase_protection=bool(passphrase_protection), pin_protection=bool(pin_protection), label=label) resp = self.call(msg) if not isinstance(resp, proto.EntropyRequest): raise Exception("Invalid response, expected EntropyRequest") external_entropy = self._get_local_entropy() print "Computer generated entropy:", binascii.hexlify(external_entropy) resp = self.call(proto.EntropyAck(entropy=external_entropy)) return isinstance(resp, proto.Success) def load_device_by_mnemonic(self, mnemonic, pin, passphrase_protection, label, language): if self.features.initialized: raise Exception("Device is initialized already. Call wipe_device() and try again.") resp = self.call(proto.LoadDevice(mnemonic=mnemonic, pin=pin, passphrase_protection=passphrase_protection, language=language, label=label)) self.init_device() return isinstance(resp, proto.Success) def load_device_by_xprv(self, xprv, pin, passphrase_protection, label): if self.features.initialized: raise Exception("Device is initialized already. Call wipe_device() and try again.") if xprv[0:4] not in ('xprv', 'tprv'): raise Exception("Unknown type of xprv") if len(xprv) < 100 and len(xprv) > 112: raise Exception("Invalid length of xprv") node = types.HDNodeType() data = tools.b58decode(xprv, None).encode('hex') if data[90:92] != '00': raise Exception("Contain invalid private key") checksum = hashlib.sha256(hashlib.sha256(binascii.unhexlify(data[:156])).digest()).hexdigest()[:8] if checksum != data[156:]: raise Exception("Checksum doesn't match") # version 0488ade4 # depth 00 # fingerprint 00000000 # child_num 00000000 # chaincode 873dff81c02f525623fd1fe5167eac3a55a049de3d314bb42ee227ffed37d508 # privkey 00e8f32e723decf4051aefac8e2c93c9c5b214313817cdb01a1494b917c8436b35 # checksum e77e9d71 node.version = int(data[0:8], 16) node.depth = int(data[8:10], 16) node.fingerprint = int(data[10:18], 16) node.child_num = int(data[18:26], 16) node.chain_code = data[26:90].decode('hex') node.private_key = data[92:156].decode('hex') # skip 0x00 indicating privkey resp = self.call(proto.LoadDevice(node=node, pin=pin, passphrase_protection=passphrase_protection, language='english', label=label)) self.init_device() return isinstance(resp, proto.Success) def firmware_update(self, fp): if self.features.bootloader_mode == False: raise Exception("Device must be in bootloader mode") resp = self.call(proto.FirmwareErase()) if isinstance(resp, proto.Failure) and resp.code == types.Failure_FirmwareError: return False resp = self.call(proto.FirmwareUpload(payload=fp.read())) if isinstance(resp, proto.Success): return True elif isinstance(resp, proto.Failure) and resp.code == types.Failure_FirmwareError: return False raise Exception("Unexpected result " % resp) # class TrezorDebugClient(TrezorClient):