#!/usr/bin/env python
import socket
from thread import start_new_thread
from datetime import datetime
from os import urandom
from random import randrange as random
import RSAvulnerableKeyGenerator as keygen

FLAG = 'SlashRootCTF{W13nn3r_w1nn3r_RSA_d1nn3r}'
HOST = '0.0.0.0'
PORT = 6070
BUFF = 1024
BIT = 256
MAX = 10

def banner():
    return '''\
_______________________________________________
|  ____ ____ ____    ___ ____ _  _ ____ _  _  |
|  |__/ [__  |__|     |  |  | |_/  |___ |\ |  |
|  |  \ ___] |  |     |  |__| | \_ |___ | \|  |
| ____ ____ _  _ ____ ____ ____ ___ ____ ____ |
| | __ |___ |\ | |___ |__/ |__|  |  |  | |__/ |
| |__] |___ | \| |___ |  \ |  |  |  |__| |  \ |
|_____________________________________________|                                          
| [1] Generate RSA                            |
| [2] Generate Token                          |
| [3] Generate Flag                           |
|_____________________________________________|
>>> '''

def log(message, address, filename='RSATG.log'):
    with open(filename, 'a') as log:
        timestamp = datetime.now().strftime('%d/%m/%Y %H:%M:%S')
        log.write('[%s][%s:%d] %s\n' % (timestamp, address[0], address[1], str(message)))

def serve_client(client, address, receive=''):
    try:
        client.send(banner())
        token = []

        while True:
            data = ''
            receive = client.recv(BUFF).strip()

            if receive == '1':
                log('<<< Generating RSA', address)
                data = genRSA(client, address)
            elif receive == '2':
                log('<<< Generating Token', address)
                data, token = genToken(client, address)
                log('>>> Sending Token: %s' % format(token), address)
            elif receive == '3':
                log('<<< Generating flag', address)
                
                if token:
                    client.send('Token : ')

                    if format(token) == client.recv(BUFF).strip():
                        ID = urandom(16).encode('hex')
                        log('<-> RSA ID: %s[%s]' % (ID, format(token)), address)
                        log('>>> Sending Flag: %s[%s]' % (ID, format(token)), address)
                        client.send('''\
RSA ID\t : %s
FLAG\t : %s
*Sertakan TOKEN dan RSA ID pada writeup agar poin dihitung!\n''' % (ID, FLAG))
                        break
                    else:
                        log('>-< Wrong Token: %s|%s' % (receive, format(token)), address)
                        data = 'Try Again!\nYour token is %s\n' % format(token)
                else:
                    log('>-< Empty Token!', address)
                    data = 'Generate your token!\n'

                token = []

            client.send(data + '[1|2|3]>>> ')

        log('Disconnected', address)
    except Exception as message:
        log(message, address, 'error.log')
        log('>-< Disconnected because error: %s' % message, address)
    finally:
        client.close()
        log('--- Disconnected', address)

def format(token):
    return '-'.join(token)

def RSA(bit):
    e, n, d = keygen.generateKeys(bit)
    p = random(1000, 9999)
    c = pow(p, e, n)
    return e, n, d, p, c

def genRSA(client, address):
    e, n ,d, p, c = RSA(BIT)
    client.send('e = %s \nn = %s \nc = %s \np = ' % (e, n, c))
    r = client.recv(BUFF).strip()

    if r == str(p):
        log('<-> Correct %s{e:%s,n:%s,d:%s,p:%s,c:%s)' % (r, e, n, d, p, c), address)
        return '\n\m/ Correct \m/\n'
    else:
        log('>-< Wrong %s{e:%s,n:%s,d:%s,p:%s,c:%s)' % (r, e, n, d, p, c), address)
        return '\nWrong :P is %s\n' % p

def genToken(client, address):
    token = []
    data = 'Token has been generated!\n'

    for i in range(1, 6):
        e, n ,d, p, c = RSA(BIT)
        client.send('Token #%i\ne = %s \nn = %s \nc = %s \np = ' % (i, e, n, c))
        r = client.recv(BUFF).strip()

        if r == str(p):
            log('<-> Correct #%d: %s{e:%s,n:%s,d:%s,p:%s,c:%s)' % (i, r, e, n, d, p, c), address)
            token.append(str(p))
        else:
            token = []
            data = '\nWrong :P is %s\n' % p
            log('>-< Wrong #%d: %s{e:%s,n:%s,d:%s,p:%s,c:%s)' % (i, r, e, n, d, p, c), address)
            break

    return data, token

def main():
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind((HOST, PORT))
    server.listen(MAX)
    log('(+) Power On Server', [HOST, PORT])

    while True:
        try:
            client, address = server.accept()
            log('<<< Client connected from IP %s with PORT %d' % (address), address)
            start_new_thread(serve_client, (client, address))
        except Exception as message:
            log(message, [HOST, PORT], 'error.log')
        except KeyboardInterrupt:
            log('(-) Power Off Server', [HOST, PORT])
            break

    server.close()

if __name__ == '__main__':
    main()