tests: extract more common functionality

tests/device_tests/common.py

@@ -82,6 +82,16 @@ def generate_entropy(strength, internal_entropy, external_entropy):
 
 
 def recovery_enter_shares(debug, shares, groups=False):
+    """Perform the recovery flow for a set of Shamir shares.
+
+    For use in an input flow function.
+    Example:
+
+    def input_flow():
+        yield  # start recovery
+        client.debug.press_yes()
+        yield from recovery_enter_shares(client.debug, SOME_SHARES)
+    """
     word_count = len(shares[0].split(" "))
 
     # Homescreen - proceed to word number selection
@@ -118,3 +128,54 @@ def recovery_enter_shares(debug, shares, groups=False):
         # or Homescreen - confirm success
         yield
         debug.press_yes()
+
+
+def click_through(debug, screens, code=None):
+    """Click through N dialog screens.
+
+    For use in an input flow function.
+    Example:
+
+    def input_flow():
+        # 1. Confirm reset
+        # 2. Backup your seed
+        # 3. Confirm warning
+        # 4. Shares info
+        yield from click_through(client.debug, screens=4, code=B.ResetDevice)
+    """
+    for _ in range(screens):
+        received = yield
+        if code is not None:
+            assert received == code
+        debug.press_yes()
+
+
+def read_and_confirm_mnemonic(debug, words):
+    """Read a given number of mnemonic words from Trezor T screen and correctly
+    answer confirmation questions. Return the full mnemonic.
+
+    For use in an input flow function.
+    Example:
+
+    def input_flow():
+        yield from click_through(client.debug, screens=3)
+
+        yield  # confirm mnemonic entry
+        mnemonic = read_and_confirm_mnemonic(client.debug, words=20)
+    """
+    mnemonic = []
+    while True:
+        mnemonic.extend(debug.read_reset_word().split())
+        if len(mnemonic) < words:
+            debug.swipe_down()
+        else:
+            # last page is confirmation
+            debug.press_yes()
+            break
+
+    # check share
+    for _ in range(3):
+        index = debug.read_reset_word_pos()
+        debug.input(mnemonic[index])
+
+    return " ".join(mnemonic)

tests/device_tests/test_msg_resetdevice_shamir.py

@@ -8,13 +8,13 @@ from shamir_mnemonic import MnemonicError
 from trezorlib import device, messages as proto
 from trezorlib.messages import ButtonRequestType as B, ResetDeviceBackupType
 
-from .common import TrezorTest, generate_entropy
+from .common import click_through, generate_entropy, read_and_confirm_mnemonic
 
 EXTERNAL_ENTROPY = b"zlutoucky kun upel divoke ody" * 2
 
 
 @pytest.mark.skip_t1
-class TestMsgResetDeviceT2(TrezorTest):
+class TestMsgResetDeviceT2:
     # TODO: test with different options
     @pytest.mark.setup_client(uninitialized=True)
     def test_reset_device_shamir(self, client):
@@ -23,68 +23,23 @@ class TestMsgResetDeviceT2(TrezorTest):
         all_mnemonics = []
 
         def input_flow():
-            # Confirm Reset
+            # 1. Confirm Reset
-            btn_code = yield
+            # 2. Backup your seed
-            assert btn_code == B.ResetDevice
+            # 3. Confirm warning
-            client.debug.press_yes()
+            # 4. shares info
-
+            # 5. Set & Confirm number of shares
-            # Backup your seed
+            # 6. threshold info
-            btn_code = yield
+            # 7. Set & confirm threshold value
-            assert btn_code == B.ResetDevice
+            # 8. Confirm show seeds
-            client.debug.press_yes()
+            yield from click_through(client.debug, screens=8, code=B.ResetDevice)
-
-            # Confirm warning
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # shares info
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # Set & Confirm number of shares
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # threshold info
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # Set & confirm threshold value
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # Confirm show seeds
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
 
             # show & confirm shares
             for h in range(5):
-                words = []
+                # mnemonic phrases
                 btn_code = yield
                 assert btn_code == B.Other
