Implement a quick and probably bad RSA keypair, and factor it with yafu.

rsa.py

#!/usr/bin/env python

from math import lcm
from Crypto.Util.number import long_to_bytes, bytes_to_long, getPrime as get_prime
from Crypto.PublicKey.RSA import construct, import_key
from base64 import b64encode, b64decode
import subprocess
import re


class RSA:
    def __init__(self, bit_depth):
        self.public_key, self.private_key = self.generate_keys(bit_depth)
        with open('private_key.pem', 'wb') as fp:
            fp.write(self.private_key.export_key())
        with open('public_key.pem', 'wb') as fp:
            fp.write(self.public_key.export_key())


    def generate_keys(self, bit_depth):
        e = 0x10001                                 # !! 1 < e < lambda_n && e is not a divisor of lambda_n
        p = get_prime(bit_depth // 2)
        q = get_prime(bit_depth // 2)
        lambda_n = self.lam(p, q)                   # Carmichael's totient function
        n = p * q                                   # Product of primes
        d = self.mod_inv(e, lambda_n)               # Find the multiplicative inverse of e
        crt = self.mod_inv(p, q)

        private_key = construct((n, e, d, p, q, crt))
        public_key  = construct((n, e))
        return (public_key, private_key)


    def encrypt(self, message):
        m = bytes_to_long(message.encode())
        ct = pow(m, self.public_key.e, self.public_key.n)
        return b64encode(long_to_bytes(ct)).decode()


    def decrypt(self, ct):
        pt = pow(ct, self.private_key.d, self.private_key.n)
        return long_to_bytes(pt).decode()


    def lam(self, p, q):
        return lcm(p - 1, q - 1)


    def mod_inv(self, e, mod):
        return pow(e, -1, mod)


class Cracker:
    def __init__(self, ct, public_key_file):
        self.ct = bytes_to_long(b64decode(ct))
        with open(public_key_file) as fp:
            self.public_key = import_key(fp.read())


    def crack(self):
        n = self.public_key.n
        print(f'[*] Factoring public key...')
        factors = self.factor(n)
        if not factors:
            print(f"[-] YAFU couldn't factor {n}")
            return None
        p, q = factors
        print(f'[+] Done:')
        print(f"{n} =\n {p}\n* {q}")

        lambda_n = lcm(p - 1, q - 1)
        d = pow(self.public_key.e, -1, lambda_n)
        pt = pow(self.ct, d, self.public_key.n)
        return long_to_bytes(pt).decode()


    def factor(self, n):
        yafu_cmd = ['yafu', f'factor({n})']
        output = subprocess.check_output(yafu_cmd).decode()
        matches = re.findall(r'^P\d\d = (\d+)', output, re.MULTILINE)
        if matches:
            return map(lambda n: int(n), matches)
        else:
            return None


if __name__ == '__main__':
    bits = 256
    print(f'[*] Generating a random {bits}-bit RSA Keypair')
    rsa = RSA(256)

    print(f'[*] Encrypting a message')
    ct = rsa.encrypt('This is a big secret')
    print(f'[*] Message: {ct}')

    print(f'[*] Load public key into cracker')
    cracker = Cracker(ct, 'public_key.pem')
    decrypted = cracker.crack()

    if decrypted:
        print(f'[+] Message was: {decrypted}')