1

from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Util.number import inverse, bytes_to_long, long_to_bytes, isPrime
from sympy import nextprime, primerange
from decimal import Decimal, getcontext

def solve_layer2():
    print("[+] Starting Layer 2 decryption...")

    try:

        with open('public.pem', 'rb') as f:
            pubkey_pem = f.read()
        pubkey = RSA.import_key(pubkey_pem)
        n_rsa = pubkey.n
        e_rsa = pubkey.e

        with open('encrypted', 'rb') as f:
            encrypted_data = f.read()
    except FileNotFoundError as e:
        print(f"[!] Error: {e}. Make sure 'public.pem' and 'encrypted' are in the same directory.")
        return None, None

    print(f"  [-] Layer 2 n_rsa (modulus) loaded.")
    print(f"  [-] Layer 2 e_rsa (public exponent) = {e_rsa}")


    print("  [-] Searching for the 16-bit prime factor 'q'...")
    found_q = None
    for q_candidate in primerange(2, 2**16):
        if n_rsa % q_candidate == 0:
            found_q = q_candidate
            break

    if not found_q:
        print("[!] Failed to find the 16-bit prime factor q. Exiting.")
        return None, None

    q = found_q
    p = n_rsa // q
    print(f"  [-] Found 16-bit prime q = {q}")
    print(f"  [-] Calculated large prime p.")

    # Verify that p is indeed prime, as a sanity check.
    if not isPrime(p):
        print("[!] Calculated p is not prime. Something is wrong. Exiting.")
        return None, None
    print("  [-] Verified that p is prime.")

    # Reconstruct the private key for layer 2
    phi_rsa = (p - 1) * (q - 1)
    d_rsa = inverse(e_rsa, phi_rsa)

    private_key_layer2 = RSA.construct((n_rsa, e_rsa, d_rsa, p, q))
    print("  [-] Reconstructed Layer 2 private key.")


    cipher_decrypt = PKCS1_OAEP.new(private_key_layer2)
    key_size = private_key_layer2.size_in_bytes()

    encrypted_chunks = [encrypted_data[i:i + key_size] for i in range(0, len(encrypted_data), key_size)]

    try:
        decrypted_chunks = [cipher_decrypt.decrypt(chunk) for chunk in encrypted_chunks]
        decrypted_data = b''.join(decrypted_chunks)
    except ValueError as e:
        print(f"[!] Decryption failed. This may happen if the factored primes are incorrect. Error: {e}")
        return None, None

    ct_intermediate = bytes_to_long(decrypted_data)
    print("[+] Layer 2 decryption successful!")


    return ct_intermediate, p


def solve_layer1(ct_intermediate, p_for_key):
    """
    Recreates the large 'n' and 'phi' from the first encryption layer by
    reverse-engineering the starting prime from the p recovered from layer 2.
    """
    print("\n[+] Starting Layer 1 decryption...")
    print("  [-] Using the large prime p recovered from Layer 2 as our target.")


    getcontext().prec = 3000
    p_for_key_dec = Decimal(p_for_key)

    for num_primes_guess in [17, 16, 18]:
        p_start_approx = int(p_for_key_dec ** (Decimal(1)/Decimal(num_primes_guess)))
        print(f"\n  [-] Approximating based on {num_primes_guess} primes...")
        print(f"  [-] Approximate p_start is near {p_start_approx}")
        print(f"  [-] Searching for the correct 512-bit p_start around this value...")


        for i in range(-10000, 10000):
            p_start_cand = p_start_approx + i


            if p_start_cand.bit_length() != 512 or not isPrime(p_start_cand):
                continue


            p = p_start_cand
            n_candidate = 1
            phi_candidate = 1


            while True:
                n_candidate *= p
                phi_candidate *= (p - 1)
                if n_candidate.bit_length() >= 8192:
                    break
                p = nextprime(nextprime(p))

            # Now, check if the n we generated is the correct one.
            # The correct n is the one for which nextprime(n) equals the p we recovered.
            if nextprime(n_candidate) == p_for_key:
                print(f"\n  [+] SUCCESS! Found correct p_start: {p_start_cand}")
                print(f"  [+] Verified correct n_layer1 with bit length {n_candidate.bit_length()}")
                n_layer1 = n_candidate
                phi_layer1 = phi_candidate

                # With n_layer1 and phi_layer1, we can decrypt the first layer.
                e_layer1 = 65537

                print("  [-] Calculating private exponent d for Layer 1...")
                d_layer1 = inverse(e_layer1, phi_layer1)

                print("  [-] Decrypting final layer to get the flag...")
                ct_original = pow(ct_intermediate, d_layer1, n_layer1)

                flag = long_to_bytes(ct_original)

                print("\n[+] DECRYPTION COMPLETE!")
                try:
                    print(f"  [+] FLAG: {flag.decode()}")
                except UnicodeDecodeError:
                    print(f"  [+] FLAG (raw bytes): {flag}")
                return


    print("\n[!] Search for p_start failed. The approximation or search range may be off.")
    print("    Cannot proceed to decrypt Layer 1.")


def main():
    ct_intermediate, p_for_key = solve_layer2()
    if ct_intermediate is not None and p_for_key is not None:
        solve_layer1(ct_intermediate, p_for_key)

if __name__ == "__main__":
    main()