-
+                mnemonic = read_and_confirm_mnemonic(client.debug, words=20)
-                # mnemonic phrases
+                all_mnemonics.append(mnemonic)
-                # 20 word over 6 pages for strength 128, 33 words over 9 pages for strength 256
-                for i in range(6):
-                    words.extend(client.debug.read_reset_word().split())
-                    if i < 5:
-                        client.debug.swipe_down()
-                    else:
-                        # last page is confirmation
-                        client.debug.press_yes()
-
-                # check share
-                for _ in range(3):
-                    index = client.debug.read_reset_word_pos()
-                    client.debug.input(words[index])
-
-                all_mnemonics.extend([" ".join(words)])
 
                 # Confirm continue to next share
                 btn_code = yield

tests/device_tests/test_msg_resetdevice_supershamir.py

@@ -6,95 +6,43 @@ import shamir_mnemonic as shamir
 from trezorlib import device, messages as proto
 from trezorlib.messages import ButtonRequestType as B, ResetDeviceBackupType
 
-from .common import TrezorTest, generate_entropy
+from .common import click_through, generate_entropy, read_and_confirm_mnemonic
 
 EXTERNAL_ENTROPY = b"zlutoucky kun upel divoke ody" * 2
 
 
 @pytest.mark.skip_t1
-class TestMsgResetDeviceT2(TrezorTest):
+class TestMsgResetDeviceT2:
     # TODO: test with different options
     @pytest.mark.setup_client(uninitialized=True)
     def test_reset_device_supershamir(self, client):
         strength = 128
+        word_count = 20
         member_threshold = 3
         all_mnemonics = []
 
         def input_flow():
-            # Confirm Reset
+            # 1. Confirm Reset
-            btn_code = yield
+            # 2. Backup your seed
-            assert btn_code == B.ResetDevice
+            # 3. Confirm warning
-            client.debug.press_yes()
+            # 4. shares info
-
+            # 5. Set & Confirm number of groups
-            # Backup your seed
+            # 6. threshold info
-            btn_code = yield
+            # 7. Set & confirm group threshold value
-            assert btn_code == B.ResetDevice
+            # 8-17: for each of 5 groups:
-            client.debug.press_yes()
+            #   1. Set & Confirm number of shares
-
+            #   2. Set & confirm share threshold value
-            # Confirm warning
+            # 18. Confirm show seeds
-            btn_code = yield
+            yield from click_through(client.debug, screens=18, code=B.ResetDevice)
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # shares info
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # Set & Confirm number of groups
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # threshold info
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # Set & confirm group threshold value
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            for _ in range(5):
-                # Set & Confirm number of share
-                btn_code = yield
-                assert btn_code == B.ResetDevice
-                client.debug.press_yes()
-
-                # Set & confirm share threshold value
-                btn_code = yield
-                assert btn_code == B.ResetDevice
-                client.debug.press_yes()
-
-            # Confirm show seeds
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
 
             # show & confirm shares for all groups
             for g in range(5):
                 for h in range(5):
-                    words = []
+                    # mnemonic phrases
                     btn_code = yield
                     assert btn_code == B.Other
-
+                    mnemonic = read_and_confirm_mnemonic(client.debug, words=word_count)
-                    # mnemonic phrases
+                    all_mnemonics.append(mnemonic)
-                    # 20 word over 6 pages for strength 128, 33 words over 9 pages for strength 256
-                    for i in range(6):
-                        words.extend(client.debug.read_reset_word().split())
-                        if i < 5:
-                            client.debug.swipe_down()
-                        else:
-                            # last page is confirmation
-                            client.debug.press_yes()
-
-                    # check share
-                    for _ in range(3):
-                        index = client.debug.read_reset_word_pos()
-                        client.debug.input(words[index])
-
-                    all_mnemonics.extend([" ".join(words)])
 
                     # Confirm continue to next share
                     btn_code = yield
@@ -119,25 +67,15 @@ class TestMsgResetDeviceT2(TrezorTest):
                     proto.ButtonRequest(code=B.ResetDevice),
                     proto.ButtonRequest(code=B.ResetDevice),
                     proto.ButtonRequest(code=B.ResetDevice),
-                    proto.ButtonRequest(
+                    proto.ButtonRequest(code=B.ResetDevice),  # group #1 counts
-                        code=B.ResetDevice
-                    ),  # group #1 shares& thresholds
                     proto.ButtonRequest(code=B.ResetDevice),
