import socketserver
import signal
import os
from Crypto.Util.number import * # type: ignore


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 question(self,flag):
        p = getPrime(512)
        q = getPrime(512)
        n = p * q
        e = 65537
        phi = (p - 1) * (q - 1)
        d = inverse(e,phi)
        flag = bytes_to_long(flag)
        c = pow(flag, e, n)
        self.send(b"Here is N")
        self.send(str(n).encode())
        self.send(b"You can't decrypt this!")
        self.send(str(c).encode())
        self.send(b"OMG d leaked!")
        self.send(str(d).encode())
        

    def handle(self):
        signal.alarm(60)
        self.question(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"testflag"
    HOST, PORT = '0.0.0.0', 10001
    server = ForkedServer((HOST, PORT), Task)
    server.allow_reuse_address = True
    print(HOST, PORT)
    server.serve_forever()