Added parallel.py module and ability to use multiprocessing when generating keys

3 files changed, 154 insertions, 22 deletions

diff --git a/rsa/key.py b/rsa/key.py
index 031c7e9..489500a 100644
--- a/rsa/key.py
+++ b/rsa/key.py
@@ -421,15 +421,20 @@ def extended_gcd(a, b):
     if (ly < 0): ly += oa              #If neg wrap modulo orignal a
     return (a, lx, ly)                 #Return only positive values
 
-def find_p_q(nbits, accurate=True):
+def find_p_q(nbits, getprime_func=rsa.prime.getprime, accurate=True):
     ''''Returns a tuple of two different primes of nbits bits each.
     
     The resulting p * q has exacty 2 * nbits bits, and the returned p and q
     will not be equal.
 
-    @param nbits: the number of bits in each of p and q.
-    @param accurate: whether to enable accurate mode or not.
-    @returns (p, q), where p > q
+    :param nbits: the number of bits in each of p and q.
+    :param getprime_func: the getprime function, defaults to
+        :py:func:`rsa.prime.getprime`.
+
+        *Introduced in Python-RSA 3.1*
+
+    :param accurate: whether to enable accurate mode or not.
+    :returns: (p, q), where p > q
 
     >>> (p, q) = find_p_q(128)
     >>> from rsa import common
@@ -457,9 +462,9 @@ def find_p_q(nbits, accurate=True):
     
     # Choose the two initial primes
     log.debug('find_p_q(%i): Finding p', nbits)
-    p = rsa.prime.getprime(pbits)
+    p = getprime_func(pbits)
     log.debug('find_p_q(%i): Finding q', nbits)
-    q = rsa.prime.getprime(qbits)
+    q = getprime_func(qbits)
 
     def is_acceptable(p, q):
         '''Returns True iff p and q are acceptable:
@@ -480,16 +485,12 @@ def find_p_q(nbits, accurate=True):
 
     # Keep choosing other primes until they match our requirements.
     change_p = False
-    tries = 0
     while not is_acceptable(p, q):
-        tries += 1
         # Change p on one iteration and q on the other
         if change_p:
-            log.debug('   find another p')
-            p = rsa.prime.getprime(pbits)
+            p = getprime_func(pbits)
         else:
-            log.debug('   find another q')
-            q = rsa.prime.getprime(qbits)
+            q = getprime_func(qbits)
 
         change_p = not change_p
 
@@ -522,41 +523,62 @@ def calculate_keys(p, q, nbits):
 
     return (e, d)
 
-def gen_keys(nbits, accurate=True):
+def gen_keys(nbits, getprime_func, accurate=True):
     """Generate RSA keys of nbits bits. Returns (p, q, e, d).
 
     Note: this can take a long time, depending on the key size.
     
-    @param nbits: the total number of bits in ``p`` and ``q``. Both ``p`` and
+    :param nbits: the total number of bits in ``p`` and ``q``. Both ``p`` and
         ``q`` will use ``nbits/2`` bits.
+    :param getprime_func: either :py:func:`rsa.prime.getprime` or a function
+        with similar signature.
     """
 
