fixup! feat(core): add N1W1 backup/recovery demo flows

Author: romanz

tests/device_tests/reset_recovery/__init__.py

@@ -0,0 +1,13 @@
+import pytest
+
+from trezorlib.debuglink import DebugSession as Session
+from trezorlib.messages import BackupMethod, Capability
+
+
+
+def check_support(session: Session, method: BackupMethod) -> None:
+    if (
+        method is BackupMethod.N4W1
+        and Capability.N4W1 not in session.features.capabilities
+    ):
+        pytest.skip("N4W1 not supported")

tests/device_tests/reset_recovery/test_recovery_slip39_basic.py

@@ -19,6 +19,7 @@
 from trezorlib import device, exceptions, messages
 from trezorlib.debuglink import DebugSession as Session
 
+from . import  check_support
 from ...common import (
     MNEMONIC_SLIP39_BASIC_20_3of6,
     MNEMONIC_SLIP39_BASIC_20_3of6_SECRET,
@@ -75,9 +76,9 @@ def test_secret(
     session: Session, shares: list[str], secret: str, backup_type: messages.BackupType
 ):
     with session.test_ctx as client:
-        IF = InputFlowSlip39BasicRecovery(session, shares)
+        IF = InputFlowSlip39BasicRecovery(session, shares, method=messages.BackupMethod.N4W1)
         client.set_input_flow(IF.get())
-        device.recover(session, pin_protection=False, label="label")
+        device.recover(session, pin_protection=False, label="label", backup_method=messages.BackupMethod.N4W1)
 
     # Workflow successfully ended
     assert session.features.pin_protection is False

tests/device_tests/reset_recovery/test_reset_recovery_slip39_basic.py

@@ -22,21 +22,32 @@
 from trezorlib import btc, device, messages
 from trezorlib.debuglink import DebugSession as Session
 from trezorlib.debuglink import TrezorTestContext as Client
-from trezorlib.messages import BackupType
+from trezorlib.messages import BackupMethod, BackupType
 from trezorlib.tools import parse_path
 
+from . import check_support
 from ...common import MOCK_GET_ENTROPY
 from ...input_flows import (
     InputFlowSlip39BasicRecovery,
     InputFlowSlip39BasicResetRecovery,
 )
 from ...translations import set_language
 
+METHODS = [
+    BackupMethod.Display,
+    BackupMethod.N4W1,
+]
+
+TEST_PARAMS = [pytest.param(method, id=method.name) for method in METHODS]
+
 
 @pytest.mark.models("core")
 @pytest.mark.setup_client(uninitialized=True)
-def test_reset_recovery(client: Client):
+@pytest.mark.parametrize("method", TEST_PARAMS)
+def test_reset_recovery(client: Client, method: BackupMethod):
     session = client.get_seedless_session()
+    check_support(session, method)
+
     mnemonics = reset(session)
     session = client.get_session()
     address_before = btc.get_address(session, "Bitcoin", parse_path("m/44h/0h/0h/0/0"))
@@ -48,7 +59,7 @@ def test_reset_recovery(client: Client):
         session = client.get_seedless_session()
         set_language(session, lang[:2])
         selected_mnemonics = share_subset
-        recover(session, selected_mnemonics)
+        recover(session, selected_mnemonics, method)
         session = client.get_session()
         address_after = btc.get_address(
             session, "Bitcoin", parse_path("m/44h/0h/0h/0/0")
@@ -85,11 +96,13 @@ def reset(session: Session, strength: int = 128) -> list[str]:
     return IF.mnemonics
 
 
-def recover(session: Session, shares: t.Sequence[str]):
+def recover(session: Session, shares: t.Sequence[str], method: BackupMethod):
     with session.test_ctx as client:
-        IF = InputFlowSlip39BasicRecovery(session, shares)
+        IF = InputFlowSlip39BasicRecovery(session, shares, method=method)
         client.set_input_flow(IF.get())
-        device.recover(session, pin_protection=False, label="label")
+        device.recover(
+            session, pin_protection=False, label="label", backup_method=method
+        )
 
     # Workflow successfully ended
     assert session.features.pin_protection is False

tests/device_tests/reset_recovery/test_reset_slip39_basic.py

@@ -14,7 +14,7 @@
 # You should have received a copy of the License along with this library.
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
-from itertools import combinations
+from itertools import combinations, product
 
 import pytest
 from shamir_mnemonic import MnemonicError, shamir
@@ -23,8 +23,9 @@
 from trezorlib.btc import get_public_node
 from trezorlib.debuglink import DebugSession as Session
 from trezorlib.exceptions import TrezorFailure
-from trezorlib.messages import BackupAvailability, BackupType
+from trezorlib.messages import BackupAvailability, BackupMethod, BackupType
 
+from . import check_support
 from ...common import EXTERNAL_ENTROPY, MOCK_GET_ENTROPY, generate_entropy
 from ...input_flows import (
     FlowAdapter,
@@ -40,12 +41,21 @@
     try_to_cancel({"backup_device", "setup_device", "success_backup"}),
 ]
 
+METHODS = [
+    BackupMethod.Display,
+    BackupMethod.N4W1,
+]
+
+
+def reset_device(
+    session: Session, strength: int, adapt_flow: FlowAdapter, method: BackupMethod
+):
+    check_support(session, method)
 
-def reset_device(session: Session, strength: int, adapt_flow: FlowAdapter):
     member_threshold = 3
 
     with session.test_ctx as client:
-        IF = InputFlowSlip39BasicResetRecovery(session)
+        IF = InputFlowSlip39BasicResetRecovery(session, method)
         client.set_input_flow(adapt_flow(session, IF.get()))
 
         # No PIN, no passphrase, don't display random
@@ -56,6 +66,7 @@ def reset_device(session: Session, strength: int, adapt_flow: FlowAdapter):
             pin_protection=False,
             label="test",
             backup_type=BackupType.Slip39_Basic,
+            backup_method=method,
             entropy_check_count=0,
             _get_entropy=MOCK_GET_ENTROPY,
         )
@@ -80,27 +91,41 @@ def reset_device(session: Session, strength: int, adapt_flow: FlowAdapter):
         device.backup(session)
 
 
+TEST_PARAMS = [
+    pytest.param(adapt_flow, method, id=f"{adapt_flow.__name__}-{method.name}")
+    for adapt_flow, method in product(FLOW_ADAPTERS, METHODS)
+]
+
+
 @pytest.mark.setup_client(uninitialized=True)
-@pytest.mark.parametrize("adapt_flow", FLOW_ADAPTERS, ids=lambda f: f.__name__)
-def test_reset_device_slip39_basic(session: Session, adapt_flow: FlowAdapter):
-    reset_device(session, 128, adapt_flow)
+@pytest.mark.parametrize("adapt_flow,method", TEST_PARAMS)
+def test_reset_device_slip39_basic(
+    session: Session, adapt_flow: FlowAdapter, method: BackupMethod
+):
+    reset_device(session, 128, adapt_flow, method=method)
 
 
-@pytest.mark.parametrize("adapt_flow", FLOW_ADAPTERS, ids=lambda f: f.__name__)
+@pytest.mark.parametrize("adapt_flow,method", TEST_PARAMS)
 @pytest.mark.setup_client(uninitialized=True)
-def test_reset_device_slip39_basic_256(session: Session, adapt_flow: FlowAdapter):
-    reset_device(session, 256, adapt_flow)
+def test_reset_device_slip39_basic_256(
+    session: Session, adapt_flow: FlowAdapter, method: BackupMethod
+):
+    reset_device(session, 256, adapt_flow, method=method)
 
 
 @pytest.mark.setup_client(uninitialized=True)
-@pytest.mark.parametrize("adapt_flow", FLOW_ADAPTERS, ids=lambda f: f.__name__)
-def test_reset_entropy_check(session: Session, adapt_flow: FlowAdapter):
+@pytest.mark.parametrize("adapt_flow,method", TEST_PARAMS)
+def test_reset_entropy_check(
+    session: Session, adapt_flow: FlowAdapter, method: BackupMethod
+):
+    check_support(session, method)
+
     member_threshold = 3
 
     strength = 128  # 20 words
 
     with session.test_ctx as client:
-        IF = InputFlowSlip39BasicResetRecovery(session)
+        IF = InputFlowSlip39BasicResetRecovery(session, method)
         client.set_input_flow(adapt_flow(session, IF.get()))
 
         # No PIN, no passphrase.
@@ -111,6 +136,7 @@ def test_reset_entropy_check(session: Session, adapt_flow: FlowAdapter):
             pin_protection=False,
             label="test",
             backup_type=BackupType.Slip39_Basic,
+            backup_method=method,
             entropy_check_count=3,
             _get_entropy=MOCK_GET_ENTROPY,
         )

