Fix issue #98: Generate p and q of exactly half the number of bits of Nissue-98-fips-186-4-prime-selection

This commit implements the FIPS 186-4 recommendation for lower & upper
bounds on p & q, and on a minimum distance |p-q|.

@joostrijneveld and @RichardThiessen I'd appreciate your eyes on this
before it's pushed to master. If you can, please test & give me some
feedback on #98.

2 files changed, 41 insertions, 21 deletions

diff --git a/rsa/key.py b/rsa/key.py
index b4d902b..2667e28 100644
--- a/rsa/key.py
+++ b/rsa/key.py
@@ -33,6 +33,7 @@ of pyasn1.
 
 """
 
+import decimal
 import logging
 import typing
 import warnings
@@ -47,6 +48,17 @@ import rsa.core
 log = logging.getLogger(__name__)
 DEFAULT_EXPONENT = 65537
 
+SQRT_2 = 2 ** 0.5
+
+# See the find_p_q() function below.
+# We have to resolve to using the Decimal module for really large
+# numbers, as otherwise Python refuses to convert to float for the
+# sqrt(2) multiplication. This causes a loss in precision of the
+# lower bound. To ensure we don't generate primes below the lower
+# bound described by FIPS.186-4, we nudge the value of sqrt(2)
+# slightly higher.
+SQRT_2_DEC = decimal.Decimal(SQRT_2 * 1.000000000000011)
+
 
 class AbstractKey:
     """Abstract superclass for private and public keys."""
@@ -602,17 +614,19 @@ def find_p_q(nbits: int, getprime_func=rsa.prime.getprime, accurate=True) -> typ
 
     total_bits = nbits * 2
 
-    # Make sure that p and q aren't too close or the factoring programs can
-    # factor n.
-    shift = nbits // 16
-    pbits = nbits + shift
-    qbits = nbits - shift
+    # From https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.186-4.pdf
+    # section "B.3.1  Criteria for IFC Key Pairs", page 53.
+    sqrt_2 = SQRT_2 if nbits < 1024 else SQRT_2_DEC
+    lower_bound = int(sqrt_2 * 2 ** (nbits - 1))
+
+    upper_bound = 2 ** nbits - 1
+    distance_threshold = max(1, 2 ** (nbits - 100))
 
     # Choose the two initial primes
     log.debug('find_p_q(%i): Finding p', nbits)
-    p = getprime_func(pbits)
+    p = getprime_func(nbits)
     log.debug('find_p_q(%i): Finding q', nbits)
-    q = getprime_func(qbits)
+    q = getprime_func(nbits)
 
     def is_acceptable(p, q):
         """Returns True iff p and q are acceptable:
@@ -621,7 +635,11 @@ def find_p_q(nbits: int, getprime_func=rsa.prime.getprime, accurate=True) -> typ
             - (p * q) has the right nr of bits (when accurate=True)
         """
 
-        if p == q:
+        if abs(p - q) < distance_threshold:
+            return False
+        if not lower_bound <= p <= upper_bound:
+            return False
+        if not lower_bound <= q <= upper_bound:
             return False
 
         if not accurate:
@@ -636,9 +654,9 @@ def find_p_q(nbits: int, getprime_func=rsa.prime.getprime, accurate=True) -> typ
     while not is_acceptable(p, q):
         # Change p on one iteration and q on the other
         if change_p:
-            p = getprime_func(pbits)
+            p = getprime_func(nbits)
         else:
-            q = getprime_func(qbits)
+            q = getprime_func(nbits)
 
         change_p = not change_p
 
diff --git a/tests/test_key.py b/tests/test_key.py
index 9db30ce..0ae105d 100644
--- a/tests/test_key.py
+++ b/tests/test_key.py
@@ -50,23 +50,25 @@ class KeyGenTest(unittest.TestCase):
 
     def test_custom_getprime_func(self):
         # List of primes to test with, in order [p, q, p, q, ....]
-        # By starting with two of the same primes, we test that this is
-        # properly rejected.
-        primes = [64123, 64123, 64123, 50957, 39317, 33107]
+        # By starting with two of the same primes, we test for minimum distance.
+        # The 3rd prime is too low, and falls below the lower threshold.
+        # The 4th and 5th prime form a correct pair.
+        primes = [64123, 64123, 33751, 50957, 56821, 51109]
 
         def getprime(_):
             return primes.pop(0)
 
-        # This exponent will cause two other primes to be generated.
-        exponent = 136407
+        (p, q, e, d) = rsa.key.gen_keys(32,
+                                        accurate=True,
+                                        getprime_func=getprime)
+        self.assertEqual(56821, p)
+        self.assertEqual(50957, q)
 
-        (p, q, e, d) = rsa.key.gen_keys(64,
-                                        accurate=False,
-                                        getprime_func=getprime,
-                                        exponent=exponent)
-        self.assertEqual(39317, p)
-        self.assertEqual(33107, q)
+    def test_gen_usable_sizes_1024(self):
+        priv, pub = rsa.key.newkeys(1024)
 
+    def test_gen_usable_sizes_2048(self):
+        priv, pub = rsa.key.newkeys(2048)
 
 class HashTest(unittest.TestCase):
     """Test hashing of keys"""