#sample code for the server
#author : fl4nk3r
import socketserver
import threading
import os
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad

HOST = ...
PORT = 1337
FLAG = b"CTK{...}" #put the flag 
KEY = os.urandom(16) 

def encrypt_oracle(plaintext):
    cipher = AES.new(KEY, AES.MODE_ECB)
    payload = plaintext + FLAG
    padded_payload = pad(payload, 16)
    ciphertext = cipher.encrypt(padded_payload)
    return ciphertext.hex()

class ChallengeHandler(socketserver.BaseRequestHandler):
    def handle(self):
        self.request.sendall(b"Welcome to the AES ECB Oracle Challenge!\n")
        self.request.sendall(b"Send me a string, and I will encrypt for you.\n")
        self.request.sendall(b"Format: Just send raw bytes or text. I return hex.\n\n")

        while True:
            try:
                self.request.sendall(b"> ")
                data = self.request.recv(1024).strip()

                if not data:
                    break
                encrypted_hex = encrypt_oracle(data)
                response = f"Ciphertext: {encrypted_hex}\n"
                self.request.sendall(response.encode())
            except Exception as e:
                self.request.sendall(b"Error processing request.\n")
                break

class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    pass