-                    proto.ButtonRequest(
+                    proto.ButtonRequest(code=B.ResetDevice),  # group #2 counts
-                        code=B.ResetDevice
-                    ),  # group #2 shares& thresholds
                     proto.ButtonRequest(code=B.ResetDevice),
-                    proto.ButtonRequest(
+                    proto.ButtonRequest(code=B.ResetDevice),  # group #3 counts
-                        code=B.ResetDevice
-                    ),  # group #3 shares& thresholds
                     proto.ButtonRequest(code=B.ResetDevice),
-                    proto.ButtonRequest(
+                    proto.ButtonRequest(code=B.ResetDevice),  # group #4 counts
-                        code=B.ResetDevice
-                    ),  # group #4 shares& thresholds
                     proto.ButtonRequest(code=B.ResetDevice),
-                    proto.ButtonRequest(
+                    proto.ButtonRequest(code=B.ResetDevice),  # group #5 counts
-                        code=B.ResetDevice
-                    ),  # group #5 shares& thresholds
                     proto.ButtonRequest(code=B.ResetDevice),
                     proto.ButtonRequest(code=B.Other),  # show seeds
                     proto.ButtonRequest(code=B.Success),

tests/device_tests/test_msg_resetdevice_t2.py

@@ -22,50 +22,34 @@ from mnemonic import Mnemonic
 from trezorlib import device, messages as proto
 from trezorlib.messages import ButtonRequestType as B
 
-from .common import MNEMONIC12, TrezorTest, generate_entropy
+from .common import (
+    MNEMONIC12,
+    click_through,
+    generate_entropy,
+    read_and_confirm_mnemonic,
+)
 
 EXTERNAL_ENTROPY = b"zlutoucky kun upel divoke ody" * 2
 
 
 @pytest.mark.skip_t1
-class TestMsgResetDeviceT2(TrezorTest):
+class TestMsgResetDeviceT2:
     @pytest.mark.setup_client(uninitialized=True)
     def test_reset_device(self, client):
-        words = []
+        mnemonic = None
         strength = 128
 
         def input_flow():
-            # Confirm Reset
+            nonlocal mnemonic
-            btn_code = yield
+            # 1. Confirm Reset
-            assert btn_code == B.ResetDevice
+            # 2. Backup your seed
-            client.debug.press_yes()
+            # 3. Confirm warning
-
+            yield from click_through(client.debug, screens=3, code=B.ResetDevice)
-            # Backup your seed
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # Confirm warning
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
 
             # mnemonic phrases
             btn_code = yield
             assert btn_code == B.ResetDevice
-            # 12 words, 3 pages
+            mnemonic = read_and_confirm_mnemonic(client.debug, words=12)
-            for i in range(3):
-                words.extend(client.debug.read_reset_word().split())
-                if i < 2:
-                    client.debug.swipe_down()
-                else:
-                    # last page is confirmation
-                    client.debug.press_yes()
-
-            # check backup words
-            for _ in range(3):
-                index = client.debug.read_reset_word_pos()
-                client.debug.input(words[index])
 
             # confirm recovery seed check
             btn_code = yield
@@ -111,7 +95,7 @@ class TestMsgResetDeviceT2(TrezorTest):
         expected_mnemonic = Mnemonic("english").to_mnemonic(entropy)
 
         # Compare that device generated proper mnemonic for given entropies
-        assert " ".join(words) == expected_mnemonic
+        assert mnemonic == expected_mnemonic
 
         # Check if device is properly initialized
         resp = client.call_raw(proto.Initialize())
@@ -122,10 +106,12 @@ class TestMsgResetDeviceT2(TrezorTest):
 
     @pytest.mark.setup_client(uninitialized=True)
     def test_reset_device_pin(self, client):
-        words = []
+        mnemonic = None
         strength = 128
 
         def input_flow():
+            nonlocal mnemonic
+
             # Confirm Reset
             btn_code = yield
             assert btn_code == B.ResetDevice
@@ -157,19 +143,7 @@ class TestMsgResetDeviceT2(TrezorTest):
             # mnemonic phrases
             btn_code = yield
             assert btn_code == B.ResetDevice
