shamirss/shamirss.py

# Copyright 2023 John-Mark Gurney.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#

import functools
import operator
import secrets
import unittest.mock

random = secrets.SystemRandom()

def _makered(x, y):
	'''Make reduction table entry.

	given x * 2^8, reduce it assuming polynomial y.
	'''

	x = x << 8

	for i in range(3, -1, -1):
		if x & (1 << (i + 8)):
			x ^= (0x100 + y) << i

	assert x < 256

	return x

def evalpoly(polynomial, powers):
	return sum(( x * y for x, y in zip(polynomial, powers)), 0)

def create_shares(data, k, nshares):
	'''Given data, create nshares, such that given any k shares,
	data can be recovered.

	data must be bytes, or able to be converted to bytes.

	The return value will be a list of length nshares.  Each element
	will be a tuple of (<int in range [1, nshares + 1)>, <bytes>).'''

	data = bytes(data)

	powers = { x: GF2p8(x).powerseries(k - 1) for x in range(1, nshares + 1) }

	coeffs = [ [ x ] + [ random.randint(1, 255) for y in range(k - 1) ] for idx, x in enumerate(data) ]

	return [ (x, bytes([ int(evalpoly(coeffs[idx], powers[x])) for idx, val in enumerate(data) ])) for x in range(1, nshares + 1) ]

def recover_data(shares, k):
	'''Recover the value given shares, where k is needed.

	shares must be as least length of k.'''

	if len(shares) < k:
		raise ValueError('not enough shares to recover')

	return bytes([ int(sum([ GF2p8(y[idx]) *
	    functools.reduce(operator.mul, [ pix * ((GF2p8(pix) - x) ** -1) for
	    pix, piy in shares[:k] if pix != x ], 1) for x, y in shares[:k] ],
	    0)) for idx in range(len(shares[0][1]))])

class GF2p8:
	_invcache = {1: 1, 2: 195, 3: 130, 4: 162, 5: 126, 6: 65, 7: 90, 8: 81, 9: 54, 10: 63, 11: 172, 12: 227, 13: 104, 14: 45, 15: 42, 16: 235, 17: 155, 18: 27, 19: 53, 20: 220, 21: 30, 22: 86, 23: 165, 24: 178, 25: 116, 26: 52, 27: 18, 28: 213, 29: 100, 30: 21, 31: 221, 32: 182, 33: 75, 34: 142, 35: 251, 36: 206, 37: 233, 38: 217, 39: 161, 40: 110, 41: 219, 42: 15, 43: 44, 44: 43, 45: 14, 46: 145, 47: 241, 48: 89, 49: 215, 50: 58, 51: 244, 52: 26, 53: 19, 54: 9, 55: 80, 56: 169, 57: 99, 58: 50, 59: 245, 60: 201, 61: 204, 62: 173, 63: 10, 64: 91, 65: 6, 66: 230, 67: 247, 68: 71, 69: 191, 70: 190, 71: 68, 72: 103, 73: 123, 74: 183, 75: 33, 76: 175, 77: 83, 78: 147, 79: 255, 80: 55, 81: 8, 82: 174, 83: 77, 84: 196, 85: 209, 86: 22, 87: 164, 88: 214, 89: 48, 90: 7, 91: 64, 92: 139, 93: 157, 94: 187, 95: 140, 96: 239, 97: 129, 98: 168, 99: 57, 100: 29, 101: 212, 102: 122, 103: 72, 104: 13, 105: 226, 106: 202, 107: 176, 108: 199, 109: 222, 110: 40, 111: 218, 112: 151, 113: 210, 114: 242, 115: 132, 116: 25, 117: 179, 118: 185, 119: 135, 120: 167, 121: 228, 122: 102, 123: 73, 124: 149, 125: 153, 126: 5, 127: 163, 128: 238, 129: 97, 130: 3, 131: 194, 132: 115, 133: 243, 134: 184, 135: 119, 136: 224, 137: 248, 138: 156, 139: 92, 140: 95, 141: 186, 142: 34, 143: 250, 144: 240, 145: 46, 146: 254, 147: 78, 148: 152, 149: 124, 150: 211, 151: 112, 152: 148, 153: 125, 154: 234, 155: 17, 156: 138, 157: 93, 158: 188, 159: 236, 160: 216, 161: 39, 162: 4, 163: 127, 164: 87, 165: 23, 166: 229, 167: 120, 168: 98, 169: 56, 170: 171, 171: 170, 172: 11, 173: 62, 174: 82, 175: 76, 176: 107, 177: 203, 178: 24, 179: 117, 180: 192, 181: 253, 182: 32, 183: 74, 184: 134, 185: 118, 186: 141, 187: 94, 188: 158, 189: 237, 190: 70, 191: 69, 192: 180, 193: 252, 194: 131, 195: 2, 196: 84, 197: 208, 198: 223, 199: 108, 200: 205, 201: 60, 202: 106, 203: 177, 204: 61, 205: 200, 206: 36, 207: 232, 208: 197, 209: 85, 210: 113, 211: 150, 212: 101, 213: 28, 214: 88, 215: 49, 216: 160, 217: 38, 218: 111, 219: 41, 220: 20, 221: 31, 222: 109, 223: 198, 224: 136, 225: 249, 226: 105, 227: 12, 228: 121, 229: 166, 230: 66, 231: 246, 232: 207, 233: 37, 234: 154, 235: 16, 236: 159, 237: 189, 238: 128, 239: 96, 240: 144, 241: 47, 242: 114, 243: 133, 244: 51, 245: 59, 246: 231, 247: 67, 248: 137, 249: 225, 250: 143, 251: 35, 252: 193, 253: 181, 254: 146, 255: 79}

	@staticmethod
	def _primativemul(a, b):
		masks = [ 0, 0xff ]

		r = 0

		for i in range(0, 8):
			mask = a & 1
			r ^= (masks[mask] & b) << i
			a = a >> 1

		return r

	# polynomial 0x187
	_reduce = { x: _makered(x, 0x87) for x in range(0, 16) }

	def __init__(self, v):
		if v >= 256:
			raise ValueError('%d is not a member of GF(2^8)' % v)

		self._v = v

	# basic operations
	def __add__(self, o):
		if isinstance(o, int):
			return self + self.__class__(o)

		return self.__class__(self._v ^ o._v)

	def __radd__(self, o):
		return self.__add__(o)

	def __sub__(self, o):
		return self.__add__(o)

	def __rsub__(self, o):
		return self.__sub__(o)

	def __mul__(self, o):

		if isinstance(o, int):
			o = self.__class__(o)

		m = o._v

		# multiply
		r = self._primativemul(self._v, m)

		# reduce
		r ^= self._reduce[r >> 12] << 4
		r ^= self._reduce[(r >> 8) & 0xf ]

		r &= 0xff

		return self.__class__(r)

	def __rmul__(self, o):
		return self.__mul__(o)

	def __pow__(self, x):
		if x == -1 and self._invcache:
			return GF2p8(self._invcache[self._v])

		if x < 0:
			x += 255

		v = self.__class__(1)

		# TODO - make faster via caching and squaring
		for i in range(x):
			v *= self

		return v

	def powerseries(self, cnt):
		'''Generate [ self ** 0, self ** 1, ..., self ** cnt ].'''

		r = [ 1 ]

		if cnt > 0:
			r.append(self)

		for i in range(2, cnt):
			r.append(r[-1] * self)

		return r

	def __eq__(self, o):
		if isinstance(o, int):
			return self._v == o

		return self._v == o._v

	def __int__(self):
		return self._v

	def __repr__(self):
		return '%s(%d)' % (self.__class__.__name__, self._v)

