Merge pull request #651 from trezor/tsusanka/upgrade-shamir-test

Add Shamir recovery test to upgrade tests

tests/click_tests/recovery.py

@@ -0,0 +1,58 @@
+from .. import buttons
+
+
+def enter_word(debug, word):
+    word = word[:4]
+    for coords in buttons.type_word(word):
+        debug.click(coords)
+    return debug.click(buttons.CONFIRM_WORD, wait=True)
+
+
+def select_number_of_words(debug, num_of_words=20):
+    # confirm recovery
+    layout = debug.wait_layout()
+    assert layout.text.startswith("Recovery mode")
+    layout = debug.click(buttons.OK, wait=True)
+
+    # select number of words
+    assert "Select number of words" in layout.text
+    layout = debug.click(buttons.OK, wait=True)
+    assert layout.text == "WordSelector"
+
+    # click the number
+    word_option_offset = 6
+    word_options = (12, 18, 20, 24, 33)
+    index = word_option_offset + word_options.index(
+        num_of_words
+    )  # raises if num of words is invalid
+    coords = buttons.grid34(index % 3, index // 3)
+    layout = debug.click(coords, wait=True)
+    assert "Enter any share" in layout.text
+
+
+def enter_share(debug, share: str):
+    layout = debug.click(buttons.OK, wait=True)
+
+    assert layout.text == "Slip39Keyboard"
+    for word in share.split(" "):
+        layout = enter_word(debug, word)
+
+    return layout
+
+
+def enter_shares(debug, shares: list):
+    layout = debug.read_layout()
+    expected_text = "Enter any share"
+    remaining = len(shares)
+    for share in shares:
+        assert expected_text in layout.text
+        layout = enter_share(debug, share)
+        remaining -= 1
+        expected_text = "RecoveryHomescreen {} more".format(remaining)
+
+    assert "You have successfully recovered your wallet" in layout.text
+
+
+def finalize(debug):
+    layout = debug.click(buttons.OK, wait=True)
+    assert layout.text == "Homescreen"

tests/click_tests/test_recovery.py

@@ -18,15 +18,8 @@ import pytest
 
 from trezorlib import device, messages
 
-from .. import buttons
 from ..common import MNEMONIC_SLIP39_BASIC_20_3of6
-
+from . import recovery
-
-def enter_word(debug, word):
-    word = word[:4]
-    for coords in buttons.type_word(word):
-        debug.click(coords)
-    return debug.click(buttons.CONFIRM_WORD, wait=True)
 
 
 @pytest.mark.skip_t1
@@ -38,37 +31,9 @@ def test_recovery(device_handler):
     assert features.initialized is False
     device_handler.run(device.recover, pin_protection=False)
 
-    # select number of words
+    recovery.select_number_of_words(debug)
-    layout = debug.wait_layout()
+    recovery.enter_shares(debug, MNEMONIC_SLIP39_BASIC_20_3of6)
-    assert layout.text.startswith("Recovery mode")
+    recovery.finalize(debug)
-    layout = debug.click(buttons.OK, wait=True)
-
-    assert "Select number of words" in layout.text
-    layout = debug.click(buttons.OK, wait=True)
-
-    assert layout.text == "WordSelector"
-    # click "20" at 2, 2
-    coords = buttons.grid34(2, 2)
-    layout = debug.click(coords, wait=True)
-    assert "Enter any share" in layout.text
-
-    expected_text = "Enter any share (20 words)"
-    remaining = len(MNEMONIC_SLIP39_BASIC_20_3of6)
-    for share in MNEMONIC_SLIP39_BASIC_20_3of6:
-        assert expected_text in layout.text
-        layout = debug.click(buttons.OK, wait=True)
-
-        assert layout.text == "Slip39Keyboard"
-        for word in share.split(" "):
-            layout = enter_word(debug, word)
-
-        remaining -= 1
-        expected_text = "RecoveryHomescreen {} more".format(remaining)
-
-    assert "You have successfully recovered your wallet" in layout.text
-    layout = debug.click(buttons.OK, wait=True)
-
-    assert layout.text == "Homescreen"
 
     assert isinstance(device_handler.result(), messages.Success)
     features = device_handler.features()

tests/common.py

@@ -24,6 +24,7 @@ MNEMONIC_SLIP39_BASIC_20_3of6 = [
     "extra extend academic acne away best indicate impact square oasis prospect painting voting guest either argue username racism enemy eclipse",
     "extra extend academic arcade born dive legal hush gross briefing talent drug much home firefly toxic analysis idea umbrella slice",
 ]
+MNEMONIC_SLIP39_BASIC_20_3of6_SECRET = "491b795b80fc21ccdf466c0fbc98c8fc"
 # Shamir shares (128 bits, 2 groups from 1 of 1, 1 of 1, 3 of 5, 2 of 6)
 MNEMONIC_SLIP39_ADVANCED_20 = [
     "eraser senior beard romp adorn nuclear spill corner cradle style ancient family general leader ambition exchange unusual garlic promise voice",

tests/device_tests/test_msg_recoverydevice_slip39_basic.py

@@ -18,7 +18,11 @@ import pytest
 
 from trezorlib import device, exceptions, messages
 
-from ..common import MNEMONIC_SLIP39_BASIC_20_3of6, recovery_enter_shares
+from ..common import (
+    MNEMONIC_SLIP39_BASIC_20_3of6,
+    MNEMONIC_SLIP39_BASIC_20_3of6_SECRET,
+    recovery_enter_shares,
+)
 
 pytestmark = pytest.mark.skip_t1
 
@@ -33,7 +37,7 @@ MNEMONIC_SLIP39_BASIC_33_2of5 = [
 ]
 
 VECTORS = (
-    (MNEMONIC_SLIP39_BASIC_20_3of6, "491b795b80fc21ccdf466c0fbc98c8fc"),
+    (MNEMONIC_SLIP39_BASIC_20_3of6, MNEMONIC_SLIP39_BASIC_20_3of6_SECRET),
     (
         MNEMONIC_SLIP39_BASIC_33_2of5,
         "b770e0da1363247652de97a39bdbf2463be087848d709ecbf28e84508e31202a",

tests/upgrade_tests/test_firmware_upgrades.py

@@ -17,8 +17,12 @@
 import pytest
 
 from trezorlib import MINIMUM_FIRMWARE_VERSION, btc, debuglink, device
+from trezorlib.messages import BackupType
 from trezorlib.tools import H_
 
+from ..click_tests import recovery
+from ..common import MNEMONIC_SLIP39_BASIC_20_3of6, MNEMONIC_SLIP39_BASIC_20_3of6_SECRET
+from ..device_handler import BackgroundDeviceHandler
 from ..emulators import ALL_TAGS, EmulatorWrapper
 from . import SELECTED_GENS
 
@@ -194,6 +198,47 @@ def test_upgrade_reset_no_backup(gen, from_tag, to_tag):
         assert btc.get_address(emu.client, "Bitcoin", PATH) == address
 
 
+# Although Shamir was introduced in 2.1.2 already, the debug instrumentation was not present until 2.1.9.
+@for_all("core", minimum_version=(2, 1, 9))
+def test_upgrade_shamir_recovery(gen, from_tag, to_tag):
+    with EmulatorWrapper(gen, from_tag) as emu, BackgroundDeviceHandler(
+        emu.client
+    ) as device_handler:
+        assert emu.client.features.recovery_mode is False
+        debug = device_handler.debuglink()
+
+        device_handler.run(device.recover, pin_protection=False)
+
+        recovery.select_number_of_words(debug)
+        layout = recovery.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[0])
+        assert "2 more shares" in layout.text
+
+        device_id = emu.client.features.device_id
+        storage = emu.storage()
+        device_handler.check_finalize()
+
+    with EmulatorWrapper(gen, to_tag, storage=storage) as emu, BackgroundDeviceHandler(
+        emu.client
+    ) as device_handler:
+        assert device_id == emu.client.features.device_id
+        assert emu.client.features.recovery_mode
+        debug = device_handler.debuglink()
+
+        # second share
+        layout = recovery.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[2])
+        assert "1 more share" in layout.text
+
+        # last one
+        layout = recovery.enter_share(debug, MNEMONIC_SLIP39_BASIC_20_3of6[1])
+        assert "You have successfully" in layout.text
+
+        # Check the result
+        state = debug.state()
+        assert state.mnemonic_secret.hex() == MNEMONIC_SLIP39_BASIC_20_3of6_SECRET
+        assert state.mnemonic_type == BackupType.Slip39_Basic
+        device_handler.check_finalize()
+
+
 if __name__ == "__main__":
     if not ALL_TAGS:
         print("No versions found. Remember to run download_emulators.sh")