from cStringIO import StringIO
import os, random
from Crypto.Util import number


def egcd(a, b):
    if a == 0:
        return (b, 0, 1)
    else:
        g, y, x = egcd(b % a, a)
        return (g, x - (b // a) * y, y)

def modinv(a, m):
    g, x, y = egcd(a, m)
    if g != 1:
        raise Exception('modular inverse does not exist')
    else:
        return x % m

def randleetify(s, random):
    i = list(s)
    for c in range(len(i)):
        if random.random() > 0.5:
            if i[c] == 'i':
                if random.random() > 0.66:
                    i[c] = '1'
                elif random.random() > 0.33:
                    i[c] = 'I'
            elif i[c] == 'o':
                if random.random() > 0.66:
                    i[c] = '0'
                elif random.random() > 0.33:
                    i[c] = 'O'
            elif i[c] == 'a':
                if random.random() > 0.66:
                    i[c] = '4'
                elif random.random() > 0.33:
                    i[c] = 'A'
            elif i[c] == 'l':
                if random.random() > 0.66:
                    i[c] = '1'
                elif random.random() > 0.33:
                    i[c] = 'L'
            elif i[c] == 's':
                if random.random() > 0.66:
                    i[c] = '5'
                elif random.random() > 0.33:
                    i[c] = 'S'
            elif i[c] == 'e':
                if random.random() > 0.66:
                    i[c] = '3'
                elif random.random() > 0.33:
                    i[c] = 'E'
            elif i[c].isalpha():
                if random.random() > 0.5:
                    i[c] = i[c].upper()
        return ''.join(i)

flag = randleetify("i_probably_shouldnta_leeked_d",random)

def get_problem(random):
    n_length = 2048

    p = number.getPrime(n_length, os.urandom)
    q = number.getPrime(n_length, os.urandom)
    n = p*q
    e = 65537
    phi = (p-1)*(q-1)
    d = modinv(e, phi)
    cut = len(bin(d).strip('0b').strip('L'))//2
    given = bin(d).strip('0b').strip('L')[cut:]
    deez = hex(int(given,2)).zfill(len(given)/4)
    salt = random.choice("!?.")
    return (n, e, cut, deez, salt)

def generate_ciphertext(random):
    n, e, cut, deez, salt = get_problem(random)
    encoded = int(("easyctf{%s%s}" % (flag, salt)).encode('hex'),16)
    ciphertext =  'n: '+hex(n))+'\n'
    ciphertext += 'e: '+hex(e)+'\n'
    ciphertext += 'd: [??? '+str(cut)+' REDACTED BITS ???]'+str(deez)+'\n'
    ciphertext += 'c: '+hex(pow(encoded, e, n))+'\n'


    return StringIO(ciphertext)

def generate(random):
    return dict(files={
        "ciphertext.txt": generate_ciphertext
    })

def grade(random, key):
    n, e, deez, salt = get_problem(random)
    if key.find("%s_%s" % (flag, salt)) >= 0:
        return True, "Correct!"
    return False, "Nope."