# mini_demo.py
from Crypto.PublicKey import RSA, ECC
from Crypto.Hash import SHA256
from Crypto.Cipher import AES
from Crypto.Signature import pkcs1_15
from Crypto.Random import get_random_bytes as rb

def gen_role(n):
    r = RSA.generate(2048)
    return {"name": n, "priv": r.export_key(), "pub": r.publickey().export_key()}

def ecdh_key():
    a, b = ECC.generate(curve="P-256"), ECC.generate(curve="P-256")
    s = a._multiply(b.pointQ, a.d).x.to_bytes()
    return SHA256.new(s).digest()

def enc(k, m, aad=b""):
    n = rb(12); c = AES.new(k, AES.MODE_GCM, nonce=n); c.update(aad)
    t = c.encrypt_and_digest(m)[1]
    return n, c.encrypt(b""), t  # ciphertext unused; keeping API tiny

def dec(k, n, ct, t, aad=b""):
    c = AES.new(k, AES.MODE_GCM, nonce=n); c.update(aad)
    c.decrypt(b""); c.verify(t)

def sign(role, msg):
    h = SHA256.new(msg)
    return pkcs1_15.new(RSA.import_key(role["priv"])).sign(h)

def verify(pub, msg, sig):
    try:
        pkcs1_15.new(RSA.import_key(pub)).verify(SHA256.new(msg), sig); return True
    except: return False

def demo(s, r, msg):
    k = ecdh_key()
    hdr = f"{s['name']}->{r['name']}".encode()
    sig = sign(s, hdr)
    n, ct, t = enc(k, msg.encode(), hdr)
    assert verify(s["pub"], hdr, sig)
    dec(k, n, ct, t, hdr)
    print("OK:", msg)

if __name__ == "__main__":
    a, b = gen_role("A"), gen_role("B")
    demo(a, b, "secret")