From e9c7a67c734298251775d40054af1f4c7f58d2a0 Mon Sep 17 00:00:00 2001 From: Jan Pochyla Date: Mon, 26 Sep 2016 16:11:38 +0200 Subject: [PATCH] make imports more local, remove trezor.workflows, minor tweaks --- src/apps/homescreen/layout_homescreen.py | 9 ++- .../workflows => apps/management}/confirm.py | 2 +- src/apps/management/layout_load_device.py | 4 +- src/apps/management/layout_reset_device.py | 36 ++++++------ src/apps/management/layout_wipe_device.py | 8 +-- .../management}/request_pin.py | 2 +- src/apps/wallet/layout_sign_message.py | 5 +- src/trezor/wire/__init__.py | 56 +++++++------------ src/trezor/workflows/__init__.py | 0 9 files changed, 54 insertions(+), 68 deletions(-) rename src/{trezor/workflows => apps/management}/confirm.py (94%) rename src/{trezor/workflows => apps/management}/request_pin.py (96%) delete mode 100644 src/trezor/workflows/__init__.py diff --git a/src/apps/homescreen/layout_homescreen.py b/src/apps/homescreen/layout_homescreen.py index 36f735447..4e9cdf238 100644 --- a/src/apps/homescreen/layout_homescreen.py +++ b/src/apps/homescreen/layout_homescreen.py @@ -1,8 +1,10 @@ from trezor import ui, loop, res -from trezor.ui.swipe import Swipe +from trezor.utils import unimport async def swipe_to_rotate(): + from trezor.ui.swipe import Swipe + while True: degrees = await Swipe(absolute=True) ui.display.orientation(degrees) @@ -10,9 +12,12 @@ async def swipe_to_rotate(): async def animate_logo(): icon = res.load('apps/homescreen/res/trezor.toig') - async for fg in ui.pulse_animation(ui.WHITE, ui.GREY, speed=400000): + + def render(fg): ui.display.icon(0, 0, icon, fg, ui.BLACK) + await ui.animate_pulse(render, ui.WHITE, ui.GREY, speed=400000) +@unimport async def layout_homescreen(): await loop.Wait([swipe_to_rotate(), animate_logo()]) diff --git a/src/trezor/workflows/confirm.py b/src/apps/management/confirm.py similarity index 94% rename from src/trezor/workflows/confirm.py rename to src/apps/management/confirm.py index 1103dbf9b..62da9da5e 100644 --- a/src/trezor/workflows/confirm.py +++ b/src/apps/management/confirm.py @@ -19,7 +19,7 @@ async def confirm(session_id, content=None, code=None, **kwargs): @unimport -async def protect_with_confirm(*args, **kwargs): +async def require_confirm(*args, **kwargs): from trezor.messages.FailureType import ActionCancelled confirmed = await confirm(*args, **kwargs) diff --git a/src/apps/management/layout_load_device.py b/src/apps/management/layout_load_device.py index 500f040e0..b96dc20e7 100644 --- a/src/apps/management/layout_load_device.py +++ b/src/apps/management/layout_load_device.py @@ -4,8 +4,8 @@ from trezor.utils import unimport @unimport async def layout_load_device(session_id, message): - from trezor.workflows.confirm import protect_with_confirm from trezor.messages.Success import Success + from .confirm import require_confirm ui.clear() ui.display.text_center( @@ -13,7 +13,7 @@ async def layout_load_device(session_id, message): ui.display.text_center( 120, 100, 'Never do this, please.', ui.NORMAL, ui.WHITE, ui.BLACK) - await protect_with_confirm(session_id) + await require_confirm(session_id) # TODO diff --git a/src/apps/management/layout_reset_device.py b/src/apps/management/layout_reset_device.py index 2868429d2..1d5bd4d20 100644 --- a/src/apps/management/layout_reset_device.py +++ b/src/apps/management/layout_reset_device.py @@ -1,31 +1,24 @@ from trezor import wire, ui -from trezor.workflows.request_pin import request_pin_repeatedly -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, chunks -from .storage import get_storage, set_storage - @unimport async def layout_reset_device(message, session_id): from trezor.messages.Success import Success from trezor.messages.Storage import Storage from trezor.messages.FailureType import UnexpectedMessage + from .storage import get_storage, set_storage + from .request_pin import request_pin_twice if get_storage(session_id): - raise wire.FailureError( - UnexpectedMessage, 'Device is already initialized') + raise wire.FailureError(UnexpectedMessage, 'Already initialized') mnemonic = await generate_mnemonic( message.strength, message.display_random, session_id) - await show_mnemonic(mnemonic) if message.pin_protection: - pin = await request_pin_repeatedly(session_id) + pin = await request_pin_twice(session_id) else: pin = '' @@ -40,8 +33,10 @@ async def layout_reset_device(message, session_id): @unimport async def generate_mnemonic(strength, display_random, session_id): + from trezor.crypto import hashlib, random, bip39 from trezor.messages.EntropyRequest import EntropyRequest from trezor.messages.FailureType import Other + from trezor.messages.wire_types import EntropyAck if strength not in (128, 192, 256): raise wire.FailureError( @@ -50,19 +45,23 @@ async def generate_mnemonic(strength, display_random, session_id): # if display_random: # raise wire.FailureError(Other, 'Entropy display not implemented') - ack = await wire.reply_message( - session_id, EntropyRequest(), EntropyAck) + ack = await wire.reply_message(session_id, EntropyRequest(), EntropyAck) + + if len(ack.entropy) != 32: + raise wire.FailureError(Other, 'Invalid entropy (has to be 32 bytes)') - strength_bytes = strength // 8 ctx = hashlib.sha256() - ctx.update(random.bytes(strength_bytes)) - ctx.update(ack.entropy[:strength_bytes]) + ctx.update(random.bytes(32)) + ctx.update(ack.entropy) entropy = ctx.digest() - return bip39.from_data(entropy) + return bip39.from_data(entropy[:strength // 8]) +@unimport async def show_mnemonic(mnemonic): + from trezor.ui.scroll import paginate + first_page = const(0) words_per_page = const(4) words = list(enumerate(mnemonic.split())) @@ -71,6 +70,9 @@ async def show_mnemonic(mnemonic): async def show_mnemonic_page(page, page_count, mnemonic): + from trezor.ui.button import Button, CONFIRM_BUTTON, CONFIRM_BUTTON_ACTIVE + from trezor.ui.scroll import render_scrollbar, animate_swipe + ui.clear() ui.display.text( 10, 30, 'Write down your seed', ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) diff --git a/src/apps/management/layout_wipe_device.py b/src/apps/management/layout_wipe_device.py index f097bb1af..123ed9eab 100644 --- a/src/apps/management/layout_wipe_device.py +++ b/src/apps/management/layout_wipe_device.py @@ -1,14 +1,12 @@ -from trezor import wire from trezor import ui from trezor.utils import unimport -from trezor.workflows.confirm import protect_with_confirm - -from .storage import clear_storage @unimport async def layout_wipe_device(message, session_id): from trezor.messages.Success import Success + from .confirm import require_confirm + from .storage import clear_storage ui.clear() ui.display.text(10, 30, 'Wiping device', ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) @@ -18,7 +16,7 @@ async def layout_wipe_device(message, session_id): ui.display.text(10, 164, 'All data will be lost.', ui.NORMAL, ui.WHITE, ui.BLACK) - await protect_with_confirm(session_id) + await require_confirm(session_id) clear_storage(session_id) diff --git a/src/trezor/workflows/request_pin.py b/src/apps/management/request_pin.py similarity index 96% rename from src/trezor/workflows/request_pin.py rename to src/apps/management/request_pin.py index f5f1a0dc8..fc1753266 100644 --- a/src/trezor/workflows/request_pin.py +++ b/src/apps/management/request_pin.py @@ -26,7 +26,7 @@ async def request_pin(session_id, *args, **kwargs): @unimport -async def request_pin_repeatedly(session_id): +async def request_pin_twice(session_id): from trezor.messages.FailureType import PinInvalid pin_first = await request_pin(session_id) diff --git a/src/apps/wallet/layout_sign_message.py b/src/apps/wallet/layout_sign_message.py index 748398dc5..227c4f374 100644 --- a/src/apps/wallet/layout_sign_message.py +++ b/src/apps/wallet/layout_sign_message.py @@ -3,8 +3,7 @@ from trezor.utils import unimport @unimport -async def layout_sign_message(message): - from trezor.workflows.confirm import protect_with_confirm +async def layout_sign_message(message, session_id): from trezor.messages.Success import Success ui.clear() @@ -12,8 +11,6 @@ async def layout_sign_message(message): ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) ui.display.text(10, 60, message.message, ui.MONO, ui.WHITE, ui.BLACK) - await protect_with_confirm(confirm='Sign') - # TODO return Success(message='Signed') diff --git a/src/trezor/wire/__init__.py b/src/trezor/wire/__init__.py index 66896b9dd..9ca5a356a 100644 --- a/src/trezor/wire/__init__.py +++ b/src/trezor/wire/__init__.py @@ -16,6 +16,10 @@ _session_handlers = {} # session id -> generator _workflow_genfuncs = {} # wire type -> (generator function, args) _opened_sessions = set() # session ids +# TODO: get rid of this, use callbacks instead +report_writer = write_report_stream() +report_writer.send(None) + def generate_session_id(): while True: @@ -54,30 +58,17 @@ def register_session(session_id, handler): def setup(): - report_writer = write_report_stream() - report_writer.send(None) - - open_session_handler = _handle_open_session(report_writer) - open_session_handler.send(None) - - close_session_handler = _handle_close_session(report_writer) - close_session_handler.send(None) - - fallback_session_handler = _handle_unknown_session() - fallback_session_handler.send(None) - session_dispatcher = dispatch_reports_by_session( _session_handlers, - open_session_handler, - close_session_handler, - fallback_session_handler) + _handle_open_session, + _handle_close_session, + _handle_unknown_session) session_dispatcher.send(None) - schedule_task(read_report_stream(session_dispatcher)) async def read_message(session_id, *exp_types): - log.info(__name__, 'reading message, one of %s', exp_types) + log.info(__name__, 'reading message of types %s', exp_types) future = Future() wire_decoder = decode_wire_stream( _dispatch_and_build_protobuf, session_id, exp_types, future) @@ -90,9 +81,7 @@ async def write_message(session_id, pbuf_message): log.info(__name__, 'writing message %s', pbuf_message) msg_data = await pbuf_message.dumps() msg_type = pbuf_message.message_type.wire_type - writer = write_report_stream() - writer.send(None) - encode_wire_message(msg_type, msg_data, session_id, writer) + encode_wire_message(msg_type, msg_data, session_id, report_writer) async def reply_message(session_id, pbuf_message, *exp_types): @@ -151,26 +140,21 @@ def protobuf_handler(msg_type, data_len, session_id, callback, *args): return pbuf_type.load(builder) -def _handle_open_session(write_target): - while True: - yield - session_id = open_session() - wire_decoder = decode_wire_stream(_handle_registered_type, session_id) - wire_decoder.send(None) - register_session(session_id, wire_decoder) - encode_session_open_message(session_id, write_target) +def _handle_open_session(): + session_id = open_session() + wire_decoder = decode_wire_stream(_handle_registered_type, session_id) + wire_decoder.send(None) + register_session(session_id, wire_decoder) + encode_session_open_message(session_id, report_writer) -def _handle_close_session(write_target): - while True: - session_id = yield - close_session(session_id) - encode_session_close_message(session_id, write_target) +def _handle_close_session(session_id): + close_session(session_id) + encode_session_close_message(session_id, report_writer) -def _handle_unknown_session(): - while True: - yield # TODO +def _handle_unknown_session(session_id, report_data): + pass # TODO def _dispatch_and_build_protobuf(msg_type, data_len, session_id, exp_types, future): diff --git a/src/trezor/workflows/__init__.py b/src/trezor/workflows/__init__.py deleted file mode 100644 index e69de29bb..000000000