import ecpy.curves as curves
import secrets
import pickle
import Crypto.Random
import Crypto.Cipher.AES
import Crypto.Protocol.KDF

m0 = b"alice"
m1 = b"bob"
c = False

q = 224
curve = curves.Curve.get_curve('NIST-P224')
g = curve.generator

def H(p: curves.Point) -> bytes:
    secret = pickle.dumps((p.x, p.y), protocol=4)
    salt = Crypto.Random.get_random_bytes(16)
    key = Crypto.Protocol.KDF.scrypt(secret, salt, 16, N=2**14, r=8, p=1)
    return key[:32] # first 32 bytes of generated key

def E(key: bytes, message: bytes) -> tuple[bytes, bytes]:
    cipher = Crypto.Cipher.AES.new(key, Crypto.Cipher.AES.MODE_CTR)
    ct = cipher.encrypt(message)
    return (ct, cipher.nonce)

def D(key: bytes, encrypted_with_nonce: tuple[bytes, bytes]) -> bytes:
    ct, nonce = encrypted_with_nonce
    cipher = Crypto.Cipher.AES.new(key, Crypto.Cipher.AES.MODE_CTR, nonce=nonce)
    return cipher.decrypt(ct)

a = 1 + secrets.randbelow(q)
b = 1 + secrets.randbelow(q)

A = curve.mul_point(a, g)

B = curve.mul_point(b, g)
if c:
    B = curve.add_point(A, B)

k0 = H(curve.mul_point(a, B))
k1 = H(curve.mul_point(a, curve.sub_point(B, A)))
e0 = E(k0, m0)
e1 = E(k1, m1)

kc = H(curve.mul_point(b, A))
print(D(kc, e1 if c else e0))