# Alice and Bob are close together, likely because they have a lot of things in common.
# This is why Alice asked him a small *q*uestion, about something cooler than a wiener.

import Crypto
import shutil
import base64
import subprocess
from Crypto.PublicKey import RSA
import os
import zipfile
import math

FLAG = "BKPCTF{Its_not_you,_its_rsa_(that_is_broken)}"

def fermatfactor(N):  # http://facthacks.cr.yp.to/fermat.html
    if N <= 0: return [N]
    if is_even(N): return [2,N/2]
    a = ceil(sqrt(N))
    while not is_square(a^2-N):
      a = a + 1
    b = sqrt(a^2-N)
    return [a - b,a + b]

def bruteforce(N):
  s = 2 ^ 33
  for i in xrange(3, s, 2):
    if N % i == 0:
      return i
  return 0


def rsa_recover_prime_factors(n, e, d):
    """
    Compute factors p and q from the private exponent d. We assume that n has
    no more than two factors. This function is adapted from code in PyCrypto.
    """
    # See 8.2.2(i) in Handbook of Applied Cryptography.
    ktot = d * e - 1
    # The quantity d*e-1 is a multiple of phi(n), even,
    # and can be represented as t*2^s.
    t = ktot
    while t % 2 == 0:
        t = t // 2
    # Cycle through all multiplicative inverses in Zn.
    # The algorithm is non-deterministic, but there is a 50% chance
    # any candidate a leads to successful factoring.
    # See "Digitalized Signatures and Public Key Functions as Intractable
    # as Factorization", M. Rabin, 1979
    spotted = False
    a = 2
    while not spotted and a < 10000:
        k = t
        # Cycle through all values a^{t*2^i}=a^k
        while k < ktot:
            cand = power_mod(a, k, n)
            # Check if a^k is a non-trivial root of unity (mod n)
            if cand != 1 and cand != (n - 1) and power_mod(cand, 2, n) == 1:
                # We have found a number such that (cand-1)(cand+1)=0 (mod n).
                # Either of the terms divides n.
                p = gcd(cand + 1, n)
                spotted = True
                break
            k *= 2
        # This value was not any good... let's try another!
        a += 2
    if not spotted:
        raise ValueError("Unable to compute factors p and q from exponent d.")
    # Found !
    q, r = divmod(n, p)
    assert r == 0

    return (p, q)

def generate_pubkeys():
    e = 65537
    p1 = random_prime(2^512)
    q1 = next_prime(p1)
    N1 = p1 * q1
    d1 = inverse_mod(e, (p1-1)*(q1-1))
    phi1 = (p1-1)*(q1-1)

    if set([q1, p1]) == set(fermatfactor(N1)) and gcd(phi1, e) == 1:
        print '[+] Close primes generated'
        key = RSA.construct((long(N1),long(e),long(d1)))
        try:
            os.remove('close.pub')
        except OSError:
            pass
        with open('close.pub', 'w') as f:
            f.write(key.publickey().exportKey())


    p2 = random_prime(2^512)
    q2 = q1
    N2 = p2 * q2
    e = 65537
    d2 = inverse_mod(e, (p2-1)*(q2-1))
    phi2 = (p2-1)*(q2-1)
    if gcd(N1, N2) == q2 and set([p2,q2]) == set(rsa_recover_prime_factors(N2,e,d2)) and gcd(phi2, e) == 1:
        print '[+] Shared factor generated'
        key = RSA.construct((long(N2),long(e),long(d2)))
        try:
            os.remove('shared.pub')
        except OSError:
            pass
        with open('shared.pub', 'w') as f:
            f.write(key.publickey().exportKey())

    e = 65537
    p3 = random_prime(2^16)
    q3 = random_prime(2^1024)
    N3 = p3 * q3
    d3 = inverse_mod(e, (p3-1)*(q3-1))
    phi3 = (p3-1)*(q3-1)
    if bruteforce(N3) in [q3, p3] and set([p3,q3]) == set(rsa_recover_prime_factors(N3,e,d3)) and gcd(phi3, e) == 1:
        print '[+] Small p generated'
        key = RSA.construct((long(N3),long(e),long(d3)))
        try:
            os.remove('small.pub')
        except OSError:
            pass
        with open('small.pub', 'w') as f:
            f.write(key.publickey().exportKey())

    p4 = 10114792273660656874618568712406420344176220457790563178092222929337786916374923318745284718351487926620784106195715878875311958793629905453919697155685507
    q4 = 10843221374140991753173625949764386011485161421520044246309105053489500519257941272796681417497061734054081478280518835582353321569961722963922828311576983
    N4 = p4*q4
    e = 49446678600051379228760906286031155509742239832659705731559249988210578539211813543612425990507831160407165259046991194935262200565953842567148786053040450198919753834397378188932524599840027093290217612285214105791999673535556558448523448336314401414644879827127064929878383237432895170442176211946286617205
    d4 = 21780352155588618020563641971337344243907391969899764877790673891831527301137
    phi4 = (p4-1)*(q4-1)

    X = ceil(N4^(1/4)/2)
    Y = ceil(N4^(1/2))
    M = matrix([[X*Y, -X*(N4+1), -1], [0, e*X, 0], [0,0,e]])
    B = M.LLL()

    if (B[0][2]*(N4+1)-B[0][1]/X)/e == d4 and set([p4,q4]) == set(rsa_recover_prime_factors(N4,e,d4)) and gcd(phi4, e) == 1:
        print '[+] Boneh Dufree generated'
        key = RSA.construct((long(N4),long(e),long(d4)))
        try:
            os.remove('bonnet.pub')
        except OSError:
            pass
        with open('bonnet.pub', 'w') as f:
            f.write(key.publickey().exportKey())

