Entropy check workflow in ResetDevice

common/protob/messages-management.proto

@@ -423,6 +423,7 @@ message ResetDevice {
     optional bool skip_backup = 8;                          // postpone seed backup to BackupDevice workflow
     optional bool no_backup = 9;                            // indicate that no backup is going to be made
     optional BackupType backup_type = 10 [default=Bip39];   // type of the mnemonic backup
+    optional bool entropy_check = 11;                       // run with entropy check protocol
 }
 
 /**
@@ -444,14 +445,34 @@ message BackupDevice {
  * @next EntropyAck
  */
 message EntropyRequest {
+    optional bytes entropy_commitment = 1;  // HMAC-SHA256 of Trezor's internal entropy used in entropy check.
+    optional bytes prev_entropy = 2;        // Trezor's internal entropy from the previous round of entropy check.
 }
 
 /**
  * Request: Provide additional entropy for seed generation function
  * @next Success
+ * @next EntropyCheckReady
  */
 message EntropyAck {
-    required bytes entropy = 1;     // 256 bits (32 bytes) of random data
+    required bytes entropy = 1;     // 256 bits (32 bytes) of the host's random data
+}
+
+/**
+ * Response: Trezor is ready for the next phase of the entropy check protocol.
+ * @next EntropyCheckContinue
+ * @next GetPublicKey
+ */
+message EntropyCheckReady {
+}
+
+/**
+ * Request: Proceed with the next phase of the entropy check protocol, asking Trezor to either reveal its internal entropy or to finish and store the seed.
+ * @next Success
+ * @next EntropyRequest
+ */
+message EntropyCheckContinue {
+    optional bool finish = 1 [default=false];  // finish the entropy check protocol, store the seed
 }
 
 /**

common/protob/messages.proto

@@ -45,6 +45,8 @@ enum MessageType {
     MessageType_BackupDevice = 34 [(bitcoin_only) = true, (wire_in) = true];
     MessageType_EntropyRequest = 35 [(bitcoin_only) = true, (wire_out) = true];
     MessageType_EntropyAck = 36 [(bitcoin_only) = true, (wire_in) = true];
+    MessageType_EntropyCheckReady = 994 [(bitcoin_only) = true, (wire_out) = true];
+    MessageType_EntropyCheckContinue = 995 [(bitcoin_only) = true, (wire_in) = true];
     MessageType_PassphraseRequest = 41 [(bitcoin_only) = true, (wire_out) = true];
     MessageType_PassphraseAck = 42 [(bitcoin_only) = true, (wire_in) = true, (wire_tiny) = true, (wire_no_fsm) = true];
     MessageType_RecoveryDevice = 45 [(bitcoin_only) = true, (wire_in) = true];

core/.changelog.d/4155.added

@@ -0,0 +1 @@
+Entropy check workflow in ResetDevice.

core/src/apps/bitcoin/get_public_key.py

@@ -6,9 +6,13 @@
     from trezor.messages import GetPublicKey, PublicKey
     from trezor.protobuf import MessageType
 
+    from apps.common.keychain import Keychain
+
 
 async def get_public_key(
-    msg: GetPublicKey, auth_msg: MessageType | None = None
+    msg: GetPublicKey,
+    auth_msg: MessageType | None = None,
+    keychain: Keychain | None = None,
 ) -> PublicKey:
     from trezor import TR, wire
     from trezor.enums import InputScriptType
@@ -34,7 +38,8 @@ async def get_public_key(
         if auth_msg.address_n != address_n[: len(auth_msg.address_n)]:
             raise FORBIDDEN_KEY_PATH
 
-    keychain = await get_keychain(curve_name, [paths.AlwaysMatchingSchema])
+    if not keychain:
+        keychain = await get_keychain(curve_name, [paths.AlwaysMatchingSchema])
 
     node = keychain.derive(address_n)
 

core/src/apps/common/mnemonic.py

@@ -32,8 +32,12 @@ def is_bip39() -> bool:
     return get_type() == BackupType.Bip39
 
 
-def get_seed(passphrase: str = "", progress_bar: bool = True) -> bytes:
-    mnemonic_secret = get_secret()
+def get_seed(
+    passphrase: str = "",
+    progress_bar: bool = True,
+    mnemonic_secret: bytes | None = None,
+) -> bytes:
+    mnemonic_secret = mnemonic_secret or get_secret()
     if mnemonic_secret is None:
         raise ValueError  # Mnemonic not set
 

core/src/apps/debug/load_device.py

@@ -62,10 +62,10 @@ async def load_device(msg: LoadDevice) -> Success:
 
     storage_device.store_mnemonic_secret(
         secret,
-        backup_type,
         needs_backup=msg.needs_backup is True,
         no_backup=msg.no_backup is True,
     )
+    storage_device.set_backup_type(backup_type)
     storage_device.set_passphrase_enabled(bool(msg.passphrase_protection))
     storage_device.set_label(msg.label or "")
     if msg.pin:

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

@@ -212,9 +212,8 @@ async def _finish_recovery(secret: bytes, backup_type: BackupType) -> Success:
     if backup_type is None:
         raise RuntimeError
 
-    storage_device.store_mnemonic_secret(
-        secret, backup_type, needs_backup=False, no_backup=False
-    )
+    storage_device.store_mnemonic_secret(secret, needs_backup=False, no_backup=False)
+    storage_device.set_backup_type(backup_type)
     if backup_types.is_slip39_backup_type(backup_type):
         if not backup_types.is_extendable_backup_type(backup_type):
             identifier = storage_recovery.get_slip39_identifier()

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

@@ -3,8 +3,8 @@
 import storage
 import storage.device as storage_device
 from trezor import TR
-from trezor.crypto import slip39
-from trezor.enums import BackupType
+from trezor.crypto import hmac, slip39
+from trezor.enums import BackupType, MessageType
 from trezor.ui.layouts import confirm_action
 from trezor.wire import ProcessError
 
@@ -63,33 +63,52 @@ async def reset_device(msg: ResetDevice) -> Success:
     # wipe storage to make sure the device is in a clear state
     storage.reset()
 
+    # Check backup type, perform type-specific handling
+    if backup_types.is_slip39_backup_type(backup_type):
+        # set SLIP39 parameters
+        storage_device.set_slip39_iteration_exponent(slip39.DEFAULT_ITERATION_EXPONENT)
+    elif backup_type != BAK_T_BIP39:
+        # Unknown backup type.
+        raise RuntimeError
+
+    storage_device.set_backup_type(backup_type)
+
     # request and set new PIN
     if msg.pin_protection:
         newpin = await request_pin_confirm()
         if not config.change_pin("", newpin, None, None):
             raise ProcessError("Failed to set PIN")
 
-    # generate and display internal entropy
-    int_entropy = random.bytes(32, True)
-    if __debug__:
-        storage.debug.reset_internal_entropy = int_entropy
+    prev_int_entropy = None
+    while True:
+        # generate internal entropy
+        int_entropy = random.bytes(32, True)
+        if __debug__:
+            storage.debug.reset_internal_entropy = int_entropy
 
-    # request external entropy and compute the master secret
-    entropy_ack = await call(EntropyRequest(), EntropyAck)
-    ext_entropy = entropy_ack.entropy
-    # For SLIP-39 this is the Encrypted Master Secret
-    secret = _compute_secret_from_entropy(int_entropy, ext_entropy, msg.strength)
+        entropy_commitment = (
+            hmac(hmac.SHA256, int_entropy, b"").digest() if msg.entropy_check else None
+        )
 
-    # Check backup type, perform type-specific handling
-    if backup_type == BAK_T_BIP39:
-        # in BIP-39 we store mnemonic string instead of the secret
-        secret = bip39.from_data(secret).encode()
-    elif backup_types.is_slip39_backup_type(backup_type):
-        # generate and set SLIP39 parameters
-        storage_device.set_slip39_iteration_exponent(slip39.DEFAULT_ITERATION_EXPONENT)
-    else:
-        # Unknown backup type.
-        raise RuntimeError
+        # request external entropy and compute the master secret
+        entropy_ack = await call(
+            EntropyRequest(
+                entropy_commitment=entropy_commitment, prev_entropy=prev_int_entropy
+            ),
+            EntropyAck,
+        )
+        ext_entropy = entropy_ack.entropy
+        # For SLIP-39 this is the Encrypted Master Secret
+        secret = _compute_secret_from_entropy(int_entropy, ext_entropy, msg.strength)
+
+        if backup_type == BAK_T_BIP39:
+            # in BIP-39 we store mnemonic string instead of the secret
+            secret = bip39.from_data(secret).encode()
+
+        if not msg.entropy_check or await _entropy_check(secret):
+            break
+
+        prev_int_entropy = int_entropy
 
     # If either of skip_backup or no_backup is specified, we are not doing backup now.
     # Otherwise, we try to do it.
@@ -112,7 +131,6 @@ async def reset_device(msg: ResetDevice) -> Success:
     storage_device.set_passphrase_enabled(bool(msg.passphrase_protection))
     storage_device.store_mnemonic_secret(
         secret,  # for SLIP-39, this is the EMS
-        backup_type,
         needs_backup=not perform_backup,
         no_backup=bool(msg.no_backup),
     )
@@ -124,6 +142,37 @@ async def reset_device(msg: ResetDevice) -> Success:
     return Success(message="Initialized")
 
 
+async def _entropy_check(secret: bytes) -> bool:
+    """Returns True to indicate that entropy check loop should end."""
+    from trezor.messages import EntropyCheckContinue, EntropyCheckReady, GetPublicKey
+    from trezor.wire.context import call_any
+
+    from apps.bitcoin.get_public_key import get_public_key
+    from apps.common import coininfo, paths
+    from apps.common.keychain import Keychain
+    from apps.common.mnemonic import get_seed
+
+    seed = get_seed(mnemonic_secret=secret)
+
+    msg = EntropyCheckReady()
+    while True:
+        req = await call_any(
+            msg,
+            MessageType.EntropyCheckContinue,
+            MessageType.GetPublicKey,
+        )
+        assert req.MESSAGE_WIRE_TYPE is not None
+
+        if EntropyCheckContinue.is_type_of(req):
+            return req.finish
+
+        assert GetPublicKey.is_type_of(req)
+        req.show_display = False
+        curve_name = req.ecdsa_curve_name or coininfo.by_name(req.coin_name).curve_name
+        keychain = Keychain(seed, curve_name, [paths.AlwaysMatchingSchema])
+        msg = await get_public_key(req, keychain=keychain)
+
+
 async def _backup_bip39(mnemonic: str) -> None:
     words = mnemonic.split()
     await layout.show_backup_intro(single_share=True, num_of_words=len(words))
@@ -272,7 +321,7 @@ def _validate_reset_device(msg: ResetDevice) -> None:
 
 
 def _compute_secret_from_entropy(
-    int_entropy: bytes, ext_entropy: bytes, strength_in_bytes: int
+    int_entropy: bytes, ext_entropy: bytes, strength_bits: int
 ) -> bytes:
     from trezor.crypto import hashlib
 
@@ -282,7 +331,7 @@ def _compute_secret_from_entropy(
     ehash.update(ext_entropy)
     entropy = ehash.digest()
     # take a required number of bytes
-    strength = strength_in_bytes // 8
+    strength = strength_bits // 8
     secret = entropy[:strength]
     return secret
 

core/src/storage/device.py

@@ -178,13 +178,11 @@ def set_homescreen(homescreen: bytes) -> None:
 
 def store_mnemonic_secret(
     secret: bytes,
-    backup_type: BackupType,
     needs_backup: bool = False,
     no_backup: bool = False,
 ) -> None:
     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: