feat(python): Implement entropy check workflow in device.reset().

python/.changelog.d/4155.added

@@ -0,0 +1 @@
+Added support for entropy check workflow in device.reset().

python/src/trezorlib/cli/__init__.py

@@ -158,7 +158,7 @@ def trezorctl_command_with_client(
 
     # the return type of @click.pass_obj is improperly specified and pyright doesn't
     # understand that it converts f(obj, *args, **kwargs) to f(*args, **kwargs)
-    return trezorctl_command_with_client  # type: ignore [is incompatible with return type]
+    return trezorctl_command_with_client
 
 
 class AliasedGroup(click.Group):

python/src/trezorlib/cli/device.py

@@ -24,6 +24,7 @@
 import requests
 
 from .. import debuglink, device, exceptions, messages, ui
+from ..tools import format_path
 from . import ChoiceType, with_client
 
 if t.TYPE_CHECKING:
@@ -214,6 +215,7 @@ def recover(
 @click.option("-s", "--skip-backup", is_flag=True)
 @click.option("-n", "--no-backup", is_flag=True)
 @click.option("-b", "--backup-type", type=ChoiceType(BACKUP_TYPE))
+@click.option("-e", "--entropy-check-count", type=click.IntRange(0))
 @with_client
 def setup(
     client: "TrezorClient",
@@ -225,6 +227,7 @@ def setup(
     skip_backup: bool,
     no_backup: bool,
     backup_type: messages.BackupType | None,
+    entropy_check_count: int | None,
 ) -> str:
     """Perform device setup and generate new seed."""
     if strength:
@@ -253,7 +256,7 @@ def setup(
             "backup type. Traditional BIP39 backup may be generated instead."
         )
 
-    return device.reset(
+    resp, path_xpubs = device.reset_entropy_check(
         client,
         strength=strength,
         passphrase_protection=passphrase_protection,
@@ -263,8 +266,17 @@ def setup(
         skip_backup=skip_backup,
         no_backup=no_backup,
         backup_type=backup_type,
+        entropy_check_count=entropy_check_count,
     )
 
+    if isinstance(resp, messages.Success):
+        click.echo("XPUBs for the generated seed")
+        for path, xpub in path_xpubs:
+            click.echo(f"{format_path(path)}: {xpub}")
+        return resp.message or ""
+    else:
+        raise RuntimeError(f"Received {resp.__class__}")
+
 
 @cli.command()
 @click.option("-t", "--group-threshold", type=int)

python/src/trezorlib/device.py

@@ -16,14 +16,18 @@
 
 from __future__ import annotations
 
+import hashlib
+import hmac
 import os
 import time
 import warnings
-from typing import TYPE_CHECKING, Callable, Iterable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Iterable, List, Optional, Tuple
+
+from slip10 import SLIP10
 
 from . import messages
 from .exceptions import Cancelled, TrezorException
-from .tools import Address, expect, session
+from .tools import Address, expect, parse_path, session
 
 if TYPE_CHECKING:
     from .client import TrezorClient
@@ -229,9 +233,53 @@ def recover(
     return res
 
 
+def is_slip39_backup_type(backup_type: messages.BackupType):
+    return backup_type in (
+        messages.BackupType.Slip39_Basic,
+        messages.BackupType.Slip39_Advanced,
+        messages.BackupType.Slip39_Single_Extendable,
+        messages.BackupType.Slip39_Basic_Extendable,
+        messages.BackupType.Slip39_Advanced_Extendable,
+    )
+
+
+def _seed_from_entropy(
+    internal_entropy: bytes,
+    external_entropy: bytes,
+    strength: int,
+    backup_type: messages.BackupType,
+) -> bytes:
+    entropy = hashlib.sha256(internal_entropy + external_entropy).digest()
+    secret = entropy[: strength // 8]
+
+    if len(secret) * 8 != strength:
+        raise ValueError("Entropy length mismatch")
+
+    if backup_type == messages.BackupType.Bip39:
+        import mnemonic
+
+        bip39 = mnemonic.Mnemonic("english")
+        words = bip39.to_mnemonic(secret)
+        seed = bip39.to_seed(words, passphrase="")
+    elif is_slip39_backup_type(backup_type):
+        import shamir_mnemonic
+
+        seed = shamir_mnemonic.cipher.decrypt(
+            secret, b"", iteration_exponent=1, identifier=0, extendable=True
+        )
+    else:
+        raise ValueError("Unknown backup type.")
+
+    return seed
+
+
 @expect(messages.Success, field="message", ret_type=str)
+def reset(*args: Any, **kwargs: Any) -> "MessageType":
+    return reset_entropy_check(*args, **kwargs)[0]
+
+
 @session
-def reset(
+def reset_entropy_check(
     client: "TrezorClient",
     display_random: bool = False,
     strength: Optional[int] = None,
@@ -243,7 +291,9 @@ def reset(
     skip_backup: bool = False,
     no_backup: bool = False,
     backup_type: messages.BackupType = messages.BackupType.Bip39,
-) -> "MessageType":
+    entropy_check_count: Optional[int] = None,
+    paths: List[Address] = [],
+) -> Tuple["MessageType", Iterable[Tuple[Address, str]]]:
     if display_random:
         warnings.warn(
             "display_random ignored. The feature is deprecated.",
@@ -267,6 +317,10 @@ def reset(
         else:
             strength = 128
 
+    if not paths:
+        # Get XPUBs for the first BTC SegWit v0 account and first ETH account.
+        paths = [parse_path("m/84h/0h/0h"), parse_path("m/44h/60h/0h")]
+
     # Begin with device reset workflow
     msg = messages.ResetDevice(
         strength=strength,
@@ -277,17 +331,61 @@ def reset(
         skip_backup=bool(skip_backup),
         no_backup=bool(no_backup),
         backup_type=backup_type,
+        entropy_check=entropy_check_count is not None,
     )
 
     resp = client.call(msg)
     if not isinstance(resp, messages.EntropyRequest):
         raise RuntimeError("Invalid response, expected EntropyRequest")
 
-    external_entropy = os.urandom(32)
-    # LOG.debug("Computer generated entropy: " + external_entropy.hex())
-    ret = client.call(messages.EntropyAck(entropy=external_entropy))
+    while True:
+        xpubs = []
+
+        external_entropy = os.urandom(32)
+        entropy_commitment = resp.entropy_commitment
+        resp = client.call(messages.EntropyAck(entropy=external_entropy))
+
+        if entropy_check_count is None:
+            break
+
+        if not isinstance(resp, messages.Success):
+            return resp, []
+
+        for path in paths:
+            resp = client.call(messages.GetPublicKey(address_n=path))
+            if not isinstance(resp, messages.PublicKey):
+                return resp, []
+            xpubs.append(resp.xpub)
+
+        if entropy_check_count <= 0:
+            resp = client.call(messages.ResetDeviceFinish())
+            break
+
+        entropy_check_count -= 1
+
+        resp = client.call(messages.ResetDeviceContinue())
+        if not isinstance(resp, messages.EntropyRequest):
+            raise RuntimeError("Invalid response, expected EntropyRequest")
+
+        # Check the entropy commitment from the previous round.
+        assert resp.prev_entropy
+        if (
+            hmac.HMAC(key=resp.prev_entropy, msg=b"", digestmod=hashlib.sha256).digest()
+            != entropy_commitment
+        ):
+            raise RuntimeError("Invalid entropy commitment.")
+
+        # Derive the seed and check that XPUBs match.
+        seed = _seed_from_entropy(
+            resp.prev_entropy, external_entropy, strength, backup_type
+        )
+        slip10 = SLIP10.from_seed(seed)
+        for path, xpub in zip(paths, xpubs):
+            if slip10.get_xpub_from_path(path) != xpub:
+                raise RuntimeError("Invalid XPUB in entropy check")
+
     client.init_device()
-    return ret
+    return resp, zip(paths, xpubs)
 
 
 @expect(messages.Success, field="message", ret_type=str)

python/src/trezorlib/tools.py

@@ -222,6 +222,23 @@ def str_to_harden(x: str) -> int:
         raise ValueError("Invalid BIP32 path", nstr) from e
 
 
+def format_path(path: Address, flag: str = "h") -> str:
+    """
+    Convert BIP32 path list of uint32 integers with hardened flags to string.
+    Several conventions are supported to denote the hardened flag: 1', 1h
+
+    e.g.: [0, 0x80000001, 1] -> "m/0/1h/1"
+
+    :param path: list of integers
+    :return: path string
+    """
+    nstr = "m"
+    for i in path:
+        nstr += f"/{unharden(i)}{flag if is_hardened(i) else ''}"
+
+    return nstr
+
+
 def prepare_message_bytes(txt: AnyStr) -> bytes:
     """
     Make message suitable for protobuf.