Prevent possible infinite loops.

1 files changed, 8 insertions, 3 deletions

diff --git a/rsa/prime.py b/rsa/prime.py
index c377c91..6f23f9d 100644
--- a/rsa/prime.py
+++ b/rsa/prime.py
@@ -53,10 +53,15 @@ def miller_rabin_primality_testing(n, k):
     :rtype: bool
     """
 
+    # prevent potential infinite loop when d = 0
+    if n < 2:
+        return False
+
     # Decompose (n - 1) to write it as (2 ** r) * d
     # While d is even, divide it by 2 and increase the exponent.
     d = n - 1
     r = 0
+
     while not (d & 1):
         r += 1
         d >>= 1
@@ -64,9 +69,7 @@ def miller_rabin_primality_testing(n, k):
     # Test k witnesses.
     for _ in range(k):
         # Generate random integer a, where 2 <= a <= (n - 2)
-        a = 0
-        while a < 2:
-            a = rsa.randnum.randint(n - 2)
+        a = rsa.randnum.randint(n - 4) + 2
 
         x = pow(a, d, n)
         if x == 1 or x == n - 1:
@@ -134,6 +137,8 @@ def getprime(nbits):
     True
     """
 
+    assert nbits > 3  # the loop wil hang on too small numbers
+
     while True:
         integer = rsa.randnum.read_random_odd_int(nbits)