# -*- coding: utf-8 -*-
"""
Created on Tue Apr 23 13:03:18 2019

@author: Sceri
"""

from prime_number import generate_prime_number
import random

N_LENGTH = 32

def generateN():
    p = generate_prime_number(N_LENGTH // 2)
    q = generate_prime_number(N_LENGTH // 2)
    return p * q

def egcd(a, b):
        if a == 0:
            return (b, 0, 1)
        else:
            g, x, y = egcd(b % a, a)
            return (g, y - (b // a) * x, x)

def mulinv(b, n):
    g, x, _ = egcd(b, n)
    if g == 1:
        return x % n
    return -1

if __name__ == "__main__":
    n = generateN()
    print("N = " + str(n))
    k = int(input("Enter number of keys: "))

    s_vector = []
    v_vector = []
    for i in range(0, k):
        s = random.randint(1, n)
        v = mulinv(s * s, n)
        s_vector.append(s)
        v_vector.append(v)
    print("S Vector: ", s_vector)
    print("V Vector: ", v_vector)

    t = int(input("Enter number of rounds: "))
    for i in range(0, t):
        print("----------------- Round {0} -----------------".format(i))

        r = random.randint(1, n)
        x = pow(r, 2, n)

        print("A -> B: x = {0}".format(x))

        bool_vector = [random.choice([True, False]) for i in range(0, k)]
        print("B -> A: e: {0}".format(bool_vector))

        y = r
        for i in range(0, k):
            if bool_vector[i]:
                y = (y * s_vector[i]) % n
        print("A -> B: y: {0}".format(y))

        z = pow(y, 2, n)
        for i in range(0, k):
            if bool_vector[i]:
                z = (z * v_vector[i]) % n
        passed = x == z
        print("Passed: {0}".format(passed))

        if not passed:
            print("Test Failed. Hackerman detected")
            break
    if passed:
        print("Ok")