#!/usr/bin/env python2
import SocketServer
import socket
import Crypto.Random.random
import random
import gmpy2
import hashlib

# from secrets import FLAG


welcome = """Welcome to RSA sign and verify service.
Due to high costs of generating prime numbers you need to provide them yourself.
We are still secure because we will never disclose neither public nor private exponent to you.
"""

def send(s, t):
    s.sendall(t)

def recv(s):
    res = ''
    while '\n' not in res :
        x = s.recv(1)
        if len(x) == 0:
            break
        res += x
    return res

def main(s):
    send(s, welcome)
    send(s, "Please provide p=")
    p = int(recv(s).strip())
    send(s, "Please provide q=")
    q = int(recv(s).strip())

    assert p>2 and gmpy2.is_prime(p), "p check failed"
    assert q>2 and gmpy2.is_prime(q), "q check failed"
    assert p != q, "p is equal to q"

    n = p*q

    assert 2**511 < n and 2**512 > n, "n is not 512 bit"

    e = Crypto.Random.random.getrandbits(512)
    phi = (p-1)*(q-1)
    d = gmpy2.invert(e, phi)

    for i in range(20):
        plain = random.getrandbits(512)
        cipher = pow(plain, e, n)
        decipher = pow(cipher, d, n)
        assert plain == decipher, "rsa doesn't work for this key"


    send(s, "msg=")
    msg = recv(s).strip().decode("hex")
    send(s, "sign=")
    sign = int(recv(s))

    assert pow(sign, e, n) == int(hashlib.sha512(msg).hexdigest(), 16), "failed to verify"
    assert msg.strip() == "Give me the flag", "you don't want the flag"
    send(s, "FLAG:")


class TaskHandler(SocketServer.BaseRequestHandler):
    def handle(self):
        try:
          main(self.request)
        except Exception as err:
            send(self.request, err.message)

if __name__ == '__main__':
    SocketServer.ThreadingTCPServer.allow_reuse_address = True
    server = SocketServer.ThreadingTCPServer(('0.0.0.0', 1337), TaskHandler)
    server.serve_forever()