# Player 1:
#
#   $ sage multisig.sage genshare 2 3
#   A_0 = (17885028560402702239015261002188504514192120716688655005384557793085172279782, 8478670718018646753668209243232218399157692386519596227477013206445567901574)
#   A_1 = (24159996358010728217037819815568538754372914232633035350125672407860245980511, 25930880768349059332486494233731724414031492717763788767169325464072537905825)
#   
#   18245370559787825905107196417449809168129753681991950558452633836050485263625*X + Y + 21540401009192565430578225487575742717824479971010380831145688671749585986737
#   
#   Player 1's R = (1, 18110273049677706376100070599318402040771879310880963369761162988986654645832)
#   Player 2's R = (2, 28812924799218929326885620434040569836005182110830660190988271901329532330304)
#   Player 3's R = (3, 10567554239431103421778424016590760667875428428838709632535638065279047066679)
#
# Player 2:
#
#   $ sage multisig.sage genshare 2 3
#   A_0 = (802940932777049145807706375204159819724468245569326332695920733054960506640, 16355306126743352920399387200062438463930792131448355230670435463427127600465)
#   A_1 = (18141591414190824859071415735841428554530479628199222928098733408655160229999, 8126608045648459195646305845881115977993411928279311595258374101920586157117)
#   
#   11632470451635415873896303993978816754215391793499851090581681689882331027728*X + Y + 14165208049524082319699866276606579004249018988481566200129536690961988099823
#   
#   Player 1's R = (1, 3150343808169550662296575981586581204898645699960230088968524367549043820546)
#   Player 2's R = (2, 20465895665863183644293018239779741414046310388402026378066585426060075740915)
#   Player 3's R = (3, 8833425214227767770396714245800924659830918594902175287484903736177744713187)
#
# Player 3:
#
#   $ sage multisig.sage genshare 2 3
#   A_0 = (18050147058625614196833411623290233672504963362553094195996683497184802776885, 1558020177192412780876956779873498858051955811269538614817345479760770446460)
#   A_1 = (8859924332949089269591379272271847007387528856001855566494832087904814398546, 27961455601858187272958269886544614132704520355314277906405837204662338231452)
#   
#   22603695278713704898809499991696306162603661219853447180662052608597912268649*X + Y + 26665809323852841663944379970510953158263212308412307792093486582513151778711
#   
#   Player 1's R = (1, 8626540016091551149031612542136694605859239435617539786603946305675661848834)
#   Player 2's R = (2, 14970867046706895106114858802612365406618634697705739985621636445471112528282)
#   Player 3's R = (3, 21315194077322239063198105063088036207378029959793940184639326585266563207730)

# Each player has now generated the curve, shared its commits, and distributed shares.
# Now we recover the public key using all A0s
#
#   $ sage multisig.sage pubkey "(17885028560402702239015261002188504514192120716688655005384557793085172279782, 8478670718018646753668209243232218399157692386519596227477013206445567901574)" "(802940932777049145807706375204159819724468245569326332695920733054960506640, 16355306126743352920399387200062438463930792131448355230670435463427127600465)" "(18050147058625614196833411623290233672504963362553094195996683497184802776885, 1558020177192412780876956779873498858051955811269538614817345479760770446460)"
#
#   (9803495978299341257553881350441085748898862974053305000078308077444698847765 : 7803011951094511021525891181798443393652882333120996671357964800654859381546 : 1)

# To recover the shared secret, player's 1 and 2 will work together to
# recreate all 3 curves.
#
#   $ sage multisig.sage recover "(1, 18110273049677706376100070599318402040771879310880963369761162988986654645832)" "(2, 28812924799218929326885620434040569836005182110830660190988271901329532330304)"
#   21540401009192565430578225487575742717824479971010380831145688671749585986737
#
#   $ sage multisig.sage recover "(1, 3150343808169550662296575981586581204898645699960230088968524367549043820546)" "(2, 20465895665863183644293018239779741414046310388402026378066585426060075740915)"
#   14165208049524082319699866276606579004249018988481566200129536690961988099823
#
#   $ sage multisig.sage recover "(1, 8626540016091551149031612542136694605859239435617539786603946305675661848834)" "(2, 14970867046706895106114858802612365406618634697705739985621636445471112528282)"
#   26665809323852841663944379970510953158263212308412307792093486582513151778711
#
# You can see each command returns the constant coefficient in each curve.
# Finally we combine these to get the final curve and hence shared secret.
#
#   $ sage multisig.sage combine 14165208049524082319699866276606579004249018988481566200129536690961988099823 21540401009192565430578225487575742717824479971010380831145688671749585986737 26665809323852841663944379970510953158263212308412307792093486582513151778711
#   Secret: 4475373763911391702436979230349320953610598304020960064009226448437999969077
#   Pubkey: (9803495978299341257553881350441085748898862974053305000078308077444698847765 : 7803011951094511021525891181798443393652882333120996671357964800654859381546 : 1)
#
# We can see the public key matches what we got in the 'pubkey' step.

