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

6 files changed, 229 insertions, 49 deletions

diff --git a/CHANGELOG.txt b/CHANGELOG.txt
index 64cdfa9..5918511 100644
--- a/CHANGELOG.txt
+++ b/CHANGELOG.txt
@@ -1,6 +1,12 @@
 Python-RSA changelog
 ========================================
 
+Version 3.1 - in development
+----------------------------------------
+
+- Added ability to generate keys on multiple cores simultaneously.
+
+
 Version 3.0.1 - released 2011-08-07
 ----------------------------------------
 
diff --git a/create_timing_table.py b/create_timing_table.py
new file mode 100755
index 0000000..b1b2871
--- /dev/null
+++ b/create_timing_table.py
@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+
+import time
+import rsa
+
+poolsize = 8
+accurate = True
+
+def run_speed_test(bitsize):
+
+    iterations = 0
+    start = end = time.time()
+
+    # At least a number of iterations, and at least 2 seconds
+    while iterations < 10 or end - start < 2:
+        iterations += 1
+        rsa.newkeys(bitsize, accurate=accurate, poolsize=poolsize)
+        end = time.time()
+
+    duration = end - start
+    dur_per_call = duration / iterations
+
+    print '%5i bit: %9.3f sec. (%i iterations over %.1f seconds)' % (bitsize,
+            dur_per_call, iterations, duration)
+
+for bitsize in (128, 256, 384, 512, 1024, 2048, 3072, 4096):
+    run_speed_test(bitsize)
+
+
diff --git a/doc/usage.rst b/doc/usage.rst
index e4436e4..611e868 100644
--- a/doc/usage.rst
+++ b/doc/usage.rst
@@ -37,6 +37,10 @@ Alternatively you can use :py:meth:`rsa.PrivateKey.load_pkcs1` and
     ...     keydata = privatefile.read()
     >>> pubkey = rsa.PrivateKey.load_pkcs1(keydata)
 
+
+Time to generate a key
+++++++++++++++++++++++++++++++++++++++++
+
 Generating a keypair may take a long time, depending on the number of
 bits required. The number of bits determines the cryptographic
 strength of the key, as well as the size of the message you can
@@ -44,35 +48,44 @@ encrypt. If you don't mind having a slightly smaller key than you
 requested, you can pass ``accurate=False`` to speed up the key
 generation process.
 
-These are some average timings from my netbook (Linux 2.6, 1.6 GHz
-Intel Atom N270 CPU, 2 GB RAM). Since key generation is a random
-process, times may differ.
-
-+----------------+------------------+
-| Keysize (bits) | Time to generate |
-+================+==================+
-| 32             | 0.01 sec.        |
-+----------------+------------------+
-| 64             | 0.03 sec.        |
-+----------------+------------------+
-| 96             | 0.04 sec.        |
-+----------------+------------------+
-| 128            | 0.08 sec.        |
-+----------------+------------------+
-| 256            | 0.27 sec.        |
-+----------------+------------------+
-| 384            | 0.93 sec.        |
-+----------------+------------------+
-| 512            | 1.21 sec.        |
-+----------------+------------------+
-| 1024           | 7.93 sec.        |
-+----------------+------------------+
-| 2048           | 132.97 sec.      |
-+----------------+------------------+
+Another way to speed up the key generation process is to use multiple
+processes in parallel to speed up the key generation. Use no more than
+the number of processes that your machine can run in parallel; a
+dual-core machine should use ``poolsize=2``; a quad-core
+hyperthreading machine can run two threads on each core, and thus can
+use ``poolsize=8``.
+
+    >>> (pubkey, privkey) = rsa.newkeys(512, poolsize=8)
+
+These are some average timings from my desktop machine (Linux 2.6,
+2.93 GHz quad-core Intel Core i7, 16 GB RAM) using 64-bit CPython 2.7.
+Since key generation is a random process, times may differ even on
+similar hardware. On all tests, we used the default ``accurate=True``.
+
++----------------+------------------+------------------+
+| Keysize (bits) | single process   | eight processes  |
++================+==================+==================+
+| 128            | 0.01 sec.        | 0.01 sec.        |
++----------------+------------------+------------------+
+| 256            | 0.03 sec.        | 0.02 sec.        |
++----------------+------------------+------------------+
+| 384            | 0.09 sec.        | 0.04 sec.        |
++----------------+------------------+------------------+
+| 512            | 0.11 sec.        | 0.07 sec.        |
++----------------+------------------+------------------+
+| 1024           | 0.79 sec.        | 0.30 sec.        |
++----------------+------------------+------------------+
+| 2048           | 6.55 sec.        | 1.60 sec.        |
++----------------+------------------+------------------+
+| 3072           | 23.4 sec.        | 7.14 sec.        |
++----------------+------------------+------------------+
+| 4096           | 72.0 sec.        | 24.4 sec.        |
++----------------+------------------+------------------+
 
 If key generation is too slow for you, you could use OpenSSL to
-generate them for you, then load them in your Python code. See
-:ref:`openssl` for more information.
+generate them for you, then load them in your Python code. OpenSSL
+generates a 4096-bit key in 3.5 seconds on the same machine as used
+above. See :ref:`openssl` for more information.
 
 Key size requirements
 --------------------------------------------------
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
     
     """