diff --git a/core/src/apps/management/recovery_device/word_validity.py b/core/src/apps/management/recovery_device/word_validity.py index 915cb0323..10950e0a6 100644 --- a/core/src/apps/management/recovery_device/word_validity.py +++ b/core/src/apps/management/recovery_device/word_validity.py @@ -16,10 +16,7 @@ NOK_THRESHOLD_REACHED = const(3) def check( - current_index: int, - current_word: str, - backup_type: Optional[EnumTypeBackupType], - previous_words: List[str], + backup_type: Optional[EnumTypeBackupType], partial_mnemonic: List[str] ) -> int: # we can't perform any checks if the backup type was not yet decided if backup_type is None: @@ -34,22 +31,22 @@ def check( raise RuntimeError if backup_type == BackupType.Slip39_Basic: - return check_slip39_basic(current_index, current_word, previous_mnemonics) + return check_slip39_basic(partial_mnemonic, previous_mnemonics) if backup_type == BackupType.Slip39_Advanced: - return check_slip39_advanced( - current_index, current_word, previous_words, previous_mnemonics - ) + return check_slip39_advanced(partial_mnemonic, previous_mnemonics) # there are no other backup types raise RuntimeError def check_slip39_basic( - current_index: int, current_word: str, previous_mnemonics: List[List[str]] + partial_mnemonic: List[str], previous_mnemonics: List[List[str]] ) -> int: # check if first 3 words of mnemonic match # we can check against the first one, others were checked already + current_index = len(partial_mnemonic) - 1 + current_word = partial_mnemonic[-1] if current_index < 3: share_list = previous_mnemonics[0][0].split(" ") if share_list[current_index] != current_word: @@ -65,11 +62,10 @@ def check_slip39_basic( def check_slip39_advanced( - current_index: int, - current_word: str, - previous_words: List[str], - previous_mnemonics: List[List[str]], + partial_mnemonic: List[str], previous_mnemonics: List[List[str]] ) -> int: + current_index = len(partial_mnemonic) - 1 + current_word = partial_mnemonic[-1] if current_index < 2: share_list = next(s for s in previous_mnemonics if s)[0].split(" ") if share_list[current_index] != current_word: @@ -87,7 +83,7 @@ def check_slip39_advanced( # check if share was already added for group elif current_index == 3: # we use the 3rd word from previously entered shares to find the group id - group_identifier_word = previous_words[2] + group_identifier_word = partial_mnemonic[2] group_index = None for i, group in enumerate(previous_mnemonics): if len(group) > 0: diff --git a/core/tests/test_apps.management.recovery_device.py b/core/tests/test_apps.management.recovery_device.py index eab24e7e1..497fa4d6a 100644 --- a/core/tests/test_apps.management.recovery_device.py +++ b/core/tests/test_apps.management.recovery_device.py @@ -4,6 +4,8 @@ from mock_storage import mock_storage import storage import storage.recovery from apps.management.recovery_device.recover import process_slip39 +from trezor.messages import BackupType +from apps.management.recovery_device.word_validity import check, OK, NOK_IDENTIFIER_MISMATCH, NOK_ALREADY_ADDED, NOK_THRESHOLD_REACHED MNEMONIC_SLIP39_BASIC_20_3of6 = [ "extra extend academic bishop cricket bundle tofu goat apart victim enlarge program behavior permit course armed jerky faint language modern", @@ -142,21 +144,18 @@ class TestSlip39(unittest.TestCase): @mock_storage def test_check_word_validity(self): - from trezor.messages import BackupType - from apps.management.recovery_device.word_validity import check, OK, NOK_IDENTIFIER_MISMATCH, NOK_ALREADY_ADDED, NOK_THRESHOLD_REACHED - storage.recovery.set_in_progress(True) - # nothing is stored -> should raise + # We claim to know the backup type, but nothing is stored. That is an invalid state. with self.assertRaises(RuntimeError): - check(0, "ocean", BackupType.Slip39_Advanced, []) + check(BackupType.Slip39_Advanced, ["ocean"]) # if backup type is not set we can not do any checks - result = check(0, "ocean", None, []) + result = check(None, ["ocean"]) self.assertIs(result, OK) # BIP-39 has no "on-the-fly" checks - result = check(0, "ocean", BackupType.Bip39, []) + result = check(BackupType.Bip39, ["ocean"]) self.assertIs(result, OK) # let's store two shares in the storage @@ -166,15 +165,15 @@ class TestSlip39(unittest.TestCase): self.assertIsNone(secret) # different identifier - result = check(0, "slush", BackupType.Slip39_Advanced, []) + result = check(BackupType.Slip39_Advanced, ["slush"]) self.assertIs(result, NOK_IDENTIFIER_MISMATCH) # same first word but still a different identifier - result = check(1, "slush", BackupType.Slip39_Advanced, ["trash"]) + result = check(BackupType.Slip39_Advanced, ["trash", "slush"]) self.assertIs(result, NOK_IDENTIFIER_MISMATCH) # same mnemonic found out using the index - result = check(3, "ambition", BackupType.Slip39_Advanced, ["trash", "smug", "adjust"]) + result = check(BackupType.Slip39_Advanced, ["trash", "smug", "adjust", "ambition"]) self.assertIs(result, NOK_ALREADY_ADDED) # Let's store two more. The group is 4/6 so this group is now complete. @@ -184,7 +183,7 @@ class TestSlip39(unittest.TestCase): self.assertIsNone(secret) # If trying to add another one from this group we get a warning. - result = check(2, "adjust", BackupType.Slip39_Advanced, ["trash", "smug"]) + result = check(BackupType.Slip39_Advanced, ["trash", "smug", "adjust"]) self.assertIs(result, NOK_THRESHOLD_REACHED)