Merge pull request #541 from trezor/ciny/slip39_various_fixes

core: various fixes to slip39

core/src/apps/common/storage/recovery.py

@@ -131,11 +131,10 @@ def fetch_slip39_remaining_shares() -> Optional[List[int]]:
     if not remaining:
         return None
 
-    result = []
-    for i in range(get_slip39_group_count()):
-        result.append(remaining[i])
-
-    return result
+    group_count = get_slip39_group_count()
+    if not group_count:
+        raise RuntimeError
+    return list(remaining[:group_count])
 
 
 def end_progress() -> None:

core/src/apps/common/storage/recovery_shares.py

@@ -9,12 +9,18 @@ if False:
 # Each mnemonic is stored under key = index.
 
 
-def set(index: int, mnemonic: str) -> None:
-    common.set(common.APP_RECOVERY_SHARES, index, mnemonic.encode())
+def set(index: int, mnemonic: str, group_index: int) -> None:
+    common.set(
+        common.APP_RECOVERY_SHARES,
+        index + group_index * slip39.MAX_SHARE_COUNT,
+        mnemonic.encode(),
+    )
 
 
-def get(index: int) -> Optional[str]:
-    m = common.get(common.APP_RECOVERY_SHARES, index)
+def get(index: int, group_index: int) -> Optional[str]:
+    m = common.get(
+        common.APP_RECOVERY_SHARES, index + group_index * slip39.MAX_SHARE_COUNT
+    )
     if m:
         return m.decode()
     return None
@@ -32,9 +38,8 @@ def fetch() -> List[List[str]]:
 
 def fetch_group(group_index: int) -> List[str]:
     mnemonics = []
-    starting_index = group_index * slip39.MAX_SHARE_COUNT
-    for index in range(starting_index, starting_index + slip39.MAX_SHARE_COUNT):
-        m = get(index)
+    for index in range(slip39.MAX_SHARE_COUNT):
+        m = get(index, group_index)
         if m:
             mnemonics.append(m)
 

core/src/apps/management/recovery_device/recover.py

@@ -29,8 +29,6 @@ def process_slip39(words: str) -> Tuple[Optional[bytes], slip39.Share]:
     share = slip39.decode_mnemonic(words)
 
     remaining = storage.recovery.fetch_slip39_remaining_shares()
-    # TODO: move this whole logic to storage
-    index_with_group_offset = share.index + share.group_index * slip39.MAX_SHARE_COUNT
 
     # if this is the first share, parse and store metadata
     if not remaining:
@@ -42,7 +40,7 @@ def process_slip39(words: str) -> Tuple[Optional[bytes], slip39.Share]:
         storage.recovery.set_slip39_remaining_shares(
             share.threshold - 1, share.group_index
         )
-        storage.recovery_shares.set(index_with_group_offset, words)
+        storage.recovery_shares.set(share.index, words, share.group_index)
 
         # if share threshold and group threshold are 1
         # we can calculate the secret right away
@@ -58,7 +56,7 @@ def process_slip39(words: str) -> Tuple[Optional[bytes], slip39.Share]:
     # These should be checked by UI before so it's a Runtime exception otherwise
     if share.identifier != storage.recovery.get_slip39_identifier():
         raise RuntimeError("Slip39: Share identifiers do not match")
-    if storage.recovery_shares.get(index_with_group_offset):
+    if storage.recovery_shares.get(share.index, share.group_index):
         raise RuntimeError("Slip39: This mnemonic was already entered")
 
     remaining_for_share = (
@@ -69,7 +67,7 @@ def process_slip39(words: str) -> Tuple[Optional[bytes], slip39.Share]:
         remaining_for_share - 1, share.group_index
     )
     remaining[share.group_index] = remaining_for_share - 1
-    storage.recovery_shares.set(index_with_group_offset, words)
+    storage.recovery_shares.set(share.index, words, share.group_index)
 
     if remaining.count(0) < share.group_threshold:
         # we need more shares

core/src/apps/management/reset_device/layout.py

@@ -149,40 +149,38 @@ def _split_share_into_pages(share_words):
 
 
 async def _confirm_share_words(ctx, share_index, share_words, group_index=None):
-    numbered = list(enumerate(share_words))
-
     # divide list into thirds, rounding up, so that chunking by `third` always yields
     # three parts (the last one might be shorter)
-    third = (len(numbered) + 2) // 3
+    third = (len(share_words) + 2) // 3
 
-    for part in utils.chunks(numbered, third):
-        if not await _confirm_word(
-            ctx, share_index, part, len(share_words), group_index
-        ):
+    offset = 0
+    count = len(share_words)
+    for part in utils.chunks(share_words, third):
+        if not await _confirm_word(ctx, share_index, part, offset, count, group_index):
             return False
+        offset += len(part)
 
     return True
 
 
-async def _confirm_word(
-    ctx, share_index, numbered_share_words, count, group_index=None
-):
+async def _confirm_word(ctx, share_index, share_words, offset, count, group_index=None):
+    # remove duplicates
+    non_duplicates = list(set(share_words))
+    # shuffle list
+    random.shuffle(non_duplicates)
+    # take top NUM_OF_CHOICES words
+    choices = non_duplicates[: MnemonicWordSelect.NUM_OF_CHOICES]
+    # select first of them
+    checked_word = choices[0]
+    # find its index
+    checked_index = share_words.index(checked_word) + offset
+    # shuffle again so the confirmed word is not always the first choice
+    random.shuffle(choices)
 
-    # TODO: duplicated words in the choice list
-    # shuffle the numbered seed half, slice off the choices we need
-    random.shuffle(numbered_share_words)
-    numbered_choices = numbered_share_words[: MnemonicWordSelect.NUM_OF_CHOICES]
-
-    # we always confirm the first (random) word index
-    checked_index, checked_word = numbered_choices[0]
     if __debug__:
         debug.reset_word_index.publish(checked_index)
 
-    # shuffle again so the confirmed word is not always the first choice
-    random.shuffle(numbered_choices)
-
     # let the user pick a word
-    choices = [word for _, word in numbered_choices]
     select = MnemonicWordSelect(choices, share_index, checked_index, count, group_index)
     if __debug__:
         selected_word = await ctx.wait(select, debug.input_signal())
@@ -565,7 +563,10 @@ class Slip39NumInput(ui.Component):
             elif self.step is Slip39NumInput.SET_THRESHOLD:
                 if self.group_id is None:
                     first_line_text = "For recovery you need"
-                    second_line_text = "any %s of the shares." % count
+                    if count == self.input.max_count:
+                        second_line_text = "all %s of the shares." % count
+                    else:
+                        second_line_text = "any %s of the shares." % count
                 else:
                     first_line_text = "The required number of "
                     second_line_text = "shares to form Group %s." % (self.group_id + 1)