Set Miller-Rabin rounds based on bitsize

rsa/prime.py

@@ -20,6 +20,7 @@
 Roberto Tamassia, 2002.
 """
 
+import rsa.common
 import rsa.randnum
 
 __all__ = ['getprime', 'are_relatively_prime']
@@ -116,7 +117,26 @@ def is_prime(number):
     #   * p, q bitsize: 1024; rounds: 4
     #   * p, q bitsize: 1536; rounds: 3
     # See: http://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.186-4.pdf
-    return miller_rabin_primality_testing(number, 7)
+
+    # (size, rounds) tuples, in descending order by size.
+    minimum_rounds_per_size = [
+        (1536, 3),
+        (1024, 4),
+        (512, 7),
+    ]
+    # Calculate number bitsize.
+    bitsize = rsa.common.bit_size(number)
+    # Set number of rounds.
+    for size, rounds in minimum_rounds_per_size:
+        if bitsize >= size:
+            k = rounds
+            break
+    # For smaller bitsizes, set arbitrary number of rounds.
+    else:
+        k = 10
+
+    # Run primality testing with (minimum + 1) rounds.
+    return miller_rabin_primality_testing(number, k + 1)
 
 
 def getprime(nbits):