WIP - set_mnemonic_secret

core/embed/rust/librust_qstr.h

@@ -30,6 +30,7 @@ static void _librust_qstrs(void) {
   MP_QSTR_get_label;
   MP_QSTR_set_label;
   MP_QSTR_get_mnemonic_secret;
+  MP_QSTR_set_mnemonic_secret;
   MP_QSTR_is_passphrase_enabled;
   MP_QSTR_set_passphrase_enabled;
   MP_QSTR_get_passphrase_always_on_device;
@@ -68,4 +69,5 @@ static void _librust_qstrs(void) {
   MP_QSTR_verb;
   MP_QSTR_verb_cancel;
   MP_QSTR_reverse;
+  MP_QSTR_secret;
 }

core/embed/rust/src/storagedevice/storage_device.rs

@@ -77,6 +77,8 @@ const AUTOLOCK_DELAY_MAXIMUM: u32 = 0x2000_0000; // ~6 days
 // Other SD-salt-related constants are in core/src/storage/sd_salt.py
 const SD_SALT_AUTH_KEY_LEN_BYTES: u8 = 16;
 
+const STORAGE_VERSION_CURRENT: u8 = 0x02;
+
 extern "C" {
     // storage.h
     fn storage_has(key: u16) -> secbool::Secbool;
@@ -196,6 +198,50 @@ extern "C" fn storagedevice_get_mnemonic_secret() -> Obj {
     unsafe { util::try_or_raise(block) }
 }
 
+extern "C" fn storagedevice_set_mnemonic_secret(
+    n_args: usize,
+    args: *const Obj,
+    kwargs: *mut Map,
+) -> Obj {
+    let block = |_args: &[Obj], kwargs: &Map| {
+        let secret: Buffer = kwargs.get(Qstr::MP_QSTR_secret)?.try_into()?;
+        let backup_type: u8 = kwargs.get(Qstr::MP_QSTR_backup_type)?.try_into()?;
+        let needs_backup: bool = if kwargs.contains_key(Qstr::MP_QSTR_needs_backup) {
+            kwargs.get(Qstr::MP_QSTR_needs_backup)?.try_into()?
+        } else {
+            false
+        };
+        let no_backup: bool = if kwargs.contains_key(Qstr::MP_QSTR_no_backup) {
+            kwargs.get(Qstr::MP_QSTR_needs_backup)?.try_into()?
+        } else {
+            false
+        };
+
+        let key = _get_appkey(_VERSION, false);
+        storagedevice_storage_set(key, &[STORAGE_VERSION_CURRENT] as *const _, 1);
+
+        let key = _get_appkey(_MNEMONIC_SECRET, false);
+        storagedevice_storage_set(key, secret.as_ptr(), secret.len() as u16);
+
+        let key = _get_appkey(_BACKUP_TYPE, false);
+        storagedevice_storage_set_u8(key, backup_type);
+
+        let key = _get_appkey(_NO_BACKUP, false);
+        storagedevice_storage_set_true_or_delete(key, no_backup);
+
+        let key = _get_appkey(INITIALIZED, true);
+        storagedevice_storage_set_bool(key, true);
+
+        if !no_backup {
+            let key = _get_appkey(_NEEDS_BACKUP, false);
+            storagedevice_storage_set_true_or_delete(key, needs_backup);
+        }
+
+        Ok(true.into())
+    };
+    unsafe { util::try_with_args_and_kwargs(n_args, args, kwargs, block) }
+}
+
 extern "C" fn storagedevice_is_passphrase_enabled() -> Obj {
     let block = || {
         let key = _get_appkey(_USE_PASSPHRASE, false);
@@ -545,6 +591,14 @@ pub fn storagedevice_storage_set_bool(key: u16, val: bool) -> bool {
     storagedevice_storage_set(key, &val as *const _, 1)
 }
 
+pub fn storagedevice_storage_set_true_or_delete(key: u16, val: bool) -> bool {
+    if val {
+        storagedevice_storage_set_bool(key, true)
+    } else {
+        storagedevice_storage_delete(key)
+    }
+}
+
 pub fn storagedevice_storage_set_u8(key: u16, val: u8) -> bool {
     storagedevice_storage_set(key, &val.to_be_bytes() as *const _, 1)
 }
@@ -592,6 +646,7 @@ pub static mp_module_trezorstoragedevice: Module = obj_module! {
 
     // TODO: should we return None or True in case of set_xxx?
 
+    /// BackupTypeInt = TypeVar("BackupTypeInt", bound=int)
     /// StorageSafetyCheckLevel = Literal[0, 1]
 
     /// def is_version_stored() -> bool:
@@ -630,6 +685,15 @@ pub static mp_module_trezorstoragedevice: Module = obj_module! {
     ///     """Get mnemonic secret."""
     Qstr::MP_QSTR_get_mnemonic_secret => obj_fn_0!(storagedevice_get_mnemonic_secret).as_obj(),
 
+    /// def set_mnemonic_secret(
+    ///     secret: bytes,
+    ///     backup_type: BackupTypeInt,
+    ///     needs_backup: bool = False,
+    ///     no_backup: bool = False,
+    /// ) -> None:
+    ///     """Set mnemonic secret"""
+    Qstr::MP_QSTR_set_mnemonic_secret => obj_fn_kw!(0, storagedevice_set_mnemonic_secret).as_obj(),
+
     /// def is_passphrase_enabled() -> bool:
     ///     """Whether passphrase is enabled."""
     Qstr::MP_QSTR_is_passphrase_enabled => obj_fn_0!(storagedevice_is_passphrase_enabled).as_obj(),

core/mocks/generated/trezorstoragedevice.pyi

@@ -1,4 +1,5 @@
 from typing import *
+BackupTypeInt = TypeVar("BackupTypeInt", bound=int)
 StorageSafetyCheckLevel = Literal[0, 1]
 
 
@@ -47,6 +48,16 @@ def get_mnemonic_secret() -> bytes:
     """Get mnemonic secret."""
 
 
+# rust/src/storagedevice/storage_device.rs
+def set_mnemonic_secret(
+    secret: bytes,
+    backup_type: BackupTypeInt,
+    needs_backup: bool = False,
+    no_backup: bool = False,
+) -> None:
+    """Set mnemonic secret"""
+
+
 # rust/src/storagedevice/storage_device.rs
 def is_passphrase_enabled() -> bool:
     """Whether passphrase is enabled."""

core/src/apps/debug/load_device.py

@@ -36,9 +36,9 @@ async def load_device(ctx: wire.Context, msg: LoadDevice) -> Success:
         storagedevice.set_slip39_identifier(identifier)
         storagedevice.set_slip39_iteration_exponent(iteration_exponent)
 
-    storage.device.store_mnemonic_secret(
-        secret,
-        backup_type,
+    storagedevice.set_mnemonic_secret(
+        secret=secret,
+        backup_type=backup_type,
         needs_backup=msg.needs_backup is True,
         no_backup=msg.no_backup is True,
     )

core/src/apps/management/recovery_device/homescreen.py

@@ -129,8 +129,8 @@ async def _finish_recovery(
     if backup_type is None:
         raise RuntimeError
 
-    storage.device.store_mnemonic_secret(
-        secret, backup_type, needs_backup=False, no_backup=False
+    storagedevice.set_mnemonic_secret(
+        secret=secret, backup_type=backup_type, needs_backup=False, no_backup=False
     )
     if backup_type in (BackupType.Slip39_Basic, BackupType.Slip39_Advanced):
         identifier = storage.recovery.get_slip39_identifier()

core/src/apps/management/reset_device/__init__.py

@@ -86,9 +86,9 @@ async def reset_device(ctx: wire.Context, msg: ResetDevice) -> Success:
     if msg.label is not None:
         storage.device.set_label(msg.label)
     storagedevice.set_passphrase_enabled(bool(msg.passphrase_protection))
-    storage.device.store_mnemonic_secret(
-        secret,  # for SLIP-39, this is the EMS
-        msg.backup_type,
+    storagedevice.set_mnemonic_secret(
+        secret=secret,  # for SLIP-39, this is the EMS
+        backup_type=msg.backup_type,
         needs_backup=not perform_backup,
         no_backup=bool(msg.no_backup),
     )

core/src/storage/device.py

@@ -4,7 +4,6 @@ from ubinascii import hexlify
 
 import storage.cache
 from storage import common
-from trezor import storagedevice
 
 # TODO: export all the constants also in storagedevice?
 
@@ -160,19 +159,19 @@ def get_backup_type() -> BackupType:
 #     common.set(_NAMESPACE, _HOMESCREEN, homescreen, public=True)
 
 
-def store_mnemonic_secret(
-    secret: bytes,
-    backup_type: BackupType,
-    needs_backup: bool = False,
-    no_backup: bool = False,
-) -> None:
-    storagedevice.set_version(common.STORAGE_VERSION_CURRENT)
-    common.set(_NAMESPACE, _MNEMONIC_SECRET, secret)
-    common.set_uint8(_NAMESPACE, _BACKUP_TYPE, backup_type)
-    common.set_true_or_delete(_NAMESPACE, _NO_BACKUP, no_backup)
-    common.set_bool(_NAMESPACE, INITIALIZED, True, public=True)
-    if not no_backup:
-        common.set_true_or_delete(_NAMESPACE, _NEEDS_BACKUP, needs_backup)
+# def store_mnemonic_secret(
+#     secret: bytes,
+#     backup_type: BackupType,
+#     needs_backup: bool = False,
+#     no_backup: bool = False,
+# ) -> None:
+#     storagedevice.set_version(common.STORAGE_VERSION_CURRENT)
+#     common.set(_NAMESPACE, _MNEMONIC_SECRET, secret)
+#     common.set_uint8(_NAMESPACE, _BACKUP_TYPE, backup_type)
+#     common.set_true_or_delete(_NAMESPACE, _NO_BACKUP, no_backup)
+#     common.set_bool(_NAMESPACE, INITIALIZED, True, public=True)
+#     if not no_backup:
+#         common.set_true_or_delete(_NAMESPACE, _NEEDS_BACKUP, needs_backup)
 
 
 # def needs_backup() -> bool: