diff --git a/src/apps/homescreen/layout_homescreen.py b/src/apps/homescreen/layout_homescreen.py index 6ffa0819..aefb4978 100644 --- a/src/apps/homescreen/layout_homescreen.py +++ b/src/apps/homescreen/layout_homescreen.py @@ -39,7 +39,7 @@ def layout_homescreen(initialize_msg=None): features.passphrase_cached = False features.passphrase_protection = False features.vendor = 'bitcointrezor.com' - wire.write(features) + yield from wire.write(features) yield loop.Wait([dispatcher.dispatch(), swipe_to_rotate(), animate_logo()]) diff --git a/src/apps/management/__init__.py b/src/apps/management/__init__.py index 29479f35..d8cbe763 100644 --- a/src/apps/management/__init__.py +++ b/src/apps/management/__init__.py @@ -12,6 +12,18 @@ def dispatch_LoadDevice(mtype, mbuf): return layout_load_device(message) +@unimport_func +def dispatch_WipeDevice(mtype, mbuf): + from trezor.messages.WipeDevice import WipeDevice + + message = WipeDevice.loads(mbuf) + + from .layout_wipe_device import layout_wipe_device + return layout_wipe_device(message) + + def boot(): LoadDevice = 13 register(LoadDevice, dispatch_LoadDevice) + WipeDevice = 5 + register(WipeDevice, dispatch_WipeDevice) diff --git a/src/apps/management/layout_load_device.py b/src/apps/management/layout_load_device.py index e61afca3..7ce7b22d 100644 --- a/src/apps/management/layout_load_device.py +++ b/src/apps/management/layout_load_device.py @@ -1,22 +1,7 @@ from trezor import wire from trezor import ui from trezor.utils import unimport_gen - - -@unimport_gen -def confirm(): - from trezor.ui.confirm import ConfirmDialog, CONFIRMED - from trezor.messages.ButtonRequest import ButtonRequest - from trezor.messages.ButtonRequestType import Other - from trezor.messages.ButtonAck import ButtonAck - - dialog = ConfirmDialog() - dialog.render() - - ack = yield from wire.call(ButtonRequest(code=Other), ButtonAck) - res = yield from dialog.wait() - - return res == CONFIRMED +from trezor.workflows.confirm import confirm @unimport_gen @@ -30,8 +15,8 @@ def layout_load_device(message): if confirmed: from trezor.messages.Success import Success - wire.write(Success(message='Loaded')) + yield from wire.write(Success(message='Loaded')) else: from trezor.messages.Failure import Failure from trezor.messages.FailureType import ActionCancelled - wire.write(Failure(message='Cancelled', code=ActionCancelled)) + yield from wire.write(Failure(message='Cancelled', code=ActionCancelled)) diff --git a/src/apps/management/layout_wipe_device.py b/src/apps/management/layout_wipe_device.py new file mode 100644 index 00000000..1aa658cc --- /dev/null +++ b/src/apps/management/layout_wipe_device.py @@ -0,0 +1,22 @@ +from trezor import wire +from trezor import ui +from trezor.utils import unimport_gen +from trezor.workflows.confirm import confirm + + +@unimport_gen +def layout_wipe_device(message): + + ui.clear() + ui.display.text_center(120, 40, 'Really wipe device?', ui.BOLD, ui.WHITE, ui.BLACK) + ui.display.text_center(120, 100, 'You might regret it!', ui.NORMAL, ui.WHITE, ui.BLACK) + + confirmed = yield from confirm() + + if confirmed: + from trezor.messages.Success import Success + yield from wire.write(Success(message='Wiped')) + else: + from trezor.messages.Failure import Failure + from trezor.messages.FailureType import ActionCancelled + yield from wire.write(Failure(message='Cancelled', code=ActionCancelled)) diff --git a/src/apps/wallet/layout_get_public_key.py b/src/apps/wallet/layout_get_public_key.py index 4071562b..47d0a005 100644 --- a/src/apps/wallet/layout_get_public_key.py +++ b/src/apps/wallet/layout_get_public_key.py @@ -1,23 +1,6 @@ from trezor import wire from trezor.utils import unimport_gen - - -@unimport_gen -def request_pin(): - from trezor.ui.pin import PinMatrix - from trezor.ui.confirm import ConfirmDialog, CONFIRMED - from trezor.messages.ButtonRequest import ButtonRequest - from trezor.messages.ButtonRequestType import ProtectCall - from trezor.messages.ButtonAck import ButtonAck - - matrix = PinMatrix() - dialog = ConfirmDialog(matrix) - dialog.render() - - ack = yield from wire.call(ButtonRequest(code=ProtectCall), ButtonAck) - res = yield from dialog.wait() - - return matrix.pin if res == CONFIRMED else None +from trezor.workflows.request_pin import request_pin @unimport_gen @@ -35,9 +18,9 @@ def layout_get_public_key(message): pubkey.node.fingerprint = 0 pubkey.node.chain_code = 'deadbeef' pubkey.node.public_key = 'deadbeef' - wire.write(pubkey) + yield from wire.write(pubkey) else: from trezor.messages.Failure import Failure from trezor.messages.FailureType import ActionCancelled - wire.write(Failure(message='Cancelled', code=ActionCancelled)) + yield from wire.write(Failure(message='Cancelled', code=ActionCancelled)) diff --git a/src/tests/test_wire.py b/src/tests/test_wire.py index 964fda47..dea3dda4 100644 --- a/src/tests/test_wire.py +++ b/src/tests/test_wire.py @@ -75,7 +75,15 @@ class TestWire(unittest.TestCase): content = bytes([x for x in range(0, 256)]) message = b'##\xab\xcd\x00\x00\x01\00' + content reports = [b'\x3f' + ch + '\x00' * (63 - len(ch)) for ch in chunks(message, 63)] - write_wire_msg(int('0xabcd'), content) + + writer = write_wire_msg(int('0xabcd'), content) + res = 1 # Something not None + try: + while True: + writer.send(None) + except StopIteration as e: + res = e.value + self.assertEqual(res, None) self.assertEqual(sent_reps, reports) diff --git a/src/trezor/msg.py b/src/trezor/msg.py index 405c1d61..7b5161bc 100644 --- a/src/trezor/msg.py +++ b/src/trezor/msg.py @@ -1,6 +1,3 @@ -import ustruct - -from trezor import loop from TrezorMsg import Msg _msg = Msg() diff --git a/src/trezor/wire.py b/src/trezor/wire.py index 97229c16..78b3c61f 100644 --- a/src/trezor/wire.py +++ b/src/trezor/wire.py @@ -18,6 +18,7 @@ def read_report(): def write_report(rep): size = msg.send(IFACE, rep) assert size == REPORT_LEN, 'HID write failed' + yield # write_report is a generator for the sake of consistency def read_wire_msg(): @@ -64,7 +65,7 @@ def write_wire_msg(mtype, mbuf): while i < len(data): data[i] = 0 i += 1 - write_report(rep) + yield from write_report(rep) mbuf = mbuf[n:] data = rep[1:] @@ -81,10 +82,10 @@ def read(*types): def write(m): mbuf = m.dumps() mtype = m.message_type.wire_type - write_wire_msg(mtype, mbuf) + yield from write_wire_msg(mtype, mbuf) def call(req, *types): - write(req) + yield from write(req) res = yield from read(*types) return res diff --git a/src/trezor/workflows/__init__.py b/src/trezor/workflows/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/trezor/workflows/confirm.py b/src/trezor/workflows/confirm.py new file mode 100644 index 00000000..2c1c3f0a --- /dev/null +++ b/src/trezor/workflows/confirm.py @@ -0,0 +1,21 @@ +from trezor import wire +from trezor import ui +from trezor.utils import unimport_gen + + +@unimport_gen +def confirm(content=None, code=None): + from trezor.ui.confirm import ConfirmDialog, CONFIRMED + from trezor.messages.ButtonRequest import ButtonRequest + from trezor.messages.ButtonRequestType import Other + from trezor.messages.ButtonAck import ButtonAck + + dialog = ConfirmDialog(content) + dialog.render() + + if code is None: + code = Other + ack = yield from wire.call(ButtonRequest(code=code), ButtonAck) + res = yield from dialog.wait() + + return res == CONFIRMED diff --git a/src/trezor/workflows/request_pin.py b/src/trezor/workflows/request_pin.py new file mode 100644 index 00000000..e163e72f --- /dev/null +++ b/src/trezor/workflows/request_pin.py @@ -0,0 +1,20 @@ +from trezor import wire +from trezor.utils import unimport_gen + + +@unimport_gen +def request_pin(): + from trezor.ui.pin import PinMatrix + from trezor.ui.confirm import ConfirmDialog, CONFIRMED + from trezor.messages.ButtonRequest import ButtonRequest + from trezor.messages.ButtonRequestType import ProtectCall + from trezor.messages.ButtonAck import ButtonAck + + matrix = PinMatrix() + dialog = ConfirmDialog(matrix) + dialog.render() + + ack = yield from wire.call(ButtonRequest(code=ProtectCall), ButtonAck) + res = yield from dialog.wait() + + return matrix.pin if res == CONFIRMED else None