#!/usr/bin/env python3
from Cryptotools.Groups.cyclic import Cyclic
from Cryptotools.Numbers.primeNumber import getPrimeNumber
from random import randint, choice
from math import log, log10

"""
    Here, we will try to understand why we need to have a generator when we encrypt data for Diffie-Hellman
    https://crypto.stackexchange.com/questions/25489/why-does-diffie-hellman-need-be-a-cyclic-group
"""
def operation(a, b, n):
    return (a ** b) % n

def getGenerator(gr, p):
    index = 1
    for g in range(2, p):
        z = list()
        for entry in range(1, p):
            res = operation(g, index, p)
            #if res not in z:
            z.append(res)
            index = index + 1
        print(f"{g}: {sorted(z)}")

def getGenerator2(gr, p, g):
    index = 1
    z = list()
    for entry in range(1, p):
        res = operation(g, index, p)
        if res not in z:
            z.append(res)
            index = index + 1
    return z
    

def computePublicKey(key, p, g):
    return (g ** key) % p

def computeEphemeralKey(public, secret, p):
    return (public ** secret) % p

gr = list()
# Public value
p = 43
g = 0
for i in range(1, p):
    gr.append(i)

print(f"p = {p}")
print(f"G = {gr}")
# We try with a generator which is not in list
cyclic = Cyclic(gr, p, operation)
generators = cyclic.getGenerators()
print(f"All generators: {generators}")

g = 3 # In the group
print(f"g = {g}")

# We can compute with the secret key
secretKeyA = 5
secretKeyB = 10
publicKeyA = computePublicKey(secretKeyA, p, g)
publicKeyB = computePublicKey(secretKeyB, p, g)
print(f"Public key A: {publicKeyA}")
print(f"Public key B: {publicKeyB}\n")

# Here, g is in the generator, we need to compute all values until that match with publicKeyA
for a in range(1, p):
    res = operation(g, a, p)
    if res == publicKeyA:
        print(f"Brute forced secret key of A: {a}")

for a in range(1, p):
    res = operation(g, a, p)
    if res == publicKeyB:
        print(f"Brute forced secret key of B: {a}")

print()

#for a in generators:
#    print(a)
#
#print()

# Here, we generate all key withthe same generator
keys = list()
for a in range(1, p):
    keys.append(operation(g, a, p))

print((keys))
print(sorted(keys))

print()

print(f"Keys with g = 4: {sorted(getGenerator2(gr, p, 4))}")

# Do the same, but g is not a generator
g = 4 # Not in the generator group
publicKeyA = computePublicKey(secretKeyA, p, g)
publicKeyB = computePublicKey(secretKeyB, p, g)
print(f"Public key A: {publicKeyA}")
print(f"Public key B: {publicKeyB}")

keys = list()
for a in range(1, p):
    keys.append(operation(g, a, p))
print(keys)

# Eve sniff the traffic and knows p, g and publicKeyA and B
# Eve, knows p and g, because it's public
# Eve need to guess the secretKey of A.
# For doing that, we iterate all posibility until that match with the publicKeyA

#for a in range(1, 4):
#    res = operation(g, a, p)
#    print(res)