tests/input_flows.py

@@ -34,7 +34,13 @@
     read_and_confirm_mnemonic,
     swipe_if_necessary,
 )
-from .input_flows_helpers import BackupFlow, EthereumFlow, PinFlow, RecoveryFlow
+from .input_flows_helpers import (
+    BackupFlow,
+    EthereumFlow,
+    PinFlow,
+    RecoveryFlow,
+    n4w1_handle_write,
+)
 
 B = messages.ButtonRequestType
 
@@ -1958,18 +1964,24 @@ def input_flow_common(self) -> BRGeneratorType:
 def load_N_shares(
     debug: DebugLink,
     n: int,
+    method: messages.BackupMethod = messages.BackupMethod.Display,
 ) -> Generator[None, "messages.ButtonRequest", list[str]]:
     mnemonics: list[str] = []
 
     for _ in range(n):
-        # Phrase screen
-        mnemonic = yield from read_and_confirm_mnemonic(debug)
-        assert mnemonic is not None
-        mnemonics.append(mnemonic)
-
-        br = yield  # Confirm continue to next
-        assert br.code == B.Success
-        debug.press_yes()
+        if method is messages.BackupMethod.Display:
+            # Phrase screen
+            mnemonic = yield from read_and_confirm_mnemonic(debug)
+            assert mnemonic is not None
+            mnemonics.append(mnemonic)
+            br = yield  # Confirm continue to next
+            assert br.code == B.Success
+            debug.press_yes()
+        elif method is messages.BackupMethod.N4W1:
+            assert (yield).name == "backup_write"
+            mnemonics.append(n4w1_handle_write(debug).decode())
+        else:
+            raise RuntimeError
 
     return mnemonics
 
