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

sybrenstuvel · sybrenstuvel · commit 360d04285b64 · 2011-08-10T12:52:59.000+02:00
diff --git 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
@@ -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
@@ -37,42 +37,55 @@ 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
 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
@@ -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 is_acceptable(p, q):
 
     # 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
@@ -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
