More tweaks to random number generation

1 files changed, 50 insertions, 11 deletions

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
+