import argparse, base64, sys

q = 0x40000000000000000000000000000000224698fc0994a8dd8c46eb2100000001
K = GF(q)
P.<X, Y> = K[]

p = 0x40000000000000000000000000000000224698fc094cf91b992d30ed00000001
Fp = GF(p)
E = EllipticCurve(Fp, (0, 5))
G = E(
    23241645597038891398529199502196854108878665864265357905694087894995100434173,
    14702009283686283423048268817274882285027504402886079870290245450065579125215
)
assert G.order() == q

def genshare(args):
    t, n = args.t, args.n
    assert t <= n
    C = Y
    A = []
    for i in range(t):
        a_i = K.random_element()
        C += a_i * X^i
        A_i = a_i*G
        print(f"A_{i} = ({A_i[0]}, {A_i[1]})")
        A.append(A_i)

    print()
    print(C)
    print()

    R = []
    for j in range(1, n+1):
        x_j = j
        y_j = -C(X=x_j, Y=0)
        assert C(X=x_j, Y=y_j) == 0
        R_j = (x_j, y_j)
        print(f"Player {j}'s R = {R_j}")
        R.append(R_j)

    # Each player upon receiving their shares should perform this check
    def eval_C(x):
        P = E(0)
        for i, A_i in enumerate(A):
            P += x^i * A_i
        return P

    for R_j in R:
        x_j, y_j = R_j
        assert y_j*G + eval_C(x_j) == E(0)

def pubkey(args):
    P = E(0)
    for A0str in args.A0:
        x, y = A0str.split(",")
        x = x.strip("(")
        y = y.strip(") ")
        x, y = K(x), K(y)
        A0 = E(x, y)
        P += A0
    print(P)

def recover(args):
    R = []
    for Rjstr in args.Rj:
        x, y = Rjstr.split(",")
        x = x.strip("(")
        y = y.strip(") ")
        x, y = K(x), K(y)
        R_j = (x, y)
        R.append(R_j)

    # Create the Vandermonde matrix with a₀, …, aₜ₋₁ as indeterminates.
    V = []
    t = len(R)
    y = []
    for R_j in R:
        x_j, y_j = R_j
        y.append(y_j)

        V.append([x_j^i for i in range(t)])
    V = matrix(V)
    y = vector(y)

    a = V^-1 * -y
    print(a[0])

def combine(args):
    a0 = K(0)
    for a in args.a0:
        a0 += a
    print(f"Secret: {a0}")
    print(f"Pubkey: {a0*G}")

def main():
    parser = argparse.ArgumentParser(prog="multisig.sage")
    subparsers = parser.add_subparsers(required=True)

    parser_genshare = subparsers.add_parser("genshare", help="Generate a share")
    parser_genshare.add_argument("t", type=int, help="threshold for recovery")
    parser_genshare.add_argument("n", type=int, help="total players")
    parser_genshare.set_defaults(func=genshare)

    parser_pubkey = subparsers.add_parser("pubkey",
                                          help="Compute shared pubkey")
    parser_pubkey.add_argument("A0", nargs="+")
    parser_pubkey.set_defaults(func=pubkey)

    parser_recover = subparsers.add_parser("recover",
                                           help="Recover shared secret")
    parser_recover.add_argument("Rj", nargs="+")
    parser_recover.set_defaults(func=recover)

    parser_combine = subparsers.add_parser("combine",
                                           help="Combine shared secrets")
    parser_combine.add_argument("a0", type=int, nargs="+")
    parser_combine.set_defaults(func=combine)

    args = parser.parse_args()
    args.func(args)

main()