#!/usr/bin/env python3
import random
import time

def rabin_miller(n, certainty=128):
    """Determine if a positive integer is (probably) a prime.

    Positional argument:

        n   --      The positive integer to check for primality.

    Keyword argument:

        certainty   --  How sure you want to be than n really is a prime. See
                        below for details.

    Returns:

        True        --  n is probably prime. The probability that n is not
                        prime is <= 1/2**certainty.
        False       --  n is *definitely* not prime.
    """
    rand = random.SystemRandom()
    def test_once(d,r,n):
        a = rand.randint(2, n-2)
        x = pow(a, d, n)
        if x == 1 or x == n-1: return True
        for i in range(1,r):
            x = (x*x) % n
            if x == n-1: return True
        return False

    # Take care of special cases
    if n == 2 or n == 3: return True
    if n == 1 or n % 2 == 0: return False

    # Find r, d such that n = 2**r * d + 1, with d odd
    d = n - 1
    r = 0
    while d % 2 == 0:
        d = d // 2
        r += 1
    for t in range(certainty//2+1):
        if test_once(d,r,n) == False:
            return False
    return True

def find_prime(start=0, bits=0, linear=True):
    """Find a randomly chosen prime of a given number of bits, or find the
    first prime >= a given number.

    Keyword arguments:

        bits    --  Number of bits if you want a random prime.

        start   --  Starting search position if you want a non-random prime.

        linear  --  Controls search for random primes. See below.

        Note: you should specify exactly one of "bits" and "start".

    Returns:

        prime   --  A prime number

    When searching for a random prime, "linear" specifies what happens if
    the random number it picks is not prime. If linear is True, it will add
    2 and check that number. If linear is False, it will pick another random
    number to check. When linear is True, if adding 2 would take it past
    2**bits, it will wrap to 2**(bits-1)+1.
    """
    rand = random.SystemRandom()
    if bits != 0 and start == 0:
        base = 2**(bits-1)
        wrap = 2*base
        p = rand.randint(base, wrap-1) | 1
    elif bits == 0 and start != 0:
        if start <= 2:
            return 2
        p = start | 1
        linear = True
        base = start ^ 1
        wrap = 2 * p
    else:
        print("Error: find_prime must be given exactly one of bits and start")
        return 0
    while not rabin_miller(p):
        if linear:
            p += 2
            if p >= wrap:
                p = base + 1
        else:
            p = rand.randint(base, wrap-1) | 1
    return p

def poly_at(x, poly, m):
    y = 0
    f = 1
    for c in poly:
        y += c * f
        y = y % m
        f = (f * x) % m
    return y

def row_multiply(r, a, m):
    """multiply row `r` by ``a` mod `m`"""
    for i in range(len(r)):
        r[i] = (r[i] * a) % m

def row_multiply_add(r, r2, a, m):
    for i in range(len(r)):
        r[i] = (r[i] + r2[i] * a) % m

def zero_column(eqs, c, r, m):
    row_multiply(eqs[r], pow(eqs[r][c], -1, m), m)
    for j in range(len(eqs)):
        if j != r:
            a = (m - eqs[j][c]) % m
            row_multiply_add(eqs[j], eqs[r], a, m)

def solve(eqs, m):
    for d in range(len(eqs)):
        if eqs[d][d] == 0:
            for j in range(d+1,len(eqs)):
                if eqs[j][d] != 0:
                    eqs[d], eqs[j] = eqs[j], eqs[d]
                    print("swap rows",d,"and",j)
                    break
        zero_column(eqs, d, d, m)

def recover_secret(shares, p):
    num_shares = len(shares)
    equations = []
    for r in range(num_shares):
        i, share_i = shares[r]
        g = 1
        row = [1]
        for t in range(1, num_shares):
            g = (g * i) % p
            row.append(g)
        row.append(share_i)
        equations.append(row)
    solve(equations, p)
    return [r[-1] for r in equations]

def make_sss(s, k, p=0):
    """make a secret sharing scheme to share secret `s` with `k` shares
    required to reconstruct the secret"""
    rand = random.SystemRandom()
    if p == 0:
        bits = 2
        base = 2
        while base <= s:
            base *= 2
            bits += 1
        p = find_prime(bits=bits, linear=False)
    else:
        if s >= p:
            print("secret is too large for prime")
            return (0, [])
    poly = [s]
    for i in range(k-1):
        poly.append(rand.randint(1, p-1))
    return (p, poly)

def make_share(sss, i):
    return poly_at(i, sss[1], sss[0])

if __name__ == "__main__":
    rand = random.SystemRandom()
    k = 3
    bits = 128
    sss = make_sss(rand.randint(1, 2**bits-1), k)
    print("prime", sss[0])
    print(sss[1])
    shares = [[i**2+1, make_share(sss, i**2+1)] for i in range(k)]
    poly = recover_secret(shares, sss[0])
    print(poly)
    print(shares)