-    (p, q) = find_p_q(nbits // 2, accurate)
+    (p, q) = find_p_q(nbits // 2, getprime_func, accurate)
     (e, d) = calculate_keys(p, q, nbits // 2)
 
     return (p, q, e, d)
 
-def newkeys(nbits, accurate=True):
+def newkeys(nbits, accurate=True, poolsize=1):
     """Generates public and private keys, and returns them as (pub, priv).
 
     The public key is also known as the 'encryption key', and is a
-    :py:class:`PublicKey` object. The private key is also known as the
-    'decryption key' and is a :py:class:`PrivateKey` object.
+    :py:class:`rsa.PublicKey` object. The private key is also known as the
+    'decryption key' and is a :py:class:`rsa.PrivateKey` object.
 
     :param nbits: the number of bits required to store ``n = p*q``.
     :param accurate: when True, ``n`` will have exactly the number of bits you
         asked for. However, this makes key generation much slower. When False,
         `n`` may have slightly less bits.
+    :param poolsize: the number of processes to use to generate the prime
+        numbers. If set to a number > 1, a parallel algorithm will be used.
+        This requires Python 2.6 or newer.
 
     :returns: a tuple (:py:class:`rsa.PublicKey`, :py:class:`rsa.PrivateKey`)
+
+    The ``poolsize`` parameter was added in *Python-RSA 3.1* and requires
+    Python 2.6 or newer.
     
     """
 
     if nbits < 16:
         raise ValueError('Key too small')
 
-    (p, q, e, d) = gen_keys(nbits)
+    if poolsize < 1:
+        raise ValueError('Pool size (%i) should be >= 1' % poolsize)
+
+    # Determine which getprime function to use
+    if poolsize > 1:
+        from rsa import parallel
+        import functools
+
+        getprime_func = functools.partial(parallel.getprime, poolsize=poolsize)
+    else: getprime_func = rsa.prime.getprime
+
+    # Generate the key components
+    (p, q, e, d) = gen_keys(nbits, getprime_func)
     
+    # Create the key objects
     n = p * q
 
     return (
diff --git a/rsa/parallel.py b/rsa/parallel.py
new file mode 100644
index 0000000..82042c8
--- /dev/null
+++ b/rsa/parallel.py
@@ -0,0 +1,92 @@
+# -*- coding: utf-8 -*-
+#
+#  Copyright 2011 Sybren A. Stüvel <sybren@stuvel.eu>
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+'''Functions for parallel computation on multiple cores.
+
+Introduced in Python-RSA 3.1.
+
+.. note::
+
+    Requires Python 2.6 or newer.
+
+'''
+
+import multiprocessing as mp
+
+import rsa.prime
+import rsa.randnum
+
+def _find_prime(nbits, pipe):
+    while True:
+        integer = rsa.randnum.read_random_int(nbits)
+
+        # Make sure it's odd
+        integer |= 1
+
+        # Test for primeness
+        if rsa.prime.is_prime(integer):
+            pipe.send(integer)
+            return
+
+def getprime(nbits, poolsize):
+    """Returns a prime number that can be stored in 'nbits' bits.
+
+    Works in multiple threads at the same time.
+
+    >>> p = getprime(128, 3)
+    >>> rsa.prime.is_prime(p-1)
+    False
+    >>> rsa.prime.is_prime(p)
+    True
+    >>> rsa.prime.is_prime(p+1)
+    False
+    
+    >>> from rsa import common
+    >>> common.bit_size(p) == 128
+    True
+    
+    """
+
+    (pipe_recv, pipe_send) = mp.Pipe(duplex=False)
+
+    # Create processes
+    procs = [mp.Process(target=_find_prime, args=(nbits, pipe_send))
+             for _ in range(poolsize)]
+    [p.start() for p in procs]
+
+    result = pipe_recv.recv()
+
+    [p.terminate() for p in procs]
+
+    return result
+
+__all__ = ['getprime']
+
+    
+if __name__ == '__main__':
+    print 'Running doctests 1000x or until failure'
+    import doctest
+    
+    for count in range(100):
+        (failures, tests) = doctest.testmod()
+        if failures:
+            break
+        
+        if count and count % 10 == 0:
+            print '%i times' % count
+    
+    print 'Doctests done'
+
diff --git a/rsa/prime.py b/rsa/prime.py
index 4b2cb2e..340ab94 100644
--- a/rsa/prime.py
+++ b/rsa/prime.py
@@ -14,7 +14,11 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-'''Numerical functions related to primes.'''
+'''Numerical functions related to primes.
+
+Implementation based on the book Algorithm Design by Michael T. Goodrich and
+Roberto Tamassia, 2002.
+'''
 
 __all__ = [ 'getprime', 'are_relatively_prime']
 
@@ -36,8 +40,13 @@ def gcd(p, q):
 def jacobi(a, b):
     """Calculates the value of the Jacobi symbol (a/b) where both a and b are
     positive integers, and b is odd
+
+    :returns: -1, 0 or 1
     """
 
+    assert a > 0
+    assert b > 0
+
     if a == 0: return 0
     result = 1
     while a > 1:
@@ -58,7 +67,8 @@ def jacobi_witness(x, n):
     """
 
     j = jacobi(x, n) % n
-    f = pow(x, (n - 1) // 2, n)
+
+    f = pow(x, n >> 1, n)
 
     if j == f: return False
     return True
@@ -73,6 +83,14 @@ def randomized_primality_testing(n, k):
 
     # 50% of Jacobi-witnesses can report compositness of non-prime numbers
 
+    # The implemented algorithm using the Jacobi witness function has error
+    # probability q <= 0.5, according to Goodrich et. al
+    #
+    # q = 0.5
+    # t = int(math.ceil(k / log(1 / q, 2)))
+    # So t = k / log(2, 2) = k / 1 = k
+    # this means we can use range(k) rather than range(t)
+
     for _ in range(k):
         x = rsa.randnum.randint(n-1)
         if jacobi_witness(x, n): return False
@@ -102,7 +120,7 @@ def getprime(nbits):
     False
     
     >>> from rsa import common
-    >>> common.bit_size(p) <= 128
+    >>> common.bit_size(p) == 128
     True
     
     """