from playcrypt.simulator.hide_sim import HIDESim
from playcrypt.games.game_hide import GameHIDE
from playcrypt.simulator.bind_sim import BINDSim
from playcrypt.games.game_bind import GameBIND
from playcrypt.simulator.ufcma_sign_sim import UFCMASignSim
from playcrypt.games.game_ufcma_sign import GameUFCMASign
from playcrypt.primitives import *
from playcrypt.new_tools import *
from playcrypt.tools import *
import json
import sys
import os
import itertools

sys.path.append(os.path.abspath(os.path.join('..')))


def ADD(a,b):
    return a+b
def MULT(a,b):
    return a*b
def INT_DIV(a,N):
    return (a//N, a%N)
def MOD(a,N):
    return a%N
def EXT_GCD(a,N):
    return egcd(a,N)
def MOD_INV(a,N):
    res = modinv(a,N)
    if res == None:
        raise ValueError("Inverse does not exist.")
    return res
def MOD_EXP(a,n,N):
    return exp(a,n,N)

"""
!!!READ THIS FIRST!!!
Please name your collaborator(s) for this assignment. You can also state that
you have worked completely independently. Please refer to the course website
regarding the collaboration policy.

Your response: Independent


"""

"""
Problem 1 [8 points]
Let p be a prime of bit length k >= 8 such that (p - 1)/2 is also prime, and
let g be a generators of the group G = Z_p^*. Le q = p - 1. Consider the
digital signature scheme DS = (K, S, V) whose component algorithms are below,
where the message m is in Z_q^*.
"""

def K():
    x = random_Z_N(q)
    X = MOD_EXP(g, x, p)
    y = random_Z_N(q)
    Y = MOD_EXP(g, y, p)
    return (X, Y), (x, y)           # return (pk, sk)

def S(sk, m):
    """
    :param sk: The secret key sk = (x, y) used to sign messages
    :param M: The message to be signed, must be in Z_q^*
    :return: return the signature of message M
    """
    (x, y) = sk
    if EXT_GCD(m, q)[0] != 1:
        return False
    if m > q - 1 or m < 0:
        return False
    z = MOD(y + x * m, q)
    return z

def V1(pk, m, z):
    """
    This function is named V1 to differentiate from V2 in Problem 2
    :param pk: The public key pk = (X, Y) used to verify messages
    :param M: The message to be verified, must be in Z_q^*
    :param sig: The signature to be verified
    :return: return 1 if the signature is valid and 0 otherwise
    """
    (X, Y) = pk
    if EXT_GCD(m, q)[0] != 1:
        return 0
    if m > q - 1 or m < 0:
        return 0
    if z > q - 1 or z < 0:
        return 0
    if MOD(Y * MOD_EXP(X, m, p), p) == MOD_EXP(g, z, p):
        return 1
    return 0

"""
Present an O(k^2)-time adversary A making at most two queries to its Sign oracle and
achieving Adv^{uf-cma}_DS(A) = 1.
"""
# return valid message and signature
def A(Sign, pk):
    """
    You must fill in this method. This is the adversary that the problem is
    asking for. It should return a message and a signature.

    :param Sign:  Signing oracle
    :param pk:  Public key pk
    """

    (X, Y) = pk
    M0 = 1          # gcd (1, q) == 1; exists in Z_star(q)
    M1 = 3          # gcd(3, q) == 1; exists in Z_star(q)

    sig0 = Sign(M0)
    sig1 = Sign(M1)  

    new_msg = 5
    new_sig = MOD(sig1*2 - sig0, q)
    # == MOD( MOD(2y+6x, q) - MOD(y+x, q), q) == MOD(MOD(y+5x, q), q)
    # print(f"new_sig: {new_sig}")


    if(MOD(Y * MOD_EXP(X, new_msg, p), p) == MOD_EXP(g, new_sig, p)):
        return new_msg, new_sig
    else:
        new_msg = 7
        new_sig = MOD(sig1*3 - sig0*2, q)
        pass
    
    return new_msg, new_sig

"""
Problem 2 [12 points]
Let p be a prime of bit length k >= 8 such that (p - 1)/2 is also prime. Let g,
h be two different generators of the group G = Z_p^*. Let CS= (P, C, V) be the
commitment scheme whose consituent algorithms are as follows, where the message
M is in Z_{p-1}:
"""

def P():
    pi = (g, h)
    return pi

def C(pi, M):
    """
    :param pi: Public parameters
    :param M: The message to be commited, element of Z_{p-1}
    :return: return the commital and decommital key
    """
    (g, h) = pi
    K = random_Z_N(p-1)
    A = MOD_EXP(g, K, p)
    B = MOD_EXP(h, M, p)
    C_1 = MOD(A*B, p)
    C_2 = MOD(M+K, p-1)
    return ((C_1, C_2), K)

def V2(pi, C, M, K):
    """
    This function is named V2 to differentiate from V1 in Problem 1
    :param pi: Public parameters
    :param C: The commital
    :param M: The message to be verified
    :param K: The decommital key
    :return: return 1 if the opening is valid and 0 otherwise
    """
    (g, h) = pi
    (C_1, C_2) = C
    if not 0 <= K < p-1 or not 0 <= M < p-1:
        return 0
    A = MOD_EXP(g, K, p)
    B = MOD_EXP(h, M, p)
    C_1_prime = MOD(A*B, p)
    C_2_prime = MOD(M+K, p-1)
    if (C_1 == C_1_prime) and (C_2 == C_2_prime):
        return 1
    else:
        return 0

"""
Problem 2 Part 1. [6 points] Specify an O(k^3)-time adversary A1 making one query to its LR oracle and
achieving Adv^{hide}_CS(A1) = 1.
"""

def A1(lr, pi):
    """
    You must fill in this method. This is the adversary that the problem is
    asking for. It should return 0 or 1.

    :param lr: The oracle supplied by game HIDE
    :param pi: The public parameter pi
    """

    (g, h) = pi

    M0 = 1
    M1 = 2

    c= lr(M0, M1)           # returns (C_1, C_2)
    (C1, C2) = c
    K = C2 - 1

    A = MOD_EXP(g, K, p)
    B = MOD_EXP(h, M0, p)
    if(C1 == MOD(A*B, p) and C2 == MOD(M0 + K, (p-1))):
        return 0
    return 1

"""
Problem 2 Part 2. [6 points] Specify an O(k)-time adversary A2 such that Adv^{bind}_CS(A2) = 1.
(Hint: What is the value of g^{(p-1)/2} mod p, and why?)
"""

def A2(pi):
    """
    You must fill in this method. This is the adversary that the problem is
    asking for. It should return a tuple (C, M_0, M_1, K_0, K_1).

    :param pi: The public parameter pi
    """
    (g, h) = pi
    # print(f"g^(p-1)/2 mod p: {MOD_EXP(g, (p-1)//2, p)}")            # g^{(p-1)/2} mod p = p-1
    # Ciphertext C is the safe, K_0, K_1 are the keys
    # M_0, M_1 are messages
    # Use K_0 to reveal M_0, use K_1 to reveal M_1, have same Ciphertext


    # print(f"g^(M0-1)/2 mod p: {MOD_EXP(h, (p-1)//2, p)}")
    # C(pi, M)
    K_0 = (p-1)//2
    M_0 = (p-1)//2

    C_1 = MOD((p-1)*(p-1), p)           # = 1
    # print(f"C_1: {C_1}")
    C_2 = MOD(M_0+K_0, p-1)             # = 0
    # print(f"C_2: {C_2}")
    C = (C_1, C_2)

    M_1 = 0                         
    K_1 = 0
    A_1 = MOD_EXP(g, K_1, p)
    B_1 = MOD_EXP(h, M_1, p)
    C_1_check = MOD(A_1*B_1, p)             # = 1
    C_2_check = MOD(M_1+K_1, p-1)           # = 0
    if(C_1 != C_1_check or C_2 != C_2_check):
        print("Did not form the same ciphertext/commital")
        return None

    return ((C_1, C_2), M_0, M_1, K_0, K_1)


if __name__ == '__main__':
    # Problem 1
    # Sample random parameters
    k = 12
    print('Sampling random parameters of bit length k = %d' % k)
    p = random.randint(2**(k - 1), 2**k)
    while not is_prime(p) or not is_prime((p-1)//2):
        p = random.randint(2**(k - 1), 2**k)
    q = p-1
    g = random_Z_N_star(p)
    while (MOD_EXP(g, (p-1)//2, p) == 1) or (MOD_EXP(g, 2, p) == 1):
        g = random_Z_N_star(p)
    print('p = %d, g = %d' % (p, g))

    game_ufcma = GameUFCMASign(S, V1, 0, K)
    sim_ufcma = UFCMASignSim(game_ufcma, A)

    print("The advantage of your adversary A is approx. " + str(sim_ufcma.compute_advantage()))

    # Problem 2
    # Sample random parameters
    k = 12
    print('Sampling random parameters of bit length k = %d' % k)
    p = random.randint(2**(k - 1), 2**k)
    while not is_prime(p) or not is_prime((p-1)//2):
        p = random.randint(2**(k - 1), 2**k)
    g = random_Z_N_star(p)
    while (MOD_EXP(g, (p-1)//2, p) == 1) or (MOD_EXP(g, 2, p) == 1):
        g = random_Z_N_star(p)
    h = random_Z_N_star(p)
    while (h == g) or (MOD_EXP(h, (p-1)//2, p) == 1) or (MOD_EXP(h, 2, p) == 1):
        h = random_Z_N_star(p)
    print('p = %d, g = %d, h = %d' % (p, g, h))

    game_hide = GameHIDE(P, C)
    sim_hide = HIDESim(game_hide, A1)

    game_bind = GameBIND(P, V2)
    sim_bind = BINDSim(game_bind, A2)

    print("The advantage of your adversary A1 is approx. " + str(sim_hide.compute_advantage()))
    print("The advantage of your adversary A2 is approx. " + str(sim_bind.compute_advantage()))