def make_challenge():
    with open('bonnet.pub', 'r') as k:
        pubkey = Crypto.PublicKey.RSA.importKey(k.read())

    passwd_bonnet = base64.b64encode(os.urandom(24))
    ciphertext = pubkey.encrypt(passwd_bonnet, None)[0]
    with open('almost_there.encrypted', 'w') as f:
        f.write(ciphertext)

    with open('FLAG', 'w') as f:
        f.write(FLAG)

    rc = subprocess.call(['7z', 'a', '-p%s' % passwd_bonnet, '-y', 'almost_there.zip'] + ['FLAG'])
    shutil.copy('bonnet.pub', 'almost_there.pub')

    ###

    with open('small.pub', 'r') as k:
        pubkey = Crypto.PublicKey.RSA.importKey(k.read())

    passwd_small = base64.b64encode(os.urandom(24))
    ciphertext = pubkey.encrypt(passwd_small, None)[0]
    with open('almost_almost_there.encrypted', 'w') as f:
        f.write(ciphertext)

    rc = subprocess.call(['7z', 'a', '-p%s' % passwd_small, '-y', 'almost_almost_there.zip'] + ['almost_there.zip', 'almost_there.encrypted', 'almost_there.pub'])
    shutil.copy('small.pub', 'almost_almost_there.pub')

    ###

    with open('shared.pub', 'r') as k:
        pubkey = Crypto.PublicKey.RSA.importKey(k.read())

    passwd_shared = base64.b64encode(os.urandom(24))
    ciphertext = pubkey.encrypt(passwd_shared, None)[0]
    with open('almost_almost_almost_there.encrypted', 'w') as f:
        f.write(ciphertext)

    rc = subprocess.call(['7z', 'a', '-p%s' % passwd_shared, '-y', 'almost_almost_almost_there.zip'] + ['almost_almost_there.zip', 'almost_almost_there.encrypted', 'almost_almost_there.pub'])
    shutil.copy('shared.pub', 'almost_almost_almost_there.pub')

    ###

    with open('close.pub', 'r') as k:
        pubkey = Crypto.PublicKey.RSA.importKey(k.read())

    passwd_close = base64.b64encode(os.urandom(24))
    ciphertext = pubkey.encrypt(passwd_close, None)[0]
    with open('almost_almost_almost_almost_there.encrypted', 'w') as f:
        f.write(ciphertext)

    rc = subprocess.call(['7z', 'a', '-p%s' % passwd_close, '-y', 'almost_almost_almost_almost_there.zip'] + ['almost_almost_almost_there.zip', 'almost_almost_almost_there.encrypted', 'almost_almost_almost_there.pub'])
    shutil.copy('close.pub', 'almost_almost_almost_almost_there.pub')

    rc = subprocess.call(['7z', 'a', '-y', 'almost_almost_almost_almost_almost_there.zip'] + ['almost_almost_almost_almost_there.zip', 'almost_almost_almost_almost_there.encrypted', 'almost_almost_almost_almost_there.pub'])

    print '[+] Archive created!'

