unify unimport_func and unimport_gen, asyncify

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')

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)
+    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)

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()
 

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)

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()
+            await animate_swipe()
 
-    yield from paginate(render, len(mnemonic_pages))
+    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)

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()
 

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)

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)
+    return layout_get_public_key(*args, **kwargs)
 
 
-@unimport_func
-def dispatch_SignTx(mtype, mbuf):
-    from trezor.messages.SignTx import SignTx
-
-    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)

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()

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'))

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')

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)

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]

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

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')