from hashlib import sha256
import socketserver
import signal
from flag import flag
from Crypto.Cipher import AES
import base64




class Task(socketserver.BaseRequestHandler):
    def _recvall(self):
        BUFF_SIZE = 2048
        data = b''
        while True:
            part = self.request.recv(BUFF_SIZE)
            data += part
            if len(part) < BUFF_SIZE:
                break
        return data.strip()

    def send(self, msg, newline=True):
        try:
            if newline:
                msg += b'\n'
            self.request.sendall(msg)
        except:
            pass

    def recv(self, prompt=b'[-] '):
        self.send(prompt, newline=False)
        return self._recvall()

    def task(self):
        key = sha256(flag).digest()
        key = key[:16]
        aes = AES.new(key, AES.MODE_ECB)
        data = self.recv(prompt=b'')
        data = b'cqupt'+ data + flag 
        # len(flag) == 21
        # wordlist = b'123456789qwertyuiopasdfghjklzxcvbnmQWERTYUIOPASDFGHJKLZXCVBNM{}_'
        if len(data)%16 != 0:
            pad = b'0' * (16 - len(data)%16) 
            data = data + pad
        cipher = aes.encrypt(data)
        print(aes.decrypt(cipher))
        cipher = base64.b64encode(cipher)
        self.send(cipher)
        

    
    def handle(self):
        signal.alarm(60)
        self.task()
            
        return


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


class ForkedServer(socketserver.ForkingMixIn, socketserver.TCPServer):
    pass


if __name__ == "__main__":
    #flag = bytes(os.getenv("FLAG"),"utf-8")
    HOST, PORT = '0.0.0.0', 10002
    server = ForkedServer((HOST, PORT), Task)
    server.allow_reuse_address = True
    print(HOST, PORT)
    server.serve_forever()