shamir_bech32.py

# Copyright (c) 2021 Samuel Dobson
# Copyright (c) 2021 Russell O'Connor
# Copyright (c) 2017, 2020 Pieter Wuille
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import random
import math

CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
BECH32M_CONST = 0x2bc830a3
MS32_CONST = 0x10ce0795c2fd1e62a

GF32_EXP = [1, 2, 4, 8, 16, 9, 18, 13, 26, 29, 19, 15, 30, 21, 3, 6, 12, 24, 25, 27, 31, 23, 7, 14, 28, 17, 11, 22, 5, 10, 20, 1]
GF32_LOG = [-1, 31, 1, 14, 2, 28, 15, 22, 3, 5, 29, 26, 16, 7, 23, 11, 4, 25, 6, 10, 30, 13, 27, 21, 17, 18, 8, 19, 24, 9, 12, 20]

# 41 encodes the defining polynomial of the extension GF(32) over GF(2)
GF32_EXT_MOD = 41

class ChecksumType:
    CHECKSUM_CONST = 1
    polymod = lambda x: x
    LENGTH = 0
    def __init__(self, c, pf, l):
        self.CHECKSUM_CONST = c
        self.polymod = pf
        self.LENGTH = l

def bech32_polymod(values):
    """Internal function that computes the Bech32 checksum."""
    generator = [0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3]
    chk = 1
    for value in values:
        top = chk >> 25
        chk = (chk & 0x1ffffff) << 5 ^ value
        for i in range(5):
            chk ^= generator[i] if ((top >> i) & 1) else 0
    return chk

def ms32_polymod(values):
    GEN = [0x1f28f80fffe92f842, 0x1751a20bdef255484, 0x07a316039ceda0d08, 0x0e0e2c0739da09a10, 0x1c164a0e739d13129]
    residue = 0x23181b3
    for v in values:
        b = (residue >> 60)
        residue = (residue & 0x0fffffffffffffff) << 5 ^ v
        for i in range(5):
            residue ^= GEN[i] if ((b >> i) & 1) else 0
    return residue


BECH32M_CHECKSUM = ChecksumType(BECH32M_CONST, bech32_polymod, 6)
MS32_CHECKSUM = ChecksumType(MS32_CONST, ms32_polymod, 13)

def create_checksum(ctype, values):
    """Compute the checksum values given values without a checksum."""
    polymod = ctype.polymod(values + [0] * ctype.LENGTH) ^ ctype.CHECKSUM_CONST
    return [(polymod >> 5 * (ctype.LENGTH - 1 - i)) & 31 for i in range(ctype.LENGTH)]

def bech32_charset_encode(data):
    """Encode data in the bech32 characterset."""
    return ''.join([CHARSET[d] for d in data]) 

def bech32_charset_decode(enc):
    """Decode data from the bech32 characterset."""
    return [CHARSET.find(x) for x in enc]

def convertbits(data, frombits, tobits, pad=True):
    """General power-of-2 base conversion."""
    acc = 0
    bits = 0
    ret = []
    maxv = (1 << tobits) - 1
    max_acc = (1 << (frombits + tobits - 1)) - 1
    for value in data:
        if value < 0 or (value >> frombits):
            return None
        acc = ((acc << frombits) | value) & max_acc
        bits += frombits
        while bits >= tobits:
            bits -= tobits
            ret.append((acc >> bits) & maxv)
    if pad:
        if bits:
            ret.append((acc << (tobits - bits)) & maxv)
    elif bits >= frombits or ((acc << (tobits - bits)) & maxv):
        return None
    return ret

def verify_GF32_tables():
    """Helper function for testing, verifies the GF32_LOG and _EXP tables are correct."""
    v = 1;
    for i in range(1, 32):
        v = v << 1
        if (v & 32):
            v ^= GF32_EXT_MOD
        if GF32_LOG[v] != i:
            print("ERROR: Expected LOG of {} to be {}, got {}.".format(v, i, GF32_LOG[v]))
        if GF32_EXP[i] != v:
            print("ERROR: Expected EXP of {} to be {}, got {}.".format(i, v, GF32_EXP[i]))
    print("GF32 table verification complete!\n")

