from utils import generate_secret, coin_toss, H
from utils import prime
from utils import encrypt, decrypt
from Crypto.Math.Numbers import Integer

def alice_1():
	g = Integer(2)
	a = generate_secret()
	A_secret = pow(g, Integer(int.from_bytes(a)), prime)
	return A_secret, a

def bob_1(A_secret):
	g = Integer(2)
	b = generate_secret()
	B_secret = pow(g, Integer.from_bytes(b), prime)
	coin = coin_toss()
	if coin:
		B_secret = A_secret * B_secret
	return B_secret, b, coin 

def alice_2(A_secret, B_secret, a, m0, m1):
	a = int.from_bytes(a)
	k0 = H(str(pow(int(B_secret), a, prime)).encode())
	k1 = H(str(pow((B_secret // A_secret), a, prime)).encode())
	# Nonce at this point is generated by AES built-in method, not by scrypt
	e0, e1 = encrypt(k0, m0.encode()), encrypt(k1, m1.encode())
	return *e0, *e1

def bob_2(A_secret, b, c, e0, e1, n0, n1):
	kc = H(str(pow(int(A_secret), int.from_bytes(b), prime)).encode())
	if c:
		return decrypt(e1, n1, kc)
	else:
		return decrypt(e0, n0, kc)

m0, m1 = "alice", "bob"

A, a = alice_1()
B, b, c = bob_1(A)
e0, n0, e1, n1 = alice_2(A, B, a, m0, m1)
result = bob_2(A, b, c, e0, e1, n0, n1)


print(result.decode())