add missing import, improve shamir testing in multiwallet (#56)

* add missing import, improve shamir testing in multiwallet

* simplify logic

* black

* flake8

Co-authored-by: Michael Flaxman <mflaxman@gmail.com>

Author: mflaxman

buidl/__init__.py

@@ -14,5 +14,6 @@
 from .pbkdf2 import *
 from .psbt import *
 from .script import *
+from .shamir import *
 from .tx import *
 from .witness import *

multiwallet.py

@@ -587,26 +587,43 @@ def do_split_seed(self, arg):
             maximum=16,
         )
         passphrase = b""
+        # Max possible combos is 12,870 (with 8-of-16)
+        testing_limit = 12870
         if self.ADVANCED_MODE:
             add_passphrase = _get_bool(
                 "Do you want to add a passphrase?", default=False
             )
             if add_passphrase:
                 passphrase = _get_confirmed_pw().encode("ascii")
+            if _get_bool(
+                "Do you want to limit how many of these shares you test (for faster results)?",
+                default=False,
+            ):
+                testing_limit = _get_int(
+                    "How many combinations would you like to test?",
+                    default=1000,
+                    minimum=k,
+                    maximum=testing_limit,
+                )
         shares = ShareSet.generate_shares(mnemonic, k, n, passphrase=passphrase)
         # test the shares (cap at 1000 combinations)
-        print("testing combinations...")
-        iterations = 0
-        for combo in combinations(shares, k):
-            iterations += 1
-            if iterations > 1000:
+        print_yellow(
+            "Testing share combinations to be certain they will recover your seed phrase"
+        )
+        for cnt, combo in enumerate(combinations(shares, k)):
+            if cnt > testing_limit:
                 break
-            elif iterations % 10:
-                print(".", end="", flush=True)
             calculated = ShareSet.recover_mnemonic(combo, passphrase)
             if calculated != mnemonic:
                 # we should never reach this line
                 raise RuntimeError(f"Bad shares {calculated} for mnemonic {mnemonic}")
+            if cnt % 10 == 0:
+                print(".", end="", flush=True)
+
+        print_yellow(
+            f"\nSuccesfully tested {'ALL ' if testing_limit == 12870 else ''}{cnt} combinations"
+        )
+
         print("\n")
         share_mnemonics = "\n\n".join(shares)
         if passphrase:

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="buidl",
-    version="0.2.12",
+    version="0.2.13",
     author="Example Author",
     author_email="author@example.com",
     description="An easy-to-use and fully featured bitcoin library written in pure python (no dependencies).",