# GNU GPLv3, see LICENSE

from argparse import ArgumentParser

from shamira import generate, reconstruct

def run():

	parser=ArgumentParser()

	subparsers=parser.add_subparsers()

	buildSplitParser(subparsers.add_parser("split"))

	buildJoinParser(subparsers.add_parser("join"))

	args=parser.parse_args()

	args.func(args)

def buildSplitParser(parser):

	parser.add_argument("-k",type=int,required=True,help="number of shares necessary for recovering the secret")

	parser.add_argument("-n",type=int,required=True,help="number of generated shares")

	encoding=parser.add_mutually_exclusive_group()

	encoding.add_argument("--hex",action="store_true",help="encode shares' bytes as a hexadecimal string")

	encoding.add_argument("--b32",action="store_true",help="encode shares' bytes as a base32 string")

	encoding.add_argument("--b64",action="store_true",help="encode shares' bytes as a base64 string")

	parser.add_argument("secret",help="secret to be parsed")

	parser.set_defaults(func=_generate)

def buildJoinParser(parser):

	encoding=parser.add_mutually_exclusive_group()

	encoding.add_argument("--hex",action="store_true",help="decode shares' bytes from a hexadecimal string")

	encoding.add_argument("--b32",action="store_true",help="decode shares' bytes from a base32 string")

	encoding.add_argument("--b64",action="store_true",help="decode shares' bytes from a base64 string")

	parser.add_argument("-r","--raw",action="store_true",help="return secret as raw bytes")

	parser.add_argument("share",nargs="+",help="shares to be joined")

	parser.set_defaults(func=_reconstruct)

def _generate(args):

	encoding=getEncoding(args) or "b32"

	shares=generate(args.secret,args.k,args.n,encoding)

	for s in shares:

		print(s)

def _reconstruct(args):

	encoding=getEncoding(args)

	print(reconstruct(*args.share,encoding=encoding,raw=args.raw))

def getEncoding(args):

	if args.hex: return "hex"

	elif args.b32: return "b32"

	elif args.b64: return "b64"

	else: return ""

def reconstructRaw(*shares):

	"""Tries to recover the secret from its shares.

	:param shares: ((i, (bytes) share), ...)

	:return: (bytes) reconstructed secret. Too few shares returns garbage."""

	secretLen=len(shares[0][1])

	res=[None]*secretLen

	for i in range(secretLen):

		points=[(x,s[i]) for (x,s) in shares]

		res[i]=(getConstantCoef(*points))

	return bytes(res)

def reconstruct(*shares):

	"""Wraps reconstructRaw.

	:param shares: ((str) share, ...)

	:return: (str) reconstructed secret. Too few shares returns garbage."""

	bs=reconstructRaw(*(decode(s) for s in shares))

	return bs.decode(encoding="utf-8")

def decode(share):

	try:

		(i,_,shareStr)=share.partition(".")

		i=int(i)

		if not 1<=i<=255:

			raise ValueError()

		shareBytes=base64.b32decode(shareStr)

		return (i,shareBytes)

	except (ValueError,binascii.Error):

		raise ValueError('bad share format: share="{0}"'.format(share))

###

from argparse import ArgumentParser

def run():

	parser=ArgumentParser()

	subparsers=parser.add_subparsers()

	joiner=subparsers.add_parser("join")

	joiner.add_argument("share",nargs="+",help="shares to be joined")

	parser.set_defaults(func=lambda: parser.error("missing command"))

	args=parser.parse_args()

	args.func(args)

if __name__=="__main__":

	run()

# GNU GPLv3, see LICENSE

import os

import re

import base64

import binascii

import gf256

def _shareByte(secretB,k,n):

	coefs=[int(secretB)]+[int(b) for b in os.urandom(k-1)]

	points=[gf256.evaluate(coefs,i) for i in range(1,n+1)]

	return points

def generateRaw(secret,k,n):

	"""Splits secret into shares.

	:param secret: (bytes)

	:param k: number of shares necessary for secret recovery. 1 <= k <= n

	:param n: (int) number of shares generated. 1 <= n < 255

	:return: [(i, (bytes) share), ...]"""

	shares=[_shareByte(b,k,n) for b in secret]

	return [(i+1, bytes([s[i] for s in shares])) for i in range(n)]