def solve_challenge():
    os.system('rm -Rf ./tmp')
    os.system('rm -Rf ./tmp2')
    os.system('rm -Rf ./tmp3')
    os.system('rm -Rf ./tmp4')
    os.system('rm -Rf ./tmp5')

    with zipfile.ZipFile("almost_almost_almost_almost_almost_there.zip", mode='r') as z:
        z.extractall('./tmp')
    with open('./tmp/almost_almost_almost_almost_there.pub', 'r') as f:
        pubkey = Crypto.PublicKey.RSA.importKey(f.read())
    print '[+] factorizing...'
    (p,q) = fermatfactor(pubkey.n)
    print '[+] factorization done!'
    d = inverse_mod(pubkey.e, (p-1)*(q-1))
    privkey = RSA.construct((long(pubkey.n),long(pubkey.e),long(d)))
    with open('./tmp/almost_almost_almost_almost_there.encrypted', 'r') as f:
        p = privkey.decrypt(f.read())
        zipfile.ZipFile('almost_almost_almost_almost_there.zip').extractall('./tmp2', pwd=p)


    with open('./tmp2/almost_almost_almost_there.pub', 'r') as f:
        _pubkey = Crypto.PublicKey.RSA.importKey(f.read())
    print '[+] factorizing...'
    _q = gcd(pubkey.n, _pubkey.n)
    print '[+] factorization done!'
    d = inverse_mod(_pubkey.e, (_pubkey.n / _q-1)*(_q-1))
    privkey = RSA.construct((long(_pubkey.n),long(_pubkey.e),long(d)))
    with open('./tmp2/almost_almost_almost_there.encrypted', 'r') as f:
        p = privkey.decrypt(f.read())
        zipfile.ZipFile('./tmp2/almost_almost_almost_there.zip').extractall('./tmp3', pwd=p)


    with open('./tmp3/almost_almost_there.pub', 'r') as f:
        pubkey = Crypto.PublicKey.RSA.importKey(f.read())
    print '[+] factorizing 236...'
    p = bruteforce(pubkey.n)
    q = pubkey.n / p
    print '[+] factorization done!'
    d = inverse_mod(pubkey.e, (p-1)*(q-1))
    privkey = RSA.construct((long(pubkey.n),long(pubkey.e),long(d)))
    with open('./tmp3/almost_almost_there.encrypted', 'r') as f:
        p = privkey.decrypt(f.read())
        zipfile.ZipFile('./tmp3/almost_almost_there.zip').extractall('./tmp4', pwd=p)

    with open('./tmp4/almost_there.pub', 'r') as f:
        pubkey = Crypto.PublicKey.RSA.importKey(f.read())
    print '[+] factorizing 248...'
    X = ceil(pubkey.n^(1/4)/2)
    Y = ceil(pubkey.n^(1/2))
    M = matrix([[X*Y, -X*(pubkey.n+1), -1], [0, pubkey.e*X, 0], [0,0,pubkey.e]])
    B = M.LLL()
    d = (B[0][2]*(pubkey.n+1)-B[0][1]/X)/pubkey.e
    print '[+] factorization done!'
    privkey = RSA.construct((long(pubkey.n),long(pubkey.e),long(d)))
    with open('./tmp4/almost_there.encrypted', 'r') as f:
        p = privkey.decrypt(f.read())
        zipfile.ZipFile('./tmp4/almost_there.zip').extractall('./tmp5', pwd=p)


#solve_challenge()
generate_pubkeys()
make_challenge()