trezor-firmware/src/apps/common/request_passphrase.py

from trezor import ui, wire
from trezor.messages import ButtonRequestType, MessageType, PassphraseSourceType
from trezor.messages.ButtonRequest import ButtonRequest
from trezor.messages.PassphraseRequest import PassphraseRequest
from trezor.messages.PassphraseStateRequest import PassphraseStateRequest
from trezor.ui.entry_select import DEVICE, EntrySelector
from trezor.ui.passphrase import CANCELLED, PassphraseKeyboard
from trezor.ui.text import Text

from apps.common import storage
from apps.common.cache import get_state


@ui.layout
async def request_passphrase_entry(ctx):
    text = Text("Enter passphrase", ui.ICON_CONFIG)
    text.normal("Where to enter your", "passphrase?")
    text.render()

    ack = await ctx.call(
        ButtonRequest(code=ButtonRequestType.PassphraseType),
        MessageType.ButtonAck,
        MessageType.Cancel,
    )
    if ack.MESSAGE_WIRE_TYPE == MessageType.Cancel:
        raise wire.ActionCancelled("Passphrase cancelled")

    selector = EntrySelector(text)
    return await ctx.wait(selector)


@ui.layout
async def request_passphrase_ack(ctx, on_device):
    if not on_device:
        text = Text("Passphrase entry", ui.ICON_CONFIG)
        text.normal("Please, type passphrase", "on connected host.")
        text.render()

    req = PassphraseRequest(on_device=on_device)
    ack = await ctx.call(req, MessageType.PassphraseAck, MessageType.Cancel)
    if ack.MESSAGE_WIRE_TYPE == MessageType.Cancel:
        raise wire.ActionCancelled("Passphrase cancelled")

    if on_device:
        if ack.passphrase is not None:
            raise wire.ProcessError("Passphrase provided when it should not be")
        keyboard = PassphraseKeyboard("Enter passphrase")
        passphrase = await ctx.wait(keyboard)
        if passphrase == CANCELLED:
            raise wire.ActionCancelled("Passphrase cancelled")
    else:
        if ack.passphrase is None:
            raise wire.ProcessError("Passphrase not provided")
        passphrase = ack.passphrase

    req = PassphraseStateRequest(
        state=get_state(prev_state=ack.state, passphrase=passphrase)
    )
    ack = await ctx.call(req, MessageType.PassphraseStateAck, MessageType.Cancel)

    return passphrase


async def request_passphrase(ctx):
    if storage.get_passphrase_source() == PassphraseSourceType.ASK:
        on_device = await request_passphrase_entry(ctx) == DEVICE
    else:
        on_device = storage.get_passphrase_source() == PassphraseSourceType.DEVICE
    passphrase = await request_passphrase_ack(ctx, on_device)
    return passphrase


async def protect_by_passphrase(ctx):
    if storage.has_passphrase():
        return await request_passphrase(ctx)
    else:
        return ""