From 3279c6832cb81138cc44df2cc1e9ce5cea235d00 Mon Sep 17 00:00:00 2001
From: Marko Bencun <marko@shiftcrypto.ch>
Date: Thu, 18 May 2023 05:03:32 +0200
Subject: [PATCH] ui/components/trinary_input_string: allow mode to enter only
 numbers.

In BIP-85, the user needs to specify the index at which to derive a
secret. This change allows the keyboard input to be used to enter
integers.
---
 src/rust/bitbox02/src/ui/types.rs        |  2 ++
 src/ui/components/trinary_input_string.c | 10 +++++++++-
 src/ui/components/trinary_input_string.h |  2 ++
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/rust/bitbox02/src/ui/types.rs b/src/rust/bitbox02/src/ui/types.rs
index 36ba4d9504..c51c56231d 100644
--- a/src/rust/bitbox02/src/ui/types.rs
+++ b/src/rust/bitbox02/src/ui/types.rs
@@ -102,6 +102,7 @@ pub struct TrinaryInputStringParams<'a> {
     pub title: &'a str,
     /// Currently specialized to the BIP39 wordlist. Can be extended if needed.
     pub wordlist: Option<&'a crate::keystore::Bip39Wordlist>,
+    pub number_input: bool,
     pub hide: bool,
     pub special_chars: bool,
     pub longtouch: bool,
@@ -132,6 +133,7 @@ impl<'a> TrinaryInputStringParams<'a> {
                 None => 0,
                 Some(wordlist) => wordlist.len() as _,
             },
+            number_input: self.number_input,
             hide: self.hide,
             special_chars: self.special_chars,
             longtouch: self.longtouch,
diff --git a/src/ui/components/trinary_input_string.c b/src/ui/components/trinary_input_string.c
index 0953264246..bcc38a05c7 100644
--- a/src/ui/components/trinary_input_string.c
+++ b/src/ui/components/trinary_input_string.c
@@ -64,6 +64,7 @@ typedef struct {
     // Can be NULL.
     const char* const* wordlist;
     size_t wordlist_size;
+    bool number_input;
     // Only applies if wordlist != NULL: determines if a word from the wordlist was entered.
     bool can_confirm;
     // Mask user input with '*'?
@@ -294,6 +295,8 @@ static void _set_alphabet(component_t* trinary_input_string)
         }
         // Since wordlist is sorted, charset is sorted automatically.
         trinary_input_char_set_alphabet(trinary_char, charset, 1);
+    } else if (data->number_input) {
+        trinary_input_char_set_alphabet(trinary_char, _digits, 1);
     } else {
         // Otherwise set the input charset based on the user selected keyboard mode.
         keyboard_mode_t keyboard_mode = keyboard_current_mode(data->keyboard_switch_component);
@@ -450,12 +453,17 @@ component_t* trinary_input_string_create(
     memset(component, 0, sizeof(component_t));
     memset(data, 0, sizeof(data_t));
 
+    if (params->number_input && data->wordlist != NULL) {
+        Abort("trinary_input_string: invalid params");
+    }
+
     data->confirm_cb = confirm_cb;
     data->confirm_callback_param = confirm_callback_param;
     data->cancel_cb = cancel_cb;
     data->cancel_callback_param = cancel_callback_param;
     data->wordlist = params->wordlist;
     data->wordlist_size = params->wordlist_size;
+    data->number_input = params->number_input;
     data->hide = params->hide;
     data->longtouch = params->longtouch;
     data->cancel_is_backbutton = params->cancel_is_backbutton;
@@ -481,7 +489,7 @@ component_t* trinary_input_string_create(
     data->confirm_component = confirm_button_create(params->longtouch, ICON_BUTTON_CHECK);
     ui_util_add_sub_component(component, data->confirm_component);
 
-    if (params->wordlist == NULL) {
+    if (params->wordlist == NULL && !params->number_input) {
         data->keyboard_switch_component =
             keyboard_switch_create(top_slider, params->special_chars, component);
         ui_util_add_sub_component(component, data->keyboard_switch_component);
diff --git a/src/ui/components/trinary_input_string.h b/src/ui/components/trinary_input_string.h
index 8365577d1d..96ac6602e5 100644
--- a/src/ui/components/trinary_input_string.h
+++ b/src/ui/components/trinary_input_string.h
@@ -26,6 +26,8 @@ typedef struct {
     const char* title;
     // Restrict and autocomplete to this list of words. Set to NULL to allow arbitrary input.
     const char* const* wordlist;
+    // If true, the user can enter numbers only.
+    bool number_input;
     // Set to 0 if wordlist is NULL.
     size_t wordlist_size;
     // Mask the chars entered as `*`. For password input.