wip: add wipe_code option to load_device, support in tests

common/protob/messages-management.proto

@@ -299,6 +299,7 @@ message LoadDevice {
     optional uint32 u2f_counter = 8;                        // U2F counter
     optional bool needs_backup = 9;                         // set "needs backup" flag
     optional bool no_backup = 10;                           // indicate that no backup is going to be made
+    optional string wipe_code = 11;                         // configure wipe code
 }
 
 /**

core/src/apps/debug/load_device.py

@@ -16,6 +16,8 @@ async def load_device(ctx: Context, msg: LoadDevice) -> Success:
     from trezor.ui.layouts import confirm_action
 
     mnemonics = msg.mnemonics  # local_cache_attribute
+    pin = msg.pin  # local_cache_attribute
+    wipe_code = msg.wipe_code  # local_cache_attribute
 
     # _validate
     if storage_device.is_initialized():
@@ -29,6 +31,11 @@ async def load_device(ctx: Context, msg: LoadDevice) -> Success:
             raise ProcessError(
                 "All shares are required to have the same number of words"
             )
+
+    if wipe_code and not pin:
+        raise ProcessError("PIN must be set for wipe code to work.")
+    if wipe_code == pin:
+        raise ProcessError("PIN and wipe code cannot be the same.")
     # END _validate
 
     is_slip39 = backup_types.is_slip39_word_count(word_count)
@@ -72,7 +79,13 @@ async def load_device(ctx: Context, msg: LoadDevice) -> Success:
     )
     storage_device.set_passphrase_enabled(bool(msg.passphrase_protection))
     storage_device.set_label(msg.label or "")
-    if msg.pin:
-        config.change_pin("", msg.pin, None, None)
+    if pin:
+        config.change_pin("", pin, None, None)
+
+        if wipe_code:
+            print("setting wipe code to", wipe_code)
+            config.change_wipe_code(pin, None, wipe_code)
+        else:
+            print("disabling wipe code")
 
     return Success(message="Device loaded")

core/src/trezor/messages.py

@@ -2368,6 +2368,7 @@ if TYPE_CHECKING:
         u2f_counter: "int | None"
         needs_backup: "bool | None"
         no_backup: "bool | None"
+        wipe_code: "str | None"
 
         def __init__(
             self,
@@ -2381,6 +2382,7 @@ if TYPE_CHECKING:
             u2f_counter: "int | None" = None,
             needs_backup: "bool | None" = None,
             no_backup: "bool | None" = None,
+            wipe_code: "str | None" = None,
         ) -> None:
             pass
 

legacy/firmware/config.c

@@ -506,6 +506,10 @@ void config_loadDevice(const LoadDevice *msg) {
 
   if (msg->has_pin) {
     config_changePin("", msg->pin);
+
+    if (msg->has_wipe_code) {
+      config_changeWipeCode(msg->pin, msg->wipe_code);
+    }
   }
 
   if (msg->mnemonics_count) {

legacy/firmware/fsm_msg_common.h

@@ -302,6 +302,21 @@ void fsm_msgLoadDevice(const LoadDevice *msg) {
     }
   }
 
+  if (msg->has_wipe_code && !msg->has_pin) {
+    fsm_sendFailure(FailureType_Failure_DataError,
+                    _("Wipe code provided without PIN"));
+    layoutHome();
+    return;
+  }
+
+  if (msg->has_wipe_code && msg->has_pin && strncmp(msg->pin, msg->wipe_code,
+                                                    sizeof(msg->pin)) == 0) {
+    fsm_sendFailure(FailureType_Failure_DataError,
+                    _("PIN and wipe code must not be the same"));
+    layoutHome();
+    return;
+  }
+
   config_loadDevice(msg);
   fsm_sendSuccess(_("Device loaded"));
   layoutHome();

legacy/firmware/protob/messages-management.options

@@ -21,6 +21,7 @@ LoadDevice.mnemonics                    max_count:16 max_size:241
 LoadDevice.pin                          max_size:51
 LoadDevice.language                     max_size:17
 LoadDevice.label                        max_size:33
+LoadDevice.wipe_code                    max_size:51
 
 ResetDevice.language                    max_size:17
 ResetDevice.label                       max_size:33

python/src/trezorlib/debuglink.py

@@ -1192,6 +1192,7 @@ def load_device(
     skip_checksum: bool = False,
     needs_backup: bool = False,
     no_backup: bool = False,
+    wipe_code: Optional[str] = None,
 ) -> protobuf.MessageType:
     if isinstance(mnemonic, str):
         mnemonic = [mnemonic]
@@ -1213,6 +1214,7 @@ def load_device(
             skip_checksum=skip_checksum,
             needs_backup=needs_backup,
             no_backup=no_backup,
+            wipe_code=wipe_code,
         )
     )
     client.init_device()

python/src/trezorlib/messages.py

@@ -3465,6 +3465,7 @@ class LoadDevice(protobuf.MessageType):
         8: protobuf.Field("u2f_counter", "uint32", repeated=False, required=False, default=None),
         9: protobuf.Field("needs_backup", "bool", repeated=False, required=False, default=None),
         10: protobuf.Field("no_backup", "bool", repeated=False, required=False, default=None),
+        11: protobuf.Field("wipe_code", "string", repeated=False, required=False, default=None),
     }
 
     def __init__(
@@ -3479,6 +3480,7 @@ class LoadDevice(protobuf.MessageType):
         u2f_counter: Optional["int"] = None,
         needs_backup: Optional["bool"] = None,
         no_backup: Optional["bool"] = None,
+        wipe_code: Optional["str"] = None,
     ) -> None:
         self.mnemonics: Sequence["str"] = mnemonics if mnemonics is not None else []
         self.pin = pin
@@ -3489,6 +3491,7 @@ class LoadDevice(protobuf.MessageType):
         self.u2f_counter = u2f_counter
         self.needs_backup = needs_backup
         self.no_backup = no_backup
+        self.wipe_code = wipe_code
 
 
 class ResetDevice(protobuf.MessageType):

tests/input_flows.py

@@ -76,30 +76,6 @@ class InputFlowBase:
         raise NotImplementedError
 
 
-class InputFlowSetupDevicePINWIpeCode(InputFlowBase):
-    def __init__(self, client: Client, pin: str, wipe_code: str):
-        super().__init__(client)
-        self.pin = pin
-        self.wipe_code = wipe_code
-
-    def input_flow_common(self) -> GeneratorType:
-        yield  # do you want to set/change the wipe code?
-        self.debug.press_yes()
-        yield from swipe_if_necessary(self.debug)  # wipe code info
-        self.debug.press_yes()
-
-        yield  # enter current pin
-        self.debug.input(self.pin)
-        yield  # enter new wipe code
-        self.debug.input(self.wipe_code)
-        yield  # please reenter
-        self.debug.press_yes()
-        yield  # enter new wipe code again
-        self.debug.input(self.wipe_code)
-        yield  # success
-        self.debug.press_yes()
-
-
 class InputFlowNewCodeMismatch(InputFlowBase):
     def __init__(
         self,

tests/persistence_tests/test_wipe_code.py

@@ -3,40 +3,27 @@ from trezorlib.debuglink import TrezorClientDebugLink as Client
 
 from ..common import MNEMONIC12
 from ..emulators import Emulator, EmulatorWrapper
-from ..input_flows import InputFlowSetupDevicePINWIpeCode
 from ..upgrade_tests import core_only, legacy_only
 
 PIN = "1234"
 WIPE_CODE = "9876"
 
 
-def setup_device_legacy(client: Client, pin: str, wipe_code: str) -> None:
+def setup_device(client: Client) -> None:
     device.wipe(client)
     debuglink.load_device(
-        client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE"
+        client,
+        MNEMONIC12,
+        PIN,
+        passphrase_protection=False,
+        label="WIPECODE",
+        wipe_code=WIPE_CODE,
     )
 
-    with client:
-        client.use_pin_sequence([PIN, WIPE_CODE, WIPE_CODE])
-        device.change_wipe_code(client)
-
-
-def setup_device_core(client: Client, pin: str, wipe_code: str) -> None:
-    device.wipe(client)
-    debuglink.load_device(
-        client, MNEMONIC12, pin, passphrase_protection=False, label="WIPECODE"
-    )
-
-    with client:
-        IF = InputFlowSetupDevicePINWIpeCode(client, pin, wipe_code)
-        client.set_input_flow(IF.get())
-        device.change_wipe_code(client)
-
 
 @core_only
 def test_wipe_code_activate_core(core_emulator: Emulator):
-    # set up device
-    setup_device_core(core_emulator.client, PIN, WIPE_CODE)
+    setup_device(core_emulator.client)
 
     core_emulator.client.init_device()
     device_id = core_emulator.client.features.device_id
@@ -72,7 +59,7 @@ def test_wipe_code_activate_core(core_emulator: Emulator):
 def test_wipe_code_activate_legacy():
     with EmulatorWrapper("legacy") as emu:
         # set up device
-        setup_device_legacy(emu.client, PIN, WIPE_CODE)
+        setup_device(emu.client)
 
         emu.client.init_device()
         device_id = emu.client.features.device_id