import os
import socketserver

p = 0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141
g = 2

b = int.from_bytes(os.urandom(32), "big")
B = pow(g, b, p)

FLAG = os.getenv("FLAG", "TEST_FLAG")

def encrypt(flag_bytes: bytes, shared_secret: int) -> bytes:
    key_byte = shared_secret
    return [b ^ key_byte for b in flag_bytes]

class Handler(socketserver.BaseRequestHandler):
    def handle(self) -> None:
        conn = self.request
        conn.sendall(f"Server public key: {B:x}\n".encode())
        conn.sendall(b"Send your public key A (hex):\n> ")
        line = self._readline(conn)
        if line is None:
            return
        line = line.strip()
        if not line:
            conn.sendall(b"Empty input\n")
            return
        try:
            A = int(line, 16)
        except ValueError:
            conn.sendall(b"Invalid hex\n")
            return
        shared_secret = pow(A, b, p)
        ciphertext = encrypt(FLAG.encode(), shared_secret)
        conn.sendall(f"Encrypted flag (hex): {ciphertext}\n".encode())

    @staticmethod
    def _readline(conn, limit: int = 4096):
        data = bytearray()
        while len(data) < limit:
            chunk = conn.recv(1)
            if not chunk:
                return None
            if chunk == b"\n":
                break
            if chunk != b"\r":
                data += chunk
        return data.decode(errors="ignore")


class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True


if __name__ == "__main__":
    host = "0.0.0.0"
    port = 5000
    with ThreadedTCPServer((host, port), Handler) as server:
        server.serve_forever()
