Add main codes

Wiener.py

@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+
+from typing import Tuple
+
+import rsa
+from sympy import Symbol, solve
+
+from ContinuedFraction import ContinuedFraction
+from libs.RSAvulnerableKeyGenerator import generateKeys
+
+
+class Wiener:
+    def __init__(self,
+                 n: int or None = None,
+                 e: int or None = None,
+                 nbits: int = 1024):
+        """
+        Initializes a class for demonstrating Wiener's Attack (small private
+        exponent attack) on RSA.
+
+        You can either supply the class with your own public key (N,e) to try
+        it out in a real scenario, or leave them blank to let the program
+        generate a keypair vulnerable to the attack.
+
+        In the latter case, you can use the nbits parameter to state how many
+        bits (how strong) you want the key size to be. The larger the key
+        size, the more complicated the calculation (and thus longer time). If
+        you leave it blank, a default key size value of 1024 bits will be
+        used. However in reality, it is more common to use 2048 or 4096 bits.
+
+        :param n: the modulus N of your given RSA key
+        :param e: the public exponent e of your given RSA key
+        :param nbits: size of generated vulnerable RSA key. Default to 1024.
+        """
+        print("=" * 20, "Initialization", "=" * 20)
+
+        if n is None or e is None:
+            print("[!] No public key given. Will generate a new keypair.")
+            n, e = self._generateKeys(nbits)
+
+        # the local variable n is in lowercase to comply with naming conventions
+        self.N = n
+        self.e = e
+
+        # Will be used to store encrypted ciphertext
+        self.enc_message: bytes or None = None
+
+        print("=" * 20, "Initialization", "=" * 20)
+
+    def _generateKeys(self, nbits: int) -> Tuple[int, int]:
+        """
+        An internal method called to generate an RSA keypair vulnerable to
+        Wiener's Attack. The nbits parameter is passed down from the __init__
+        method.
+
+        :param nbits: size of generated vulnerable RSA key.
+        :return: the public key (N,e)
+        """
+
+        # Note that during keypair generation, the records of p and q are
+        # destroyed as they are only intermediate variables and the leak of
+        # which will render the keypair vulnerable.
+        e, n, d = generateKeys(nbits)
+
+        print("[+] Generated an RSA keypair vulnerable to Wiener's attack.")
+        print("N:\t", n)
+        print("e:\t", e)
+        print("d:\t", d)
+
+        # We will not return the private exponent d, which is highly
+        # secretive and unknown to attacker.
+        return n, e
+
+    def encrypt(self, message: str or None = None) -> bytes:
+        """
+        Encrypts a message with the public key stored in the class.
+
+        :param message: The message for encryption. Leave it blank to prompt
+                        for user input.
+        :return: The encrypted message in binary bytes
+        """
+        print("\n" + "=" * 20, "Encryption", "=" * 20)
+
+        print("Let's say Alex owns the private key and publishes the public "
+              "key so that others can send him messages only he can decrypt.")
+
+        print("Now, Ye uses the public key to send Alex an encrypted message.")
+
+        if message is None:
+            # No message entered, prompt for user input
+            message = input("[+] Message Content: ")
+        else:
+            print("[+] Message Content:", message)
+
+        # Here the message is first encoded from a string into byte using UTF-8
+        pubkey = rsa.PublicKey(self.N, self.e)
+        enc_message = rsa.encrypt(message.encode("utf8"), pubkey)
+
+        # The encrypted bytes is converted into hexadecimal digits to print out.
+        print("[+] Encrypted Message (hex):\t", enc_message.hex())
+
+        # Only the ciphertext is stored at the class level, since only Ye
+        # knows the original text message.
+        self.enc_message = enc_message
+
+        print("=" * 20, "Encryption", "=" * 20)
+
+        return enc_message
+
+    def crack(self) -> Tuple[int, int, int, int]:
+        """
+        Cracks the RSA cipher using Wiener's Attack. It uses properties of
+        continued fractions to guess the value of k/d using the convergents
+        of k/N. Detailed proof is given in our presentation.
+
+        :return: Cracked RSA private key defined by (p, q, d, φ(N))
+        """
+        print("\n" + "=" * 20, "CRACKING", "=" * 20)
+
+        print("As attacker, we intercepted the encrypted message.")
+        print("We also have knowledge of the public key (N,e).")
+
+        print("Let us now apply Wiener's Attack to the known public key.")
+        print("To do so, we are going to approximate d through the continued "
+              "fraction expansion of e/N")
+
+        input("Press Enter to start cracking...")
+
+        # the local variable n is in lowercase to comply with naming conventions
+        e, n = self.e, self.N
+
+        cf = ContinuedFraction(e, n)
+        expansions = cf.expansion()
+        print("\n[+] Found the continued fraction expansion of e/N")
+        print(expansions)
+
+        # See the slides for detailed proof
+        print("\nAs demonstrated in slides, we will use these coefficients to "
+              "recursively calculate the convergents of e/N, and use these "
+              "convergents to approximate k/d and guess their values.")
+        print("To verify that our guess is correct, we will go through the "
+              "following steps \n"
+              "1. Calculate φ(N) = (ed-1)/k\n"
+              "2. Solve the equation x^2 - (N-φ(N)+1)x + N = 0. "
+              "Ideally, p and q will be the two roots.\n"
+              "3. We will use the property N=pq to verify that.")
+
+        input("Press Enter to start iterating over convergents...")
+
+        convergents = cf.convergents_iter()
+
+        # a flag indicating whether the attack works
+        solved = False
+
+        # The convergent efficiently approximates e/N, which is then be used to
+        # approximate k/d. That is why we name them as k_guess and d_guess.
+        # See the slides for detailed proof.
+        for k_guess, d_guess in convergents:
+            print(f"[+] Trying k/d = {k_guess} / {d_guess}", end="\t")
+
+            if k_guess == 0:
+                # invalid
+                print("INCORRECT")
+                continue
+
+            # Recall that k * φ(N) = ed - 1 (because ed ≡ 1 (mod φ(N)))
+            # With (N,e) known and (k,d) approximated, we can deduce φ(N)
+            phi_guess = (e * d_guess - 1) // k_guess
+
+            # We use sympy to solve this equation.
+            # See the slides for why p and q are roots.
+            x = Symbol('x', integer=True)
+            roots = solve(x ** 2 + (phi_guess - n - 1) * x + n, x)
+
+            # There should be exactly two roots (p and q are distinct primes)
+            # if not, we proceed to the next attempted guess of k/d
+            if len(roots) != 2:
+                print("INCORRECT")
+                continue
+
+            # We verify if the guess works by multiplying the roots, which
+            # should give us N
+            p_guess, q_guess = roots
+            if p_guess * q_guess != n:
+                # This (p,q) pair is incorrect, proceed to next attempted guess
+                print("INCORRECT")
+                continue
+
+            print('\n\n[+] This guess worked! It gives us:')
+            print("p:\t", p_guess)
+            print("q:\t", q_guess)
+            print("N:\t", n)
+            print("e:\t", e)
+            print("d:\t", d_guess)
+            print('φ(N):', phi_guess)
+
+            # Cracked the private key! Breaking out of the iteration
+            solved = True
+            break
+
+        if not solved:
+            print("[-] Wiener's Attack failed")
+            raise Exception("Wiener's Attack failed")
+
+        print("=" * 20, "CRACKING", "=" * 20)
+        return p_guess, q_guess, d_guess, phi_guess
+
+    def decrypt(self, enc_message: bytes or None = None) -> str:
+        """
+        Decrypt the message by cracking the cipher.
+
+        :param enc_message: The message to be decrypted. Leave it blank to use
+                            the message we encrypted in self.encrypt()
+        :return: The decrypted message
+        """
+
+        if enc_message is None:
+            if self.enc_message is None:
+                raise Exception("An encrypted message is needed for decryption")
+            enc_message = self.enc_message
+
+        # Call the method to crack the cipher
+        p_guess, q_guess, d_guess, _ = self.crack()
+
+        print("\n" + "=" * 20, "Decrypt the private message", "=" * 20)
+        print("Since we have now cracked the private key (N,d), let's use it "
+              "to decrypt the private message sent to Alex!")
+
+        # Technically, only (N,d) is needed to recover the message, but the
+        # library I use requires all variables to instantiate the
+        # rsa.PrivateKey class.
+        privkey = rsa.PrivateKey(self.N, self.e, d_guess, p_guess, q_guess)
+        dec_message = rsa.decrypt(enc_message, privkey).decode("utf8")
+        print("[+] Decrypted Message:\t", dec_message)
+
+        print("=" * 20, "Decrypt the private message", "=" * 20)
+        return dec_message
+
+
+if __name__ == '__main__':
+    w = Wiener()
+
+    w.encrypt()
+    w.decrypt()
+
+    print("\nYay! 🎉")