class TestShamirSS(unittest.TestCase):
	def test_evalpoly(self):
		a = GF2p8(random.randint(0, 255))

		powers = a.powerseries(5)

		vals = [ GF2p8(random.randint(0, 255)) for x in range(5) ]

		r = evalpoly(vals, powers)
		self.assertEqual(r, vals[0] + vals[1] * powers[1] + vals[2] * powers[2] + vals[3] * powers[3] + vals[4] * powers[4])

		r = evalpoly(vals[:3], powers)
		self.assertEqual(r, vals[0] + vals[1] * powers[1] + vals[2] * powers[2])

	def test_create_shares(self):
		self.assertRaises(TypeError, create_shares, '', 1, 1)

		val = bytes([ random.randint(0, 255) for x in range(100) ])

		a = create_shares(val, 2, 3)

		# that it has the number of shares
		self.assertEqual(len(a), 3)

		# that the length of the share data matches passed in data
		self.assertEqual(len(a[0][1]), len(val))

		# that one share isn't enough
		self.assertRaises(ValueError, recover_data, [ a[0] ], 2)

		self.assertEqual(val, recover_data(a[:2], 2))

	def test_gf2p8_inv(self):

		a = GF2p8(random.randint(0, 255))

		with unittest.mock.patch.object(GF2p8, '_invcache', {}) as pinvc:
			ainv = a ** -1

			self.assertEqual(a * ainv, 1)

			invcache = { x: int(GF2p8(x) ** -1) for x in range(1, 256) }

		if GF2p8._invcache != invcache:
			print('inv cache:', repr(invcache))
			self.assertEqual(GF2p8._invcache, invcache)

	def test_gf2p8_power(self):
		a = GF2p8(random.randint(0, 255))

		v = GF2p8(1)
		for i in range(10):
			self.assertEqual(a ** i, v)

			v = v * a

		for i in range(10):
			a = GF2p8(random.randint(0, 255))

			powers = a.powerseries(10)
			for j in range(10):
				self.assertEqual(powers[j], a ** j)

	def test_gf2p8(self):
		for i in range(10):
			a = GF2p8(random.randint(0, 255))
			b = GF2p8(random.randint(0, 255))
			c = GF2p8(random.randint(0, 255))

			self.assertEqual(a * 0, 0)

			# Identity
			self.assertEqual(a + 0, a)
			self.assertEqual(a * 1, a)
			self.assertEqual(0 + a, a)
			self.assertEqual(1 * a, a)
			self.assertEqual(0 - a, a)

			# Associativity
			self.assertEqual((a + b) + c, a + (b + c))
			self.assertEqual((a * b) * c, a * (b * c))

			# Communitative
			self.assertEqual(a + b, b + a)
			self.assertEqual(a * b, b * a)

			# Distributive
			self.assertEqual(a * (b + c), a * b + a * c)
			self.assertEqual((b + c) * a, b * a + c * a)

			# Basic mul
			self.assertEqual(GF2p8(0x80) * 2, 0x87)
			self.assertEqual(GF2p8(0x80) * 6, (0x80 * 6) ^ (0x187 << 1))
			self.assertEqual(GF2p8(0x80) * 8, (0x80 * 8) ^ (0x187 << 2) ^ (0x187 << 1) ^ 0x187)

			self.assertEqual(a + b - b, a)