diff --git a/src/apps/fido_u2f/layout_u2f.py b/src/apps/fido_u2f/layout_u2f.py index 27792bb4a..42888ffe4 100644 --- a/src/apps/fido_u2f/layout_u2f.py +++ b/src/apps/fido_u2f/layout_u2f.py @@ -1,6 +1,6 @@ from ubinascii import hexlify from trezor import ui, loop, res -from trezor.utils import unimport_gen +from trezor.utils import unimport from trezor.crypto import random from . import knownapps @@ -10,12 +10,14 @@ random.shuffle(ids) appid = ids[0] action = 'Register' -@unimport_gen + +@unimport def layout_u2f(): if appid in knownapps.knownapps: appname = knownapps.knownapps[appid] - appicon = res.load('apps/fido_u2f/res/u2f_%s.toif' % appname.lower().replace(' ', '_')) + appicon = res.load('apps/fido_u2f/res/u2f_%s.toif' % + appname.lower().replace(' ', '_')) else: appname = hexlify(appid[:4]) + '...' + hexlify(appid[-4:]) appicon = res.load('apps/fido_u2f/res/u2f_unknown.toif') diff --git a/src/apps/management/__init__.py b/src/apps/management/__init__.py index 5968c83c3..c8476f7d1 100644 --- a/src/apps/management/__init__.py +++ b/src/apps/management/__init__.py @@ -1,44 +1,35 @@ -from trezor.dispatcher import register -from trezor.utils import unimport_func +from trezor.wire import register_type, protobuf_handler +from trezor.utils import unimport +from trezor.messages.wire_types import \ + LoadDevice, ResetDevice, WipeDevice, RecoveryDevice -@unimport_func -def dispatch_LoadDevice(mtype, mbuf): - from trezor.messages.LoadDevice import LoadDevice - message = LoadDevice.loads(mbuf) +@unimport +def dispatch_LoadDevice(*args, **kwargs): from .layout_load_device import layout_load_device - return layout_load_device(message) + return layout_load_device(*args, **kwargs) -@unimport_func -def dispatch_ResetDevice(mtype, mbuf): - from trezor.messages.ResetDevice import ResetDevice - message = ResetDevice.loads(mbuf) +@unimport +def dispatch_ResetDevice(*args, **kwargs): from .layout_reset_device import layout_reset_device - return layout_reset_device(message) + return layout_reset_device(*args, **kwargs) -@unimport_func -def dispatch_WipeDevice(mtype, mbuf): - from trezor.messages.WipeDevice import WipeDevice - message = WipeDevice.loads(mbuf) +@unimport +def dispatch_WipeDevice(*args, **kwargs): from .layout_wipe_device import layout_wipe_device - return layout_wipe_device(message) + return layout_wipe_device(*args, **kwargs) -@unimport_func -def dispatch_RecoveryDevice(mtype, mbuf): - from trezor.messages.RecoveryDevice import RecoveryDevice - message = RecoveryDevice.loads(mbuf) +@unimport +def dispatch_RecoveryDevice(*args, **kwargs): from .layout_recovery_device import layout_recovery_device - return layout_recovery_device(message) + return layout_recovery_device(*args, **kwargs) + def boot(): - LoadDevice = 13 - register(LoadDevice, dispatch_LoadDevice) - ResetDevice = 14 - register(ResetDevice, dispatch_ResetDevice) - WipeDevice = 5 - register(WipeDevice, dispatch_WipeDevice) - RecoveryDevice = 45 - register(RecoveryDevice, dispatch_RecoveryDevice) \ No newline at end of file + register_type(LoadDevice, protobuf_handler, dispatch_LoadDevice) + register_type(ResetDevice, protobuf_handler, dispatch_ResetDevice) + register_type(WipeDevice, protobuf_handler, dispatch_WipeDevice) + register_type(RecoveryDevice, protobuf_handler, dispatch_RecoveryDevice) diff --git a/src/apps/management/layout_load_device.py b/src/apps/management/layout_load_device.py index 7ce7b22d0..fd8ba692f 100644 --- a/src/apps/management/layout_load_device.py +++ b/src/apps/management/layout_load_device.py @@ -1,15 +1,17 @@ from trezor import wire from trezor import ui -from trezor.utils import unimport_gen +from trezor.utils import unimport from trezor.workflows.confirm import confirm -@unimport_gen +@unimport def layout_load_device(message): ui.clear() - ui.display.text_center(120, 40, 'Really load device?', ui.BOLD, ui.WHITE, ui.BLACK) - ui.display.text_center(120, 100, 'Never do this, please.', ui.NORMAL, ui.WHITE, ui.BLACK) + ui.display.text_center(120, 40, 'Really load device?', + ui.BOLD, ui.WHITE, ui.BLACK) + ui.display.text_center( + 120, 100, 'Never do this, please.', ui.NORMAL, ui.WHITE, ui.BLACK) confirmed = yield from confirm() diff --git a/src/apps/management/layout_recovery_device.py b/src/apps/management/layout_recovery_device.py index d12d66ab1..19511e0c1 100644 --- a/src/apps/management/layout_recovery_device.py +++ b/src/apps/management/layout_recovery_device.py @@ -1,5 +1,5 @@ from trezor import ui, wire -from trezor.utils import unimport_gen +from trezor.utils import unimport def nth(n): @@ -10,13 +10,14 @@ def nth(n): return str(n) + sfx -@unimport_gen +@unimport def layout_recovery_device(message): msg = 'Please enter ' + nth(message.word_count) + ' word' ui.clear() - ui.display.text(10, 30, 'Recovering device', ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) + ui.display.text(10, 30, 'Recovering device', + ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) ui.display.text(10, 74, msg, ui.BOLD, ui.WHITE, ui.BLACK) ui.display.text(10, 104, 'of your mnemonic.', ui.BOLD, ui.WHITE, ui.BLACK) yield from wire.read(None) diff --git a/src/apps/management/layout_reset_device.py b/src/apps/management/layout_reset_device.py index 2941b6a49..a33ae6e99 100644 --- a/src/apps/management/layout_reset_device.py +++ b/src/apps/management/layout_reset_device.py @@ -1,15 +1,16 @@ from trezor import wire, ui +from trezor.messages.wire_types import EntropyAck from trezor.ui.button import Button, CONFIRM_BUTTON, CONFIRM_BUTTON_ACTIVE from trezor.ui.scroll import paginate, render_scrollbar, animate_swipe from trezor.crypto import hashlib, random, bip39 -from trezor.utils import unimport_gen, chunks +from trezor.utils import unimport, chunks -def generate_mnemonic(strength, display_random): +async def generate_mnemonic(strength, display_random, session_id): from trezor.messages.EntropyRequest import EntropyRequest - from trezor.messages.EntropyAck import EntropyAck - ack = yield from wire.call(EntropyRequest(), EntropyAck) + await wire.write_message(session_id, EntropyRequest()) + ack = await wire.read_message(session_id, EntropyAck) ctx = hashlib.sha256() ctx.update(random.bytes(32)) @@ -22,11 +23,11 @@ def generate_mnemonic(strength, display_random): return bip39.from_data(entropy) -def request_new_pin(): +async def request_new_pin(): from trezor.workflows.request_pin import request_pin - pin = yield from request_pin() - pin_again = yield from request_pin('Enter PIN again') + pin = await request_pin() + pin_again = await request_pin('Enter PIN again') if pin == pin_again: return pin @@ -34,49 +35,48 @@ def request_new_pin(): raise Exception() # TODO: wrong PIN should be handled in unified way -def show_mnemonic(mnemonic): +async def show_mnemonic(mnemonic): words_per_page = const(4) mnemonic_words = list(enumerate(mnemonic.split())) mnemonic_pages = list(chunks(mnemonic_words, words_per_page)) - def render(page, page_count): + async def render(page, page_count): - # Header + # render header & scrollbar ui.clear() - ui.display.text(10, 30, 'Write down your seed', ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) - + ui.display.text(10, 30, 'Write down your seed', + ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) render_scrollbar(page, page_count) - # Mnemonic page + # render mnemonic page for pi, (wi, word) in enumerate(mnemonic_pages[page]): top = pi * 30 + 74 pos = wi + 1 - ui.display.text_right(40, top, '%d.' % pos, ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) + ui.display.text_right(40, top, '%d.' % + pos, ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) ui.display.text(45, top, '%s' % word, ui.BOLD, ui.WHITE, ui.BLACK) if page + 1 == page_count: - # Finish button + # wait for the finish button finish = Button((0, 240 - 48, 240, 48), 'Finish', normal_style=CONFIRM_BUTTON, active_style=CONFIRM_BUTTON_ACTIVE) - yield from finish.wait() + await finish.wait() else: - # Swipe icon - yield from animate_swipe() - - yield from paginate(render, len(mnemonic_pages)) + await animate_swipe() + await paginate(render, len(mnemonic_pages)) -@unimport_gen -def layout_reset_device(m): +async def layout_reset_device(message, session_id): # TODO: Failure if not empty - mnemonic = yield from generate_mnemonic(m.strength, m.display_random) + mnemonic = await generate_mnemonic( + message.strength, message.display_random, session_id) # if m.pin_protection: # pin = yield from request_new_pin() # else: # pin = None - yield from show_mnemonic(mnemonic) + await show_mnemonic(mnemonic) diff --git a/src/apps/management/layout_wipe_device.py b/src/apps/management/layout_wipe_device.py index e8b9d6dc3..433b69d95 100644 --- a/src/apps/management/layout_wipe_device.py +++ b/src/apps/management/layout_wipe_device.py @@ -1,17 +1,19 @@ from trezor import wire from trezor import ui -from trezor.utils import unimport_gen +from trezor.utils import unimport from trezor.workflows.confirm import confirm -@unimport_gen +@unimport def layout_wipe_device(message): ui.clear() ui.display.text(10, 30, 'Wiping device', ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) - ui.display.text(10, 74, 'Do you really want to', ui.BOLD, ui.WHITE, ui.BLACK) + ui.display.text(10, 74, 'Do you really want to', + ui.BOLD, ui.WHITE, ui.BLACK) ui.display.text(10, 104, 'wipe the device?', ui.BOLD, ui.WHITE, ui.BLACK) - ui.display.text(10, 164, 'All data will be lost.', ui.NORMAL, ui.WHITE, ui.BLACK) + ui.display.text(10, 164, 'All data will be lost.', + ui.NORMAL, ui.WHITE, ui.BLACK) confirmed = yield from confirm() diff --git a/src/apps/playground/__init__.py b/src/apps/playground/__init__.py index 69d78b6fb..b10df50d3 100644 --- a/src/apps/playground/__init__.py +++ b/src/apps/playground/__init__.py @@ -1,11 +1,11 @@ -from trezor import loop from trezor import ui -from trezor import msg -from trezor.utils import unimport_gen from trezor import res +from trezor import msg +from trezor import loop +from trezor.utils import unimport -@unimport_gen +@unimport def layout_tap_to_confirm(address, amount, currency): # ui.display.bar(0, 0, 240, 40, ui.GREEN) diff --git a/src/apps/wallet/__init__.py b/src/apps/wallet/__init__.py index bf43c9cec..1e9c18003 100644 --- a/src/apps/wallet/__init__.py +++ b/src/apps/wallet/__init__.py @@ -1,42 +1,28 @@ -from trezor.dispatcher import register -from trezor.utils import unimport_func +from trezor.wire import register_type, protobuf_handler +from trezor.utils import unimport +from trezor.messages.wire_types import \ + GetPublicKey, SignTx, SignMessage -@unimport_func -def dispatch_GetPublicKey(mtype, mbuf): - from trezor.messages.GetPublicKey import GetPublicKey - - message = GetPublicKey.loads(mbuf) - +@unimport +def dispatch_GetPublicKey(*args, **kwargs): from .layout_get_public_key import layout_get_public_key - return layout_get_public_key(message) - - -@unimport_func -def dispatch_SignTx(mtype, mbuf): - from trezor.messages.SignTx import SignTx + return layout_get_public_key(*args, **kwargs) - message = SignTx.loads(mbuf) +@unimport +def dispatch_SignTx(*args, **kwargs): from .layout_sign_tx import layout_sign_tx - return layout_sign_tx(message) + return layout_sign_tx(*args, **kwargs) -@unimport_func -def dispatch_SignMessage(mtype, mbuf): - from trezor.messages.SignMessage import SignMessage - - message = SignMessage.loads(mbuf) - +@unimport +def dispatch_SignMessage(*args, **kwargs): from .layout_sign_message import layout_sign_message - return layout_sign_message(message) + return layout_sign_message(*args, **kwargs) def boot(): - GetPublicKey = 11 - register(GetPublicKey, dispatch_GetPublicKey) - SignTx = 15 - register(SignTx, dispatch_SignTx) - SignMessage = 38 - register(SignMessage, dispatch_SignMessage) - + register_type(GetPublicKey, protobuf_handler, dispatch_GetPublicKey) + register_type(SignTx, protobuf_handler, dispatch_SignTx) + register_type(SignMessage, protobuf_handler, dispatch_SignMessage) diff --git a/src/apps/wallet/layout_get_public_key.py b/src/apps/wallet/layout_get_public_key.py index 949534635..11d0eb12b 100644 --- a/src/apps/wallet/layout_get_public_key.py +++ b/src/apps/wallet/layout_get_public_key.py @@ -1,9 +1,9 @@ from trezor import wire, ui -from trezor.utils import unimport_gen +from trezor.utils import unimport from trezor.workflows.request_pin import request_pin -@unimport_gen +@unimport def layout_get_public_key(message): ui.clear() diff --git a/src/apps/wallet/layout_sign_message.py b/src/apps/wallet/layout_sign_message.py index 1350963ea..db6dea55d 100644 --- a/src/apps/wallet/layout_sign_message.py +++ b/src/apps/wallet/layout_sign_message.py @@ -1,14 +1,16 @@ from trezor import wire, ui -from trezor.utils import unimport_gen +from trezor.utils import unimport from trezor.workflows.confirm import confirm -@unimport_gen + +@unimport def layout_sign_message(message): ui.clear() - ui.display.text(10, 30, 'Signing message', ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) + ui.display.text(10, 30, 'Signing message', + ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) ui.display.text(10, 60, message.message, ui.MONO, ui.WHITE, ui.BLACK) confirmed = yield from confirm(confirm='Sign') - + if confirmed: from trezor.messages.Success import Success yield from wire.write(Success(message='Signed')) diff --git a/src/apps/wallet/layout_sign_tx.py b/src/apps/wallet/layout_sign_tx.py index b5b92e799..9ded35d9b 100644 --- a/src/apps/wallet/layout_sign_tx.py +++ b/src/apps/wallet/layout_sign_tx.py @@ -1,8 +1,8 @@ from trezor import wire, ui -from trezor.utils import unimport_gen +from trezor.utils import unimport -@unimport_gen +@unimport def layout_sign_tx(message): ui.clear() print('sending') diff --git a/src/trezor/ui/scroll.py b/src/trezor/ui/scroll.py index f1be92923..c3fdb52bf 100644 --- a/src/trezor/ui/scroll.py +++ b/src/trezor/ui/scroll.py @@ -2,27 +2,37 @@ from .swipe import Swipe, SWIPE_UP, SWIPE_DOWN from trezor import loop, ui -def change_page(page, page_count): +async def change_page(page, page_count): while True: - s = yield from Swipe() + s = await Swipe() if s == SWIPE_UP and page < page_count - 1: - return page + 1 # Scroll down + return page + 1 # scroll down elif s == SWIPE_DOWN and page > 0: - return page - 1 # Scroll up + return page - 1 # scroll up -def paginate(render_page, page_count, page=0): +async def paginate(render_page, page_count, page=0): while True: changer = change_page(page, page_count) renderer = render_page(page, page_count) waiter = loop.Wait([changer, renderer]) - result = yield waiter + result = await waiter if changer in waiter.finished: page = result else: return result +async def animate_swipe(): + await ui.animate_pulse(render_swipe_icon, ui.WHITE, ui.GREY, speed=300000, delay=200000) + + +def render_swipe_icon(fg): + ui.display.bar(102, 214, 36, 4, fg, ui.BLACK, 2) + ui.display.bar(106, 222, 28, 4, fg, ui.BLACK, 2) + ui.display.bar(110, 230, 20, 4, fg, ui.BLACK, 2) + + def render_scrollbar(page, page_count): screen_height = const(220) size = const(8) @@ -39,11 +49,3 @@ def render_scrollbar(page, page_count): ui.display.bar(x, y + i * padding, size, size, ui.GREY, ui.BLACK, 4) ui.display.bar(x, y + page * padding, size, size, ui.WHITE, ui.BLACK, 4) - - -def animate_swipe(): - def render(fg): - ui.display.bar(102, 214, 36, 4, fg, ui.BLACK, 2) - ui.display.bar(106, 222, 28, 4, fg, ui.BLACK, 2) - ui.display.bar(110, 230, 20, 4, fg, ui.BLACK, 2) - yield from ui.animate_pulse(render, ui.WHITE, ui.GREY, speed=300000, delay=200000) diff --git a/src/trezor/utils.py b/src/trezor/utils.py index 804e50eee..b07d9446e 100644 --- a/src/trezor/utils.py +++ b/src/trezor/utils.py @@ -1,8 +1,10 @@ import sys type_gen = type((lambda: (yield))()) +type_genfunc = type((lambda: (yield))) -def unimport_func(func): + +def _unimport_func(func): def inner(*args, **kwargs): mods = set(sys.modules) try: @@ -13,17 +15,34 @@ def unimport_func(func): return ret return inner -def unimport_gen(gen): + +def _unimport_genfunc(genfunc): def inner(*args, **kwargs): mods = set(sys.modules) try: - ret = yield from gen(*args, **kwargs) + ret = yield from genfunc(*args, **kwargs) finally: for to_remove in set(sys.modules) - mods: del sys.modules[to_remove] return ret return inner + +def unimport(func): + if isinstance(func, type_genfunc): + return _unimport_genfunc(func) + else: + return _unimport_func(func) + + +def coroutine(func): + def inner(*args, **kwargs): + gen = func(*args, **kwargs) + gen.send(None) + return gen + return inner + + def chunks(l, n): for i in range(0, len(l), n): yield l[i:i + n] diff --git a/src/trezor/workflows/confirm.py b/src/trezor/workflows/confirm.py index 8cdf8b832..89545341d 100644 --- a/src/trezor/workflows/confirm.py +++ b/src/trezor/workflows/confirm.py @@ -1,9 +1,9 @@ from trezor import wire -from trezor.utils import unimport_gen +from trezor.utils import unimport -@unimport_gen -def confirm(content=None, code=None, **kwargs): +@unimport +async def confirm(content=None, code=None, **kwargs): from trezor.ui.confirm import ConfirmDialog, CONFIRMED from trezor.messages.ButtonRequest import ButtonRequest from trezor.messages.ButtonRequestType import Other @@ -14,7 +14,7 @@ def confirm(content=None, code=None, **kwargs): if code is None: code = Other - ack = yield from wire.call(ButtonRequest(code=code), ButtonAck) - res = yield from dialog.wait() + ack = await wire.call(ButtonRequest(code=code), ButtonAck) + res = await dialog.wait() return res == CONFIRMED diff --git a/src/trezor/workflows/request_pin.py b/src/trezor/workflows/request_pin.py index 0363056c0..1dd42d219 100644 --- a/src/trezor/workflows/request_pin.py +++ b/src/trezor/workflows/request_pin.py @@ -1,9 +1,9 @@ from trezor import ui from trezor import wire from trezor import config -from trezor.utils import unimport_gen +from trezor.utils import unimport -MGMT_APP = const(1) +MANAGEMENT_APP = const(1) PASSPHRASE_PROTECT = (1) # 0 | 1 PIN_PROTECT = const(2) # 0 | 1 @@ -38,13 +38,12 @@ def change_pin(): pass -@unimport_gen def protect_with_pin(): from trezor.messages.Failure import Failure from trezor.messages.FailureType import PinInvalid from trezor.messages.FailureType import ActionCancelled - pin_protect = config.get(MGMT_APP, PIN_PROTECT) + pin_protect = config.get(MANAGEMENT_APP, PIN_PROTECT) if not pin_protect: return @@ -53,7 +52,7 @@ def protect_with_pin(): yield from wire.write(Failure(code=ActionCancelled, message='Cancelled')) raise Exception('Cancelled') - stored_pin = config.get(MGMT_APP, PIN) + stored_pin = config.get(MANAGEMENT_APP, PIN) if stored_pin != entered_pin: yield from wire.write(Failure(code=PinInvalid, message='PIN invalid')) raise Exception('PIN invalid')