support Shamir's secret sharing with minimal code (`shamir.py`)

2 files changed, 200 insertions, 3 deletions

diff --git a/keytree.py b/keytree.py
index 78cdf4d..71d5998 100755
--- a/keytree.py
+++ b/keytree.py
@@ -1,4 +1,4 @@
-#!/usr/bin/env python3
+#!/usr/bin/env python3.10
 # MIT License
 #
 # Copyright (c) 2020 Ted Yin <[email protected]>
@@ -47,6 +47,7 @@ import hmac
 import unicodedata
 import json
 from getpass import getpass as _getpass
+from itertools import zip_longest
 
 import bech32
 import mnemonic
@@ -58,6 +59,7 @@ from sha3 import keccak_256
 from uuid import uuid4
 from Cryptodome.Cipher import AES
 from Cryptodome.Util import Counter
+import shamir
 
 
 def getpass(prompt):
@@ -124,7 +126,7 @@ def is_infinity(P):
 # parse256(p): interprets a 32-byte sequence as a 256-bit number, most
 # significant byte first.
 def parse256(p):
-    assert(len(p) == 32)
+    assert len(p) == 32
     return int.from_bytes(p, byteorder='big')
 
 
@@ -346,6 +348,42 @@ def save_to_mew(priv_keys, n=1 << 18, p=1, r=8, dklen=32):
         raise KeytreeError("failed while saving")
 
 
+def to_chunks(n, iterable):
+    return zip_longest(*[iter(iterable)]*n, fillvalue=0)
+
+def shamir256_split(secret, t, n):
+    shares = [bytearray() for i in range(n)]
+    for chunk in to_chunks(32, secret):
+        secret = int.from_bytes(chunk, 'big')
+        while True:
+            points = shamir.split(secret, t, n)
+            good = True
+            for p in points:
+                if p.bit_length() > 256:
+                    good = False
+                    break
+            if good:
+                # all shares are within 256 bits
+                for (p, s) in zip(points, shares):
+                    s.extend(p.to_bytes(32, 'big'))
+                break
+    return shares
+
+
+def shamir256_combine(shares):
+    result = bytearray()
+    for shares in zip(*[[(i, g) for g in to_chunks(32, s)] for (i, s) in shares]):
+        shares = [(i, int.from_bytes(bytearray(p), 'big')) for (i, p) in shares]
+        try:
+            secret = shamir.combine(shares)
+        except ValueError:
+            raise KeytreeError("invalid Shamir recovery input")
+        if secret.bit_length() > 256:
+            raise KeytreeError("Shamir result is too long")
+        result.extend(secret.to_bytes(32, 'big'))
+    return result
+
+
 metamask_path = r"44'/60'/0'/0"
 avax_path = r"44'/9000'/0'/0"
 
@@ -355,6 +393,10 @@ if __name__ == '__main__':
     parser.add_argument('--save', type=str, default=None, help='save mnemonic to a file (AVAX Wallet compatible)')
     parser.add_argument('--export-mew', action='store_true', default=False, help='export keys to MEW keystore files (mnemonic is NOT saved, only keys are saved)')
     parser.add_argument('--show-private', action='store_true', default=False, help='also show private keys and the mnemonic')
+    parser.add_argument('--gen-shamir', action='store_true', default=False, help='generate Shamir\'s secret shares')
+    parser.add_argument('--shamir-threshold', type=int, default=2, help='Shamir\'s secret sharing threshold (number of shares to decode)')
+    parser.add_argument('--shamir-num', type=int, default=3, help='Shamir\'s secret sharing number (total number of shares)')
+    parser.add_argument('--recover-shamir', type=str, default=None, help='recover the secret from Shamir shares')
     parser.add_argument('--custom', action='store_true', default=False, help='use an arbitrary word combination as mnemonic')
     parser.add_argument('--seed', action='store_true', default=False, help='load mnemonic from seed')
     parser.add_argument('--path', default=avax_path, help="path prefix for key deriving (e.g. \"{}\" for Metamask)".format(metamask_path))
@@ -371,13 +413,40 @@ if __name__ == '__main__':
         for arg in unknown:
             if len(arg) > 0:
                 raise KeytreeError("invalid argument: `{}`".format(arg))
+        shares = []
         try:
+            mgen = mnemonic.Mnemonic(args.lang)
             if args.gen_mnemonic:
-                mgen = mnemonic.Mnemonic(args.lang)
                 words = mgen.generate(256)
             else:
                 if args.load:
                     words = load_from_keystore(args.load)
+                elif args.recover_shamir:
+                    try:
+                        idxes = [int(i) for i in args.recover_shamir.split(',')]
+                    except ValueError:
+                        raise KeytreeError("invalid Shamir share spec, should be something like \"1,2\"")
+                    custom_mnemonic = None
+                    for idx in idxes:
+                        swords = getpass('Enter the mnemonic for Shamir share #{}: '.format(idx))
+                        if len(swords) == 48:
+                            if not custom_mnemonic:
+                                raise KeytreeError("invalid Shamir share format")
+                            custom_mnemonic = True
+                            share = mgen.to_entropy(swords[:24]) + mgen.to_entropy(swords[24:])
+                        else:
+                            if custom_mnemonic:
+                                raise KeytreeError("invalid Shamir share format")
+                            custom_mnemonic = False
+                            try:
+                                share = mgen.to_entropy(swords)
+                            except ValueError:
+                                raise KeytreeError('invalid mnemonic')
+                        shares.append((idx, share))
+                    if custom_mnemonic:
+                        seed = shamir256_combine(shares)
+                    else:
+                        words = mgen.to_mnemonic(shamir256_combine(shares))
                 elif not args.seed:
                     words = getpass('Enter the mnemonic: ').strip()
                     if not args.custom:
