shamir/shamir.py

from typing import List, Tuple

import utils

UPPER = (2**8)**64
PRIME = (2**8)**64 + 75


def inverse(num: int) -> int:
    return pow(num, PRIME - 2, PRIME)


def mod(num: int) -> int:
    return num % PRIME


def rand_coefficient() -> int:
    with open('/dev/random', 'rb') as f:
        return utils.bytes_to_long(f.read(64))


def generate_polynomial(num_coeff: int, intercept: int) -> List[int]:
    assert num_coeff >= 1
    return [rand_coefficient() for _ in range(num_coeff - 1)] + [intercept]


def eval_polynomial(polynomial: List[int], point: int):
    result = 0
    for coeff in polynomial:
        result = (result * point + coeff) % PRIME
    return result


def secret_to_parts(secret: int, k: int, n: int) -> List[int]:
    assert 0 <= secret < UPPER
    assert 0 < k <= n

    while True:
        polynomial = generate_polynomial(k, secret)
        parts = [eval_polynomial(polynomial, point)
                 for point in range(1, n + 1)]
        # Check that all parts fit in 64 bytes, otherwise try with new polynomial
        if all(part < UPPER for part in parts):
            return parts


def parts_to_secret(parts: List[Tuple[int, int]]) -> int:
    secret = 0
    for part_id, secret_part in parts:
        current = secret_part
        for other_id, _ in parts:
            if other_id == part_id:
                continue
            current = mod(current * mod(-other_id))
            current = mod(current * inverse(mod(part_id - other_id)))
        secret = mod(secret + current)
    return secret