diff --git a/src/apps/management/layout_load_device.py b/src/apps/management/layout_load_device.py index fd8ba692ff..500f040e0d 100644 --- a/src/apps/management/layout_load_device.py +++ b/src/apps/management/layout_load_device.py @@ -1,24 +1,20 @@ -from trezor import wire from trezor import ui from trezor.utils import unimport -from trezor.workflows.confirm import confirm @unimport -def layout_load_device(message): +async def layout_load_device(session_id, message): + from trezor.workflows.confirm import protect_with_confirm + from trezor.messages.Success import Success ui.clear() - ui.display.text_center(120, 40, 'Really load device?', - ui.BOLD, 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() + await protect_with_confirm(session_id) - if confirmed: - from trezor.messages.Success import Success - yield from wire.write(Success(message='Loaded')) - else: - from trezor.messages.Failure import Failure - from trezor.messages.FailureType import ActionCancelled - yield from wire.write(Failure(message='Cancelled', code=ActionCancelled)) + # TODO + + return Success(message='Loaded') diff --git a/src/apps/management/layout_recovery_device.py b/src/apps/management/layout_recovery_device.py index 19511e0c18..c55cf56d5c 100644 --- a/src/apps/management/layout_recovery_device.py +++ b/src/apps/management/layout_recovery_device.py @@ -11,7 +11,7 @@ def nth(n): @unimport -def layout_recovery_device(message): +async def layout_recovery_device(session_id, message): msg = 'Please enter ' + nth(message.word_count) + ' word' @@ -20,4 +20,5 @@ def layout_recovery_device(message): 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) + + # TODO diff --git a/src/apps/management/layout_reset_device.py b/src/apps/management/layout_reset_device.py index b020d36d8f..2868429d2c 100644 --- a/src/apps/management/layout_reset_device.py +++ b/src/apps/management/layout_reset_device.py @@ -1,20 +1,23 @@ -from trezor import wire, ui, config -from trezor.workflows.request_pin import request_new_pin +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 - -APP_MANAGEMENT = const(1) -CFG_STORAGE = const(1) +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 + + if get_storage(session_id): + raise wire.FailureError( + UnexpectedMessage, 'Device is already initialized') mnemonic = await generate_mnemonic( message.strength, message.display_random, session_id) @@ -22,7 +25,7 @@ async def layout_reset_device(message, session_id): await show_mnemonic(mnemonic) if message.pin_protection: - pin = await request_new_pin(session_id) + pin = await request_pin_repeatedly(session_id) else: pin = '' @@ -30,10 +33,9 @@ async def layout_reset_device(message, session_id): version=1, pin=pin, mnemonic=mnemonic, passphrase_protection=message.passphrase_protection, language=message.language, label=message.label) + set_storage(session_id, await storage.dumps()) - config.set(session_id, APP_MANAGEMENT, CFG_STORAGE, await storage.dumps()) - - return Success() + return Success(message='Initialized') @unimport diff --git a/src/apps/management/layout_wipe_device.py b/src/apps/management/layout_wipe_device.py index 433b69d950..f097bb1af2 100644 --- a/src/apps/management/layout_wipe_device.py +++ b/src/apps/management/layout_wipe_device.py @@ -1,11 +1,14 @@ from trezor import wire from trezor import ui from trezor.utils import unimport -from trezor.workflows.confirm import confirm +from trezor.workflows.confirm import protect_with_confirm + +from .storage import clear_storage @unimport -def layout_wipe_device(message): +async def layout_wipe_device(message, session_id): + from trezor.messages.Success import Success ui.clear() ui.display.text(10, 30, 'Wiping device', ui.BOLD, ui.LIGHT_GREEN, ui.BLACK) @@ -15,12 +18,8 @@ def layout_wipe_device(message): ui.display.text(10, 164, 'All data will be lost.', ui.NORMAL, ui.WHITE, ui.BLACK) - confirmed = yield from confirm() + await protect_with_confirm(session_id) - 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)) + clear_storage(session_id) + + return Success(message='Wiped') diff --git a/src/apps/management/storage.py b/src/apps/management/storage.py new file mode 100644 index 0000000000..83d10ebf81 --- /dev/null +++ b/src/apps/management/storage.py @@ -0,0 +1,17 @@ +from trezor import config + + +APP_MANAGEMENT = const(1) +CFG_STORAGE = const(1) + + +def get_storage(session_id): + return config.get(session_id, APP_MANAGEMENT, CFG_STORAGE) + + +def set_storage(session_id, buf): + config.set(session_id, APP_MANAGEMENT, CFG_STORAGE, buf) + + +def clear_storage(session_id): + set_storage(session_id, '') diff --git a/src/apps/wallet/layout_get_public_key.py b/src/apps/wallet/layout_get_public_key.py index 11d0eb12bf..6f3e3db2cf 100644 --- a/src/apps/wallet/layout_get_public_key.py +++ b/src/apps/wallet/layout_get_public_key.py @@ -1,27 +1,22 @@ from trezor import wire, ui from trezor.utils import unimport -from trezor.workflows.request_pin import request_pin @unimport -def layout_get_public_key(message): +async def layout_get_public_key(session_id, message): + from trezor.messages.PublicKey import PublicKey + from trezor.messages.HDNodeType import HDNodeType - ui.clear() - pin = yield from request_pin() + # TODO: protect with pin + # TODO: fail if not initialized + # TODO: derive correct node - if pin is not None: - from trezor.messages.PublicKey import PublicKey - from trezor.messages.HDNodeType import HDNodeType - pubkey = PublicKey() - pubkey.node = HDNodeType() - pubkey.node.depth = 0 - pubkey.node.child_num = 0 - pubkey.node.fingerprint = 0 - pubkey.node.chain_code = 'deadbeef' - pubkey.node.public_key = 'deadbeef' - yield from wire.write(pubkey) + pubkey = PublicKey() + pubkey.node = HDNodeType() + pubkey.node.depth = 0 + pubkey.node.child_num = 0 + pubkey.node.fingerprint = 0 + pubkey.node.chain_code = 'deadbeef' + pubkey.node.public_key = 'deadbeef' - else: - from trezor.messages.Failure import Failure - from trezor.messages.FailureType import ActionCancelled - yield from wire.write(Failure(message='Cancelled', code=ActionCancelled)) + return pubkey diff --git a/src/trezor/workflows/confirm.py b/src/trezor/workflows/confirm.py index 292b33e7e9..1103dbf9bf 100644 --- a/src/trezor/workflows/confirm.py +++ b/src/trezor/workflows/confirm.py @@ -3,18 +3,26 @@ from trezor.utils import unimport @unimport -async def confirm(content=None, code=None, **kwargs): +async def confirm(session_id, content=None, code=None, **kwargs): 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 + from trezor.messages.wire_types import ButtonAck dialog = ConfirmDialog(content, **kwargs) dialog.render() if code is None: code = Other - ack = await wire.call(ButtonRequest(code=code), ButtonAck) - res = await dialog + await wire.reply_message(session_id, ButtonRequest(code=code), ButtonAck) + return await dialog == CONFIRMED - return res == CONFIRMED + +@unimport +async def protect_with_confirm(*args, **kwargs): + from trezor.messages.FailureType import ActionCancelled + + confirmed = await confirm(*args, **kwargs) + + if not confirmed: + raise wire.FailureError(ActionCancelled, 'Cancelled') diff --git a/src/trezor/workflows/request_pin.py b/src/trezor/workflows/request_pin.py index e057aef61b..f5f1a0dc8d 100644 --- a/src/trezor/workflows/request_pin.py +++ b/src/trezor/workflows/request_pin.py @@ -26,7 +26,7 @@ async def request_pin(session_id, *args, **kwargs): @unimport -async def request_new_pin(session_id): +async def request_pin_repeatedly(session_id): from trezor.messages.FailureType import PinInvalid pin_first = await request_pin(session_id)