-            # 12 words, 3 pages
+            mnemonic = read_and_confirm_mnemonic(client.debug, words=12)
-            for i in range(3):
-                words.extend(client.debug.read_reset_word().split())
-                if i < 2:
-                    client.debug.swipe_down()
-                else:
-                    # last page is confirmation
-                    client.debug.press_yes()
-
-            # check backup words
-            for _ in range(3):
-                index = client.debug.read_reset_word_pos()
-                client.debug.input(words[index])
 
             # confirm recovery seed check
             btn_code = yield
@@ -218,7 +192,7 @@ class TestMsgResetDeviceT2(TrezorTest):
         expected_mnemonic = Mnemonic("english").to_mnemonic(entropy)
 
         # Compare that device generated proper mnemonic for given entropies
-        assert " ".join(words) == expected_mnemonic
+        assert mnemonic == expected_mnemonic
 
         # Check if device is properly initialized
         resp = client.call_raw(proto.Initialize())

tests/device_tests/test_shamir_reset_recovery.py

@@ -4,7 +4,7 @@ from trezorlib import btc, device, messages
 from trezorlib.messages import ButtonRequestType as B, ResetDeviceBackupType
 from trezorlib.tools import parse_path
 
-from .common import recovery_enter_shares
+from .common import click_through, read_and_confirm_mnemonic, recovery_enter_shares
 
 
 @pytest.mark.skip_t1
@@ -24,70 +24,28 @@ def test_reset_recovery(client):
 
 def reset(client, strength=128):
     all_mnemonics = []
+    # per SLIP-39: strength in bits, rounded up to nearest multiple of 10, plus 70 bits
+    # of metadata, split into 10-bit words
+    word_count = ((strength + 9) // 10) + 7
 
     def input_flow():
-        # Confirm Reset
+        # 1. Confirm Reset
-        btn_code = yield
+        # 2. Backup your seed
-        assert btn_code == B.ResetDevice
+        # 3. Confirm warning
-        client.debug.press_yes()
+        # 4. shares info
-
+        # 5. Set & Confirm number of shares
-        # Backup your seed
+        # 6. threshold info
-        btn_code = yield
+        # 7. Set & confirm threshold value
-        assert btn_code == B.ResetDevice
+        # 8. Confirm show seeds
-        client.debug.press_yes()
+        yield from click_through(client.debug, screens=8, code=B.ResetDevice)
-
-        # Confirm warning
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
-
-        # shares info
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
-
-        # Set & Confirm number of shares
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
-
-        # threshold info
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
-
-        # Set & confirm threshold value
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
-
-        # Confirm show seeds
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
 
         # show & confirm shares
         for h in range(5):
-            words = []
+            # mnemonic phrases
             btn_code = yield
             assert btn_code == B.Other
-
+            mnemonic = read_and_confirm_mnemonic(client.debug, words=word_count)
-            # mnemonic phrases
+            all_mnemonics.append(mnemonic)
-            # 20 word over 6 pages for strength 128, 33 words over 9 pages for strength 256
-            for i in range(6):
-                words.extend(client.debug.read_reset_word().split())
-                if i < 5:
-                    client.debug.swipe_down()
-                else:
-                    # last page is confirmation
-                    client.debug.press_yes()
-
-            # check share
-            for _ in range(3):
-                index = client.debug.read_reset_word_pos()
-                client.debug.input(words[index])
-
-            all_mnemonics.extend([" ".join(words)])
 
             # Confirm continue to next share
             btn_code = yield
@@ -139,6 +97,7 @@ def reset(client, strength=128):
             language="english",
             backup_type=ResetDeviceBackupType.Slip39_Single_Group,
         )
+
     client.set_input_flow(None)
 
     # Check if device is properly initialized

tests/device_tests/test_shamir_reset_recovery_groups.py

@@ -4,7 +4,7 @@ from trezorlib import btc, device, messages
 from trezorlib.messages import ButtonRequestType as B, ResetDeviceBackupType
 from trezorlib.tools import parse_path
 
-from .common import recovery_enter_shares
+from .common import click_through, read_and_confirm_mnemonic, recovery_enter_shares
 
 
 @pytest.mark.skip_t1
@@ -32,82 +32,32 @@ def test_reset_recovery(client):
 
 def reset(client, strength=128):
     all_mnemonics = []