@@ -402,6 +471,21 @@ if __name__ == '__main__':
             if not args.seed:
                 print("KEEP THIS PRIVATE (mnemonic): {}".format(words))
             print("KEEP THIS PRIVATE (seed): {}".format(seed.hex()))
+        if args.shamir_threshold:
+            if args.shamir_num > 20:
+                raise KeytreeError('Shamir threshold should be <= 20')
+            if args.shamir_threshold < 2 or args.shamir_threshold > args.shamir_num:
+                raise KeytreeError('Shamir threshold should be (2, N]')
+        if args.gen_shamir:
+            if args.seed:
+                shares = shamir256_split(seed, args.shamir_threshold, args.shamir_num)
+                for idx, share in enumerate(shares):
+                    words = mgen.to_mnemonic(share[:32]) + mgen.to_mnemonic(share[32:])
+                    print("KEEP THIS PRIVATE: secret_{}: {}".format(idx + 1, words))
+            else:
+                shares = shamir256_split(mgen.to_entropy(words), args.shamir_threshold, args.shamir_num)
+                for idx, share in enumerate(shares):
+                    print("KEEP THIS PRIVATE: secret_{}: {}".format(idx + 1, mgen.to_mnemonic(share)))
         gen = BIP32(seed)
         if args.start_idx < 0 or args.end_idx < 0:
             raise KeytreeError("invalid start/end index")
diff --git a/shamir.py b/shamir.py
new file mode 100644
index 0000000..c7d606a
--- /dev/null
+++ b/shamir.py
@@ -0,0 +1,113 @@
+# Code modified from https://github.com/kurtbrose/shamir/tree/master
+
+'''
+An implementation of shamir secret sharing algorithm.
+
+To the extent possible under law,
+all copyright and related or neighboring rights
+are hereby waived.  (CC0, see LICENSE file.)
+
+All possible patents arising from this code
+are renounced under the terms of
+the Open Web Foundation CLA 1.0
+(http://www.openwebfoundation.org/legal/the-owf-1-0-agreements/owfa-1-0)
+'''
+
+from __future__ import division
+import random
+import functools
+
+# 12th Mersenne Prime is 2**127 - 1
+# use the prime to be compatible with the Shamir tool developed by Ava Labs:
+# https://github.com/ava-labs/mnemonic-shamir-secret-sharing-cli/tree/main
+# This is a 257-bit prime, so the returned points could be either 256-bit or
+# (infrequently) 257-bit.
+_PRIME = 187110422339161656731757292403725394067928975545356095774785896842956550853219
+# 13th Mersenne Prime is 2**521 - 1
+
+def _eval_at(poly, x, prime):
+    'evaluate polynomial (coefficient tuple) at x'
+    accum = 0
+    for coeff in reversed(poly):
+        accum *= x
+        accum += coeff
+        accum %= prime
+    return accum
+
+
+def split(secret, minimum, shares, prime=_PRIME):
+    '''
+    Generates a random shamir pool, returns
+    the secret and the share points.
+    '''
+    randint = functools.partial(random.SystemRandom().randint, 0)
+    if minimum > shares:
+        raise ValueError("pool secret would be irrecoverable")
+    poly = [randint(prime) for i in range(minimum - 1)]
+    poly.insert(0, secret)
+    return [_eval_at(poly, i, prime) for i in range(1, shares + 1)]
+
+
+# division in integers modulus p means finding the inverse of the denominator
+# modulo p and then multiplying the numerator by this inverse
+# (Note: inverse of A is B such that A*B % p == 1)
+# this can be computed via extended euclidean algorithm
+# http://en.wikipedia.org/wiki/Modular_multiplicative_inverse#Computation
+def _extended_gcd(a, b):
+    x = 0
+    last_x = 1
+    y = 1
+    last_y = 0
+    while b != 0:
+        quot = a // b
+        a, b = b,  a%b
+        x, last_x = last_x - quot * x, x
+        y, last_y = last_y - quot * y, y
+    return last_x, last_y
+
+
+def _divmod(num, den, p):
+    '''
+    compute num / den modulo prime p
+    To explain what this means, the return
+    value will be such that the following is true:
+    den * _divmod(num, den, p) % p == num
+    '''
+    inv, _ = _extended_gcd(den, p)
+    return num * inv
+
+
+def _lagrange_interpolate(x, x_s, y_s, p):
+    '''
+    Find the y-value for the given x, given n (x, y) points;
+    k points will define a polynomial of up to kth order
+    '''
+    k = len(x_s)
+    assert k == len(set(x_s)), "points must be distinct"
+    def PI(vals):  # upper-case PI -- product of inputs
+        accum = 1
+        for v in vals:
+            accum *= v
+        return accum
+    nums = []  # avoid inexact division
+    dens = []
+    for i in range(k):
+        others = list(x_s)
+        cur = others.pop(i)
+        nums.append(PI(x - o for o in others))
+        dens.append(PI(cur - o for o in others))
+    den = PI(dens)
+    num = sum([_divmod(nums[i] * den * y_s[i] % p, dens[i], p)
+               for i in range(k)])
+    return (_divmod(num, den, p) + p) % p
+
+
+def combine(shares, prime=_PRIME):
+    '''
+    Recover the secret from share points
+    (x,y points on the polynomial)
+    '''
+    if len(shares) < 2:
+        raise ValueError("need at least two shares")
+    x_s, y_s = zip(*shares)
+    return _lagrange_interpolate(0, x_s, y_s, prime)