From 487cc1975088a731436eec541403d93e68a5b4b7 Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Tue, 25 Jun 2019 18:25:32 +0200
Subject: [PATCH] core/recovery_device: fix dry_run

---
 core/src/apps/common/mnemonic/bip39.py      |  4 ++++
 core/src/apps/common/mnemonic/slip39.py     |  5 +++++
 core/src/apps/management/recovery_device.py | 23 ++++++++++++++++-----
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/core/src/apps/common/mnemonic/bip39.py b/core/src/apps/common/mnemonic/bip39.py
index bb2d68051..2bb437625 100644
--- a/core/src/apps/common/mnemonic/bip39.py
+++ b/core/src/apps/common/mnemonic/bip39.py
@@ -34,5 +34,9 @@ def get_seed(secret: bytes, passphrase: str):
     return seed
 
 
+def get_mnemonic_count(mnemonic: str) -> int:
+    return 1
+
+
 def check(secret: bytes):
     return bip39.check(secret)
diff --git a/core/src/apps/common/mnemonic/slip39.py b/core/src/apps/common/mnemonic/slip39.py
index 954edb4c7..b2a0bc704 100644
--- a/core/src/apps/common/mnemonic/slip39.py
+++ b/core/src/apps/common/mnemonic/slip39.py
@@ -84,3 +84,8 @@ def get_seed(encrypted_master_secret: bytes, passphrase: str):
     )
     mnemonic._stop_progress()
     return master_secret
+
+
+def get_mnemonic_count(mnemonic: str) -> int:
+    _, _, _, _, _, _, threshold, _ = slip39.decode_mnemonic(mnemonic)
+    return threshold
diff --git a/core/src/apps/management/recovery_device.py b/core/src/apps/management/recovery_device.py
index 33db047a8..207f41e05 100644
--- a/core/src/apps/management/recovery_device.py
+++ b/core/src/apps/management/recovery_device.py
@@ -64,6 +64,10 @@ async def recovery_device(ctx, msg):
         # show a note about the keyboard
         await show_keyboard_info(ctx)
 
+    if msg.dry_run:
+        dry_run_mnemonics = []
+        dry_run_mnemonic_count = None
+
     secret = None
     while secret is None:
         # ask for mnemonic words one by one
@@ -71,14 +75,23 @@ async def recovery_device(ctx, msg):
             ctx, wordcount, mnemonic_module == mnemonic.slip39
         )
         try:
-            secret = mnemonic_module.process_single(words)
+            if msg.dry_run:
+                if dry_run_mnemonic_count is None:
+                    dry_run_mnemonic_count = mnemonic_module.get_mnemonic_count(words)
+                dry_run_mnemonics.append(words)
+            else:
+                secret = mnemonic_module.process_single(words)
         except slip39.MnemonicError as e:
             raise wire.ProcessError("Mnemonic is not valid: " + str(e))
+        if msg.dry_run:
+            remaining = dry_run_mnemonic_count - len(dry_run_mnemonics)
+            if remaining == 0:
+                secret = mnemonic_module.process_all(dry_run_mnemonics)
+        else:
+            remaining = storage.get_slip39_remaining()
         # show a number of remaining mnemonics for SLIP39
         if secret is None and mnemonic_module == mnemonic.slip39:
-            await show_remaining_slip39_mnemonics(
-                ctx, title, storage.get_slip39_remaining()
-            )
+            await show_remaining_slip39_mnemonics(ctx, title, remaining)
 
     # check mnemonic validity
     # it is checked automatically in SLIP-39
@@ -94,7 +107,7 @@ async def recovery_device(ctx, msg):
 
     # dry run
     if msg.dry_run:
-        mnemonic.dry_run(secret)
+        return mnemonic.dry_run(secret)
 
     # save into storage
     if msg.pin_protection: