from hashlib import sha256,md5
import socketserver
import signal
import string
import random
import os


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 proof_of_work1(self):
        random.seed(os.urandom(8))
        proof = ''.join(
            [random.choice(string.ascii_lowercase+string.digits) for _ in range(20)])
        _hexdigest = md5(proof.encode()).hexdigest()
        self.send(f"[+] md5(XXXXX+{proof[5:]}) == {_hexdigest}".encode())
        x = self.recv(prompt=b'[+] Plz tell me XXXXX: ')
        if len(x) != 5 or md5(x+proof[5:].encode()).hexdigest() != _hexdigest:
            return False
        return True
    
    def proof_of_work2(self):
        random.seed(os.urandom(8))
        proof = ''.join(
            [random.choice(string.ascii_letters+string.digits) for _ in range(20)])
        _hexdigest = md5(proof.encode()).hexdigest()
        self.send(f"[+] md5(XXXX+{proof[4:]}) == {_hexdigest}".encode())
        self.send(b"[+] I believe you kown what XXXX is. Try to prove it. ")
        self.send(b"[+] I need you give me a strings and md5(XXXX+strings) ")
        x = self.recv(prompt=b"[+] Give me md5: ")
        y = self.recv(prompt=b"[+] Give me strings: ")
        if md5(proof[:4].encode() + y).hexdigest().encode() != x:
            return False
        return True
    
    def proof_of_work3(self):
        random.seed(os.urandom(8))
        proof = ''.join(
            [random.choice(string.ascii_letters+string.digits) for _ in range(20)])
        _hexdigest = md5(proof.encode()).hexdigest()
        self.send(f"[+] md5(XXXXXXXXX+{proof[8:]}) == {_hexdigest}".encode())
        self.send(b"[+] I believe you kown what XXXXXXXX is. Try to prove it. ")
        self.send(b"[+] I need you give me a strings and md5(XXXXXXXX+strings) ")
        self.send(b"[+] whatsmore, the suffix of string must be 'ilove0xfa'")
        x = self.recv(prompt=b"[+] Give me md5: ")
        y = self.recv(prompt=b"[+] Give me strings: ")
        if y[-9:] != b"ilove0xfa" or md5(proof[:8].encode() + y).hexdigest().encode() != x:
            return False
        return True

    def handle(self):
        # signal.alarm(60)
        if not self.proof_of_work1():
            self.send(b'[!] Wrong!')
            return
        if not self.proof_of_work2():
            self.send(b'[!] Wrong!')
            return
        if not self.proof_of_work3():
            self.send(b'[!] Wrong!')
            return

        self.send(b'here is your flag')
        self.send(flag)


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


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


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