+    # per SLIP-39: strength in bits, rounded up to nearest multiple of 10, plus 70 bits
+    # of metadata, split into 10-bit words
+    word_count = ((strength + 9) // 10) + 7
 
     def input_flow():
-        # Confirm Reset
+        # 1. Confirm Reset
-        btn_code = yield
+        # 2. Backup your seed
-        assert btn_code == B.ResetDevice
+        # 3. Confirm warning
-        client.debug.press_yes()
+        # 4. shares info
-
+        # 5. Set & Confirm number of groups
-        # Backup your seed
+        # 6. threshold info
-        btn_code = yield
+        # 7. Set & confirm group threshold value
-        assert btn_code == B.ResetDevice
+        # 8-17: for each of 5 groups:
-        client.debug.press_yes()
+        #   1. Set & Confirm number of shares
-
+        #   2. Set & confirm share threshold value
-        # Confirm warning
+        # 18. Confirm show seeds
-        btn_code = yield
+        yield from click_through(client.debug, screens=18, code=B.ResetDevice)
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
-
-        # shares info
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
-
-        # Set & Confirm number of groups
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
-
-        # threshold info
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
-
-        # Set & confirm group threshold value
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
-
-        for _ in range(5):
-            # Set & Confirm number of share
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-            # Set & confirm share threshold value
-            btn_code = yield
-            assert btn_code == B.ResetDevice
-            client.debug.press_yes()
-
-        # Confirm show seeds
-        btn_code = yield
-        assert btn_code == B.ResetDevice
-        client.debug.press_yes()
 
         # show & confirm shares for all groups
         for g in range(5):
             for h in range(5):
-                words = []
+                # mnemonic phrases
                 btn_code = yield
                 assert btn_code == B.Other
-
+                mnemonic = read_and_confirm_mnemonic(client.debug, words=word_count)
-                # mnemonic phrases
+                all_mnemonics.append(mnemonic)
-                # 20 word over 6 pages for strength 128, 33 words over 9 pages for strength 256
-                for i in range(6):
-                    words.extend(client.debug.read_reset_word().split())
-                    if i < 5:
-                        client.debug.swipe_down()
-                    else:
-                        # last page is confirmation
-                        client.debug.press_yes()
-
-                # check share
-                for _ in range(3):
-                    index = client.debug.read_reset_word_pos()
-                    client.debug.input(words[index])
-
-                all_mnemonics.extend([" ".join(words)])
 
                 # Confirm continue to next share
                 btn_code = yield
@@ -131,25 +81,15 @@ def reset(client, strength=128):
                 messages.ButtonRequest(code=B.ResetDevice),
                 messages.ButtonRequest(code=B.ResetDevice),
                 messages.ButtonRequest(code=B.ResetDevice),
-                messages.ButtonRequest(
+                messages.ButtonRequest(code=B.ResetDevice),  # group #1 counts
-                    code=B.ResetDevice
-                ),  # group #1 shares& thresholds
                 messages.ButtonRequest(code=B.ResetDevice),
-                messages.ButtonRequest(
+                messages.ButtonRequest(code=B.ResetDevice),  # group #2 counts
-                    code=B.ResetDevice
-                ),  # group #2 shares& thresholds
                 messages.ButtonRequest(code=B.ResetDevice),
-                messages.ButtonRequest(
+                messages.ButtonRequest(code=B.ResetDevice),  # group #3 counts
-                    code=B.ResetDevice
-                ),  # group #3 shares& thresholds
                 messages.ButtonRequest(code=B.ResetDevice),
-                messages.ButtonRequest(
+                messages.ButtonRequest(code=B.ResetDevice),  # group #4 counts
-                    code=B.ResetDevice
-                ),  # group #4 shares& thresholds
                 messages.ButtonRequest(code=B.ResetDevice),
-                messages.ButtonRequest(
+                messages.ButtonRequest(code=B.ResetDevice),  # group #5 counts
-                    code=B.ResetDevice
-                ),  # group #5 shares& thresholds
                 messages.ButtonRequest(code=B.ResetDevice),
                 messages.ButtonRequest(code=B.Other),  # show seeds
                 messages.ButtonRequest(code=B.Success),