@@ -2114,11 +2126,17 @@ def input_flow_eckhart(self) -> BRGeneratorType:
 
 
 class InputFlowSlip39BasicResetRecovery(InputFlowBase):
-    def __init__(self, client: Client | DebugSession):
+    def __init__(
+        self,
+        client: Client | DebugSession,
+        method: messages.BackupMethod = messages.BackupMethod.Display,
+    ):
         super().__init__(client)
         self.mnemonics: list[str] = []
+        self.method = method
 
     def input_flow_bolt(self) -> BRGeneratorType:
+        assert self.method is messages.BackupMethod.Display
         # 1. Confirm Reset
         # 2. Backup your seed
         # 3. Backup intro
@@ -2138,6 +2156,7 @@ def input_flow_bolt(self) -> BRGeneratorType:
         self.debug.press_yes()
 
     def input_flow_caesar(self) -> BRGeneratorType:
+        assert self.method is messages.BackupMethod.Display
         yield  # Confirm Reset
         self.debug.press_yes()
         yield  # Backup your seed
@@ -2169,6 +2188,7 @@ def input_flow_caesar(self) -> BRGeneratorType:
         self.debug.press_yes()
 
     def input_flow_delizia(self) -> BRGeneratorType:
+        assert self.method is messages.BackupMethod.Display
         # 1. Confirm Reset
         # 2. Wallet Created
         # 3. Backup your seed
@@ -2189,6 +2209,10 @@ def input_flow_delizia(self) -> BRGeneratorType:
         self.debug.press_yes()
 
     def input_flow_eckhart(self) -> BRGeneratorType:
+        num_screens = {
+            messages.BackupMethod.Display: 10,
+            messages.BackupMethod.N4W1: 8,
+        }[self.method]
         # 1. Confirm Reset
         # 2. Wallet Created
         # 3. Backup your seed
@@ -2199,10 +2223,10 @@ def input_flow_eckhart(self) -> BRGeneratorType:
         # 8. Confirm show seeds
         # 9. Warning
         # 10. Instructions
-        yield from click_through(self.debug, screens=10, code=B.ResetDevice)
+        yield from click_through(self.debug, screens=num_screens, code=B.ResetDevice)
 
         # Mnemonic phrases
-        self.mnemonics = yield from load_N_shares(self.debug, 5)
+        self.mnemonics = yield from load_N_shares(self.debug, 5, self.method)
 
         br = yield  # success screen
         assert br.code == B.Success
@@ -2800,18 +2824,22 @@ def __init__(
         client: Client | DebugSession,
         shares: Sequence[str],
         pin: str | None = None,
+        method: messages.BackupMethod = messages.BackupMethod.Display,
     ):
         super().__init__(client)
         self.shares = shares
         self.pin = pin
         self.word_count = len(shares[0].split(" "))
+        self.method = method
 
     def input_flow_common(self) -> BRGeneratorType:
         yield from self.REC.confirm_recovery()
         if self.pin is not None:
             yield from self.PIN.setup_new_pin(self.pin)
-        yield from self.REC.setup_slip39_recovery(self.word_count)
-        yield from self.REC.input_all_slip39_shares(self.shares)
+        if self.method is messages.BackupMethod.Display:
+            yield from self.REC.setup_slip39_recovery(self.word_count)
+
+        yield from self.REC.input_all_slip39_shares(self.shares, method=self.method)
         yield from self.REC.success_wallet_recovered()