def gf32_lagrange_interpolation(x, points):
    """construct the lagrange interpolation of points and evaluate at x, in GF(32)"""
    res = 0
    for i, (x_i, y_i) in enumerate(points):
        # evaulate lagrange basis polynomial l_i at x
        # l_i(x) = product of (x - x_j) / (x_i - x_j) for j != i

        # the entire term will be 0 if x == x_i is
        if y_i == 0:
            continue

        log_lix = 0
        for j, (x_j, y_j) in enumerate(points):
            if i == j: continue    
            # in characteristic 2, addition and subtraction are both XOR
            log_lix += GF32_LOG[( x ^ x_j )] + 31
            log_lix -= GF32_LOG[( x_i ^ x_j )]
            log_lix %= 31

        # Now multiply l_i(x) by y_i and add to result
        log_lix += GF32_LOG[y_i]
        ylix = GF32_EXP[log_lix % 31]
        res ^= ylix
    return res

def reconstruct_coeff(points, k):
    """Given a list of shares, reconstruct the secret using lagrange interpolation"""
    if len(points) < k:
        print("Require at least {} shares to reconstruct secret, only {} given.".format(k, len(points)))
        return -1
    return gf32_lagrange_interpolation(0, points)

def generate_coeff_shares(secret, k, n):
    """Given an element in GF(32), construct n shamir secret shares with threshold k"""
    if n < k or k < 1:
        print("Invalid sharing parameters. Must have 1 <= k <= n.")
        return -1

    # Generate coefficients of a polynomial of degree k-1
    coeffs = [random.getrandbits(5) for i in range(k-1)]
    log_secret = GF32_LOG[secret]
    log_coeffs = [log_secret] + [GF32_LOG[c] for c in coeffs]

    shares = []
    for i in range(1, n+1):
        fi = 0
        for j, log_aj in enumerate(log_coeffs):
            # compute (i^j)*aj
            if log_aj == -1:
                continue
            log_fij = (GF32_LOG[i] * j + log_aj) % 31
            fi ^= GF32_EXP[log_fij]
        share = (i, fi)
        shares.append(share)

    return shares

def generate_codeword(ctype, data_len=10):
    """Helper function to generate a valid codeword (with checksum) that contains data_len-bytes of data."""
    data = random.randbytes(data_len)
    five_bit_values = convertbits(data, 8, 5)
    checksum = create_checksum(ctype, five_bit_values)
    codeword = five_bit_values + checksum
    return codeword

def test_interpolation(ctype, k):
    print("\nTesting interpolation...")
    # data_len is the number of bytes (8-bits) encoded in each codeword
    data_len = 10

    codewords = [generate_codeword(ctype, data_len) for i in range(k)]
    for c in codewords:
        if ctype.polymod(c) != ctype.CHECKSUM_CONST:
            print("ERROR: codeword {} has invalid checksum.".format(c))

    # The length of each codeword should be (8/5)*data_len + CHECKSUM_LENGTH
    code_length = math.ceil(8*data_len/5 + ctype.LENGTH)
    for code in codewords:
        if len(code) != code_length:
            print("ERROR: Length mismatch, expected code length {}, got {}".format(code_length, len(code)))

    codewords = [zip([i+1]*(code_length), c) for i,c in enumerate(codewords)]
    coeff_shares = list(zip(*codewords))

    reconstructed_secret = []
    for cs in coeff_shares:
        secret_coeff = reconstruct_coeff(list(cs), k)
        reconstructed_secret.append(secret_coeff)
    print("Reconstructed secret:", reconstructed_secret)
    
    print("Checksum verifies:", ctype.polymod(reconstructed_secret) == ctype.CHECKSUM_CONST)

def main():
    # Verify the GF32_LOG and _EXP tables are correct
    verify_GF32_tables()

    # Test that recovering a secret from two valid codewords gives a third valid codeword
    test_interpolation(BECH32M_CHECKSUM, 2)
    test_interpolation(MS32_CHECKSUM, 2)

if __name__ == "__main__":
    main()