# socket, bind, listen, accept
from socket import *
from Cryptodome.PublicKey import RSA
from Cryptodome.Random import get_random_bytes
from Cryptodome.Cipher import AES, PKCS1_OAEP


skt = socket(AF_INET, SOCK_STREAM);
skt.bind(("0.0.0.0", 5150))
skt.listen(5)

while(1):
    cskt, caddr = skt.accept()
    print(caddr)
    # receive the public key from the client
    public_key_string = cskt.recv(1024) # Not a good practice, because we don't know the network MTU
    public_key = RSA.import_key(public_key_string);

    # generate a session key
    session_key = get_random_bytes(16)

    # encrypt the session key with the client's public key
    cipher_rsa = PKCS1_OAEP.new(public_key)
    enc_session_key = cipher_rsa.encrypt(session_key)
    
    # Send back the encrypted session key
    cskt.send(enc_session_key)

    # receive the encrypted message from the client
    nonce = cskt.recv(16)
    tag = cskt.recv(16)
    ciphertext = cskt.recv(1024) # Whatever is in the TCP receive buffer, and we'll hope that's all of it

    # decrypt the encrypted message from the client and print it out
    cipher_aes = AES.new(session_key, AES.MODE_EAX, nonce)
    data = cipher_aes.decrypt_and_verify(ciphertext, tag)
    print(data.decode("utf-8"))