More tweaks to random number generation

3 files changed, 64 insertions, 15 deletions

diff --git a/load_private.py b/load_private.py
index d6dd7bd..bda6e04 100755
--- a/load_private.py
+++ b/load_private.py
@@ -1,5 +1,9 @@
 #!/usr/bin/env python
 
+import logging
+
+logging.basicConfig(level=logging.DEBUG)
+
 from rsa import pkcs1, key
 
 # pick q, pick p
diff --git a/rsa/key.py b/rsa/key.py
index 233bb30..d4b03b8 100644
--- a/rsa/key.py
+++ b/rsa/key.py
@@ -9,10 +9,13 @@ of pyasn1.
 
 '''
 
+import logging
+
 import rsa.prime
 import rsa.pem
 import rsa.common
 
+log = logging.getLogger(__name__)
 
 class PublicKey(object):
     '''Represents a public RSA key.
@@ -283,7 +286,10 @@ def find_p_q(nbits):
     qbits = nbits - shift
     
     # Choose the two initial primes
+    log.debug('find_p_q(%i): Finding p', nbits)
     p = rsa.prime.getprime(pbits)
+
+    log.debug('find_p_q(%i): Finding q', nbits)
     q = rsa.prime.getprime(qbits)
 
     # Keep choosing other primes until they match our requirements.
@@ -304,12 +310,12 @@ def find_p_q(nbits):
             
         # Change p on one iteration and q on the other
         if tries & 1:
-            print 'Try %i, change q' % tries
-            qbits += diff
+            log.debug('Try %i, find another q', tries)
+            #qbits += diff
             q = rsa.prime.getprime(qbits)
         else:
-            print 'Try %i, change p' % tries
-            pbits += diff
+            log.debug('Try %i, find another p', tries)
+            #pbits += diff
             p = rsa.prime.getprime(pbits)
 
     # In the end we want p > q as described on
diff --git a/rsa/randnum.py b/rsa/randnum.py
index 7db6243..0d85f9f 100644
--- a/rsa/randnum.py
+++ b/rsa/randnum.py
@@ -1,29 +1,68 @@
 '''Functions for generating random numbers.'''
 
-import math
+# Source inspired by code by Yesudeep Mangalapilly <yesudeep@gmail.com>
+
 import os
 
 from rsa import common, transform
 
+def read_random_bits(nbits):
+    '''Reads 'nbits' random bits.
+
+    If nbits isn't a whole number of bytes, an extra byte will be appended with
+    only the lower bits set.
+    '''
+
+    nbytes, rbits = divmod(nbits, 8)
+
+    # Get the random bytes
+    randomdata = os.urandom(nbytes)
+
+    # Add the remaining random bits
+    if rbits > 0:
+        randomvalue = ord(os.urandom(1))
+        randomvalue >>= (8 - rbits)
+        randomdata = chr(randomvalue) + randomdata
+
+    return randomdata
+
+
 def read_random_int(nbits):
     """Reads a random integer of approximately nbits bits.
-    
-    The number of bits is rounded down to whole bytes to ensure that the
-    resulting number can be stored in ``nbits`` bits.
     """
 
-    randomdata = os.urandom(nbits // 8)
-    return transform.bytes2int(randomdata)
+    randomdata = read_random_bits(nbits)
+    value = transform.bytes2int(randomdata)
+
+    # Ensure that the number is large enough to just fill out the required
+    # number of bits.
+    value |= 1 << (nbits - 1)
+
+    return value
 
 def randint(maxvalue):
-    """Returns a random integer x with 1 <= x <= maxvalue"""
+    """Returns a random integer x with 1 <= x <= maxvalue
+    
+    May take a very long time in specific situations. If maxvalue needs N bits
+    to store, the closer maxvalue is to (2 ** N) - 1, the faster this function
+    is.
+    """
 
     bit_size = common.bit_size(maxvalue)
-    readbits = max(bit_size, 32)
-    mask = (1 << bit_size) - 1
 
+    tries = 0
     while True:
-        value = read_random_int(readbits) & mask
+        value = read_random_int(bit_size)
         if value <= maxvalue:
-            return value
+            break
+
+        if tries and tries % 10 == 0:
+            # After a lot of tries to get the right number of bits but still
+            # smaller than maxvalue, decrease the number of bits by 1. That'll
+            # dramatically increase the chances to get a large enough number.
+            bit_size -= 1
+        tries += 1
+
+    return value
+