You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
trezor-firmware/core/src/apps/management/recovery_device/keyboard_slip39.py

233 lines
7.4 KiB

from trezor import io, loop, res, ui, workflow
from trezor.crypto import slip39
from trezor.ui import display
from trezor.ui.components.tt.button import (
Button,
ButtonClear,
ButtonMono,
ButtonMonoConfirm,
)
if False:
from trezor.ui.components.tt.button import ButtonContent, ButtonStyleStateType
class KeyButton(Button):
def __init__(
self,
area: ui.Area,
content: ButtonContent,
keyboard: "Slip39Keyboard",
index: int,
):
self.keyboard = keyboard
self.index = index
super().__init__(area, content)
def on_click(self) -> None:
self.keyboard.on_key_click(self)
class InputButton(Button):
def __init__(self, area: ui.Area, keyboard: "Slip39Keyboard") -> None:
super().__init__(area, "")
self.word = ""
self.pending_button: Button | None = None
self.pending_index: int | None = None
self.keyboard = keyboard
self.disable()
def edit(
self,
text: str,
word: str,
pending_button: Button | None,
pending_index: int | None,
) -> None:
self.word = word
self.text = text
self.pending_button = pending_button
self.pending_index = pending_index
self.repaint = True
if word: # confirm button
self.enable()
self.normal_style = ButtonMonoConfirm.normal
self.active_style = ButtonMonoConfirm.active
self.icon = res.load(ui.ICON_CONFIRM)
else: # disabled button
self.disabled_style = ButtonMono.disabled
self.disable()
self.icon = b""
def render_content(
self, s: ButtonStyleStateType, ax: int, ay: int, aw: int, ah: int
) -> None:
text_style = s.text_style
fg_color = s.fg_color
bg_color = s.bg_color
tx = ax + 16 # x-offset of the content
ty = ay + ah // 2 + 8 # y-offset of the content
if not self.keyboard.is_input_final():
pending_button = self.pending_button
pending_index = self.pending_index
to_display = len(self.text) * "*"
if pending_button and pending_index is not None:
to_display = to_display[:-1] + pending_button.text[pending_index]
else:
to_display = self.word
display.text(tx, ty, to_display, text_style, fg_color, bg_color)
if self.pending_button and not self.keyboard.is_input_final():
width = display.text_width(to_display, text_style)
pw = display.text_width(self.text[-1:], text_style)
px = tx + width - pw
display.bar(px, ty + 2, pw + 1, 3, fg_color)
if self.icon:
ix = ax + aw - 16 * 2
iy = ty - 16
display.icon(ix, iy, self.icon, fg_color, bg_color)
class Prompt(ui.Component):
def __init__(self, prompt: str) -> None:
super().__init__()
self.prompt = prompt
def on_render(self) -> None:
if self.repaint:
display.bar(0, 8, ui.WIDTH, 60, ui.BG)
display.text(20, 40, self.prompt, ui.BOLD, ui.GREY, ui.BG)
self.repaint = False
class Slip39Keyboard(ui.Layout):
def __init__(self, prompt: str) -> None:
super().__init__()
self.prompt = Prompt(prompt)
icon_back = res.load(ui.ICON_BACK)
self.back = Button(ui.grid(0, n_x=3, n_y=4), icon_back, ButtonClear)
self.back.on_click = self.on_back_click # type: ignore
self.input = InputButton(ui.grid(1, n_x=3, n_y=4, cells_x=2), self)
self.input.on_click = self.on_input_click # type: ignore
self.keys = [
KeyButton(ui.grid(i + 3, n_y=4), k, self, i + 1)
for i, k in enumerate(
("ab", "cd", "ef", "ghij", "klm", "nopq", "rs", "tuv", "wxyz")
)
]
self.pending_button: Button | None = None
self.pending_index = 0
self.button_sequence = ""
self.mask = slip39.KEYBOARD_FULL_MASK
def dispatch(self, event: int, x: int, y: int) -> None:
for btn in self.keys:
btn.dispatch(event, x, y)
if self.input.text:
self.input.dispatch(event, x, y)
self.back.dispatch(event, x, y)
else:
self.prompt.dispatch(event, x, y)
def on_back_click(self) -> None:
# Backspace was clicked, let's delete the last character of input.
self.button_sequence = self.button_sequence[:-1]
self.edit()
def on_input_click(self) -> None:
# Input button was clicked. If the content matches the suggested word,
# let's confirm it, otherwise just auto-complete.
result = self.input.word
if self.is_input_final():
self.button_sequence = ""
self.edit()
self.on_confirm(result)
def on_key_click(self, btn: KeyButton) -> None:
# Key button was clicked. If this button is pending, let's cycle the
# pending character in input. If not, let's just append the first
# character.
if self.pending_button is btn:
index = (self.pending_index + 1) % len(btn.text)
else:
index = 0
self.button_sequence += str(btn.index)
self.edit(btn, index)
def on_timeout(self) -> None:
# Timeout occurred. Let's redraw to draw asterisks.
self.edit()
def on_confirm(self, word: str) -> None:
# Word was confirmed by the user.
raise ui.Result(word)
def edit(self, button: Button | None = None, index: int = 0) -> None:
self.pending_button = button
self.pending_index = index
# find the completions
word = ""
self.mask = slip39.word_completion_mask(self.button_sequence)
if self.is_input_final():
word = slip39.button_sequence_to_word(self.button_sequence)
# modify the input state
self.input.edit(
self.button_sequence, word, self.pending_button, self.pending_index
)
# enable or disable key buttons
for btn in self.keys:
if self.is_input_final():
btn.disable()
elif btn is button or self.check_mask(btn.index):
btn.enable()
else:
btn.disable()
# invalidate the prompt if we display it next frame
if not self.input.text:
self.prompt.repaint = True
def is_input_final(self) -> bool:
# returns True if mask has exactly one bit set to 1 or is 0
return not (self.mask & (self.mask - 1))
def check_mask(self, index: int) -> bool:
return bool((1 << (index - 1)) & self.mask)
async def handle_input(self) -> None:
touch = loop.wait(io.TOUCH)
timeout = loop.sleep(1000)
race_touch = loop.race(touch)
race_timeout = loop.race(touch, timeout)
while True:
if self.pending_button is not None:
race = race_timeout
else:
race = race_touch
result = await race
if touch in race.finished:
event, x, y = result
workflow.idle_timer.touch()
self.dispatch(event, x, y)
else:
self.on_timeout()
if __debug__:
def create_tasks(self) -> tuple[loop.Task, ...]:
from apps.debug import input_signal
return super().create_tasks() + (input_signal(),)