from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from secret import flag
import socket
import os
import threading

KEY = os.urandom(16)

def encrypt(pt):
    cipher = AES.new(KEY, AES.MODE_ECB)
    enc = cipher.encrypt(pad(pt + flag, 16)).hex()
    return enc

def handle_client(connection):
    try:
        while True:
            connection.sendall(b'give me message : ')
            chunk = connection.recv(2048).replace(b'\n', b'')
            print(chunk)
            if not chunk:
                connection.sendall(b'You should give any string !!!')
                connection.close()
                break
            if chunk:
                connection.sendall(b'encrypted message : ')
                connection.sendall(encrypt(chunk).encode() + b'\n')
    except:
        connection.close()

def tcp_server(host, port):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
        server_socket.bind((host, port))
        server_socket.listen(5)
        while True:
            connection, _ = server_socket.accept()
            client_thread = threading.Thread(target=handle_client, args=(connection,))
            client_thread.start()

if __name__ == "__main__":
    HOST = "0.0.0.0"
    PORT = 523
    tcp_server(HOST, PORT)