def reconstructRaw(*shares):

	"""Tries to recover the secret from its shares.

	:param shares: ((i, (bytes) share), ...)

	:return: (bytes) reconstructed secret. Too few shares returns garbage."""

	secretLen=len(shares[0][1])

	res=[None]*secretLen

	for i in range(secretLen):

		points=[(x,s[i]) for (x,s) in shares]

		res[i]=(gf256.getConstantCoef(*points))

	return bytes(res)

def generate(secret,k,n,encoding="b32"):

	"""Wraps generateRaw().

	:param secret: (str or bytes)

	:param k: number of shares necessary for secret recovery

	:param n: number of shares generated

	:param encoding: {hex, b32, b64} desired output encoding. Hexadecimal, Base32 or Base64.

	:return: [(str) share, ...]"""

	if isinstance(secret,str):

		secret=secret.encode("utf-8")

	shares=generateRaw(secret,k,n)

	return [encode(s,encoding) for s in shares]

def reconstruct(*shares,encoding="",raw=False):

	"""Wraps reconstructRaw.

	:param shares: ((str) share, ...)

	:param encoding: {hex, b32, b64, ""} encoding of share strings. If not provided or empty, the function tries to guess it.

# GNU GPLv3, see LICENSE

import random

from unittest import TestCase

from shamira import _shareByte

from shamira import *

class TestShamira(TestCase):

	_urandom=os.urandom

	@classmethod

	def setUpClass(cls):

		random.seed(17)

		os.urandom=lambda n: bytes(random.randint(0,255) for i in range(n))

	@classmethod

	def tearDownClass(cls):

		os.urandom=cls._urandom

	def test_shareByte(self):

			_shareByte(b"a",5,4)

			_shareByte(b"a",5,255)

		with self.assertRaises(ValueError): # not castable to int

			_shareByte("x",2,3)

		vals=_shareByte(ord(b"a"),2,3)

		points=list(zip(range(1,256), vals))

		self.assertEqual(gf256.getConstantCoef(*points), ord(b"a"))

		self.assertEqual(gf256.getConstantCoef(*points[:2]), ord(b"a"))

		self.assertNotEqual(gf256.getConstantCoef(*points[:1]), ord(b"a")) # underdetermined => random

	def testGenerateReconstructRaw(self):

		for (k,n) in [(2,3), (254,254)]:

			shares=generateRaw(b"abcd",k,n)

			random.shuffle(shares)

			self.assertEqual(reconstructRaw(*shares[:k]), b"abcd")

			self.assertNotEqual(reconstructRaw(*shares[:k-1]), b"abcd")

	def testGenerateReconstruct(self):

		for encoding in ["hex","b32","b64"]:

			for secret in [b"abcd","abcde","ěščřžý"]:

				for (k,n) in [(2,3), (254,254)]:

					raw=isinstance(secret,bytes)

					with self.subTest(enc=encoding,r=raw,sec=secret,k=k,n=n):

						shares=generate(secret,k,n,encoding)

						random.shuffle(shares)

						self.assertEqual(reconstruct(*shares[:k],encoding=encoding,raw=raw), secret)

						self.assertEqual(reconstruct(*shares[:k],raw=raw), secret)

						s=secret if raw else secret.encode("utf-8")

						self.assertNotEqual(reconstruct(*shares[:k-1],encoding=encoding,raw=True), s)

		shares=generate(b"\xfeaa",2,3)

		with self.assertRaises(UnicodeDecodeError):

			reconstruct(*shares)

	def testEncode(self):

		share=(2,b"\x00\x01\x02")

		for (encoding,encodedStr) in [("hex",'000102'),("b32",'AAAQE==='),("b64",'AAEC')]:

			with self.subTest(enc=encoding):

				self.assertEqual(encode(share,encoding), "2."+encodedStr)

	def testDecode(self):

		with self.assertRaises(ValueError):

			decode("AAA")

			decode("1.")

			decode(".AAA")

			decode("1AAA")

			decode("1.0001020f","hex")

			decode("1.000102030","hex")

			decode("1.AAAQEAY")