everything is working exept for decryption

main.py

@@ -1,21 +1,44 @@
-from utils import generate_secret, coin_toss
+from utils import generate_secret, coin_toss, H
-from utils import prime, FILE
+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
+	return A_secret, a
 
 def bob_1(A_secret):
 	g = Integer(2)
 	b = generate_secret()
 	B_secret = pow(g, Integer.from_bytes(b), prime)
-	if coin_toss():
+	coin = coin_toss()
+	if coin:
 		B_secret = A_secret * B_secret
-	return B_secret
+	return B_secret, b, coin 
 
-a = alice_1()
+def alice_2(A_secret, B_secret, a, m0, m1):
-b = bob_1(a)
+	a = int.from_bytes(a)
-print(b)
+	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())

utils.py

@@ -1,11 +1,9 @@
 from Crypto.PublicKey import ECC
-from Crypto.Protocol.KDF import PBKDF2
+from Crypto.Protocol.KDF import scrypt
-from Crypto.Hash import SHA512
 from Crypto.Random import get_random_bytes
+from Crypto.Cipher import AES
 from secrets import randbelow
 
-FILE = 'exchange.bin'
-
 prime = 0xffffffffffffffffffffffffffffffff000000000000000000000001
 
 def generate_secret(c = 0):
@@ -15,8 +13,8 @@ def generate_secret(c = 0):
 def H(secret):
 	# secret should be bytearray[], pref. from generate_secret() function
 	salt = get_random_bytes(16)
-	key = PBKDF2(generate_secret(), salt, 64, count=1000000, hmac_hash_module=SHA512)
+	key = scrypt(bytes(secret), salt, 16, N=2**14, r=8, p=1)
-	return key[32:] # first 32 bytes of generated key
+	return key[:32] # first 32 bytes of generated key
 
 def coin_toss():
 	x = randbelow(2 ** 64)
@@ -24,3 +22,14 @@ def coin_toss():
 		return False
 	else:
 		return True
+
+def encrypt(key, data):
+	cipher = AES.new(key, AES.MODE_CTR)
+	ct = cipher.encrypt(data)
+	nonce = cipher.nonce
+	return ct, nonce
+
+def decrypt(ct, nonce, key):
+	cipher = AES.new(key, AES.MODE_CTR, nonce=nonce)
+	pt = cipher.decrypt(ct)
+	return pt