Added support for padding blocks

3 files changed, 56 insertions, 15 deletions

diff --git a/rsa/core.py b/rsa/core.py
index f42801c..8f989df 100644..100755
--- a/rsa/core.py
+++ b/rsa/core.py
@@ -21,8 +21,8 @@ def encrypt_int(message, ekey, n):
         raise OverflowError("The message is too long")
 
     #Note: Bit exponents start at zero (bit counts start at 1) this is correct
-    safebit = rsa.common.bit_size(n) - 2                   #compute safe bit (MSB - 1)
-    message += (1 << safebit)                   #add safebit to ensure folding
+    safebit = rsa.common.bit_size(n) - 2        # compute safe bit (MSB - 1)
+    message += (1 << safebit)                   # add safebit to ensure folding
 
     return pow(message, ekey, n)
 
@@ -32,8 +32,8 @@ def decrypt_int(cyphertext, dkey, n):
 
     message = pow(cyphertext, dkey, n)
 
-    safebit = rsa.common.bit_size(n) - 2                   #compute safe bit (MSB - 1)
-    message -= (1 << safebit)                   #remove safebit before decode
+    safebit = rsa.common.bit_size(n) - 2        # compute safe bit (MSB - 1)
+    message -= (1 << safebit)                   # remove safebit before decode
 
     return message
 
diff --git a/rsa/keygen.py b/rsa/keygen.py
index 42601b8..66f2658 100644..100755
--- a/rsa/keygen.py
+++ b/rsa/keygen.py
@@ -95,7 +95,8 @@ def newkeys(nbits):
     key consists of a dict {d: ...., p: ...., q: ....).
     """
 
-    nbits = max(9, nbits)           # Don't let nbits go below 9 bits
+    # Don't let nbits go below 9 bits
+    nbits = max(9, nbits)
     (p, q, e, d) = gen_keys(nbits)
 
     return ( {'e': e, 'n': p*q}, {'d': d, 'p': p, 'q': q} )
diff --git a/rsa/transform.py b/rsa/transform.py
index 6719dcb..35b7461 100644..100755
--- a/rsa/transform.py
+++ b/rsa/transform.py
@@ -9,7 +9,7 @@ import types
 def bit_size(number):
     """Returns the number of bits required to hold a specific long number"""
 
-    return int(math.ceil(math.log(number,2)))
+    return int(math.ceil(math.log(number, 2)))
 
 def bytes2int(bytes):
     """Converts a list of bytes or an 8-bit string to an integer.
@@ -36,26 +36,60 @@ def bytes2int(bytes):
 
     return integer
 
-def int2bytes(number):
+def int2bytes(number, block_size=None):
     r'''Converts a number to a string of bytes.
 
+    @param number: the number to convert
+    @param block_size: the number of bytes to output. If the number encoded to
+        bytes is less than this, the block will be zero-padded. When not given,
+        the returned block is not padded.
+
+    @throws OverflowError when block_size is given and the number takes up more
+        bytes than fit into the block.
+
+
     >>> int2bytes(123456789)
     '\x07[\xcd\x15'
     >>> bytes2int(int2bytes(123456789))
     123456789
 
-    '''
+    >>> int2bytes(123456789, 6)
+    '\x00\x00\x07[\xcd\x15'
+    >>> bytes2int(int2bytes(123456789, 128))
+    123456789
 
-    if not (type(number) is types.LongType or type(number) is types.IntType):
-        raise TypeError("You must pass a long or an int")
+    >>> int2bytes(123456789, 3)
+    Traceback (most recent call last):
+    ...
+    OverflowError: Needed 4 bytes for number, but block size is 3
 
-    string = ""
+    '''
 
-    while number > 0:
-        string = "%s%s" % (chr(number & 0xFF), string)
-        number /= 256
+    # Type checking
+    if type(number) not in (types.LongType, types.IntType):
+        raise TypeError("You must pass an integer for 'number', not %s" %
+            number.__class__)
+
+    # Do some bounds checking
+    if block_size is not None:
+        needed_bytes = int(math.ceil(bit_size(number) / 8.0))
+        if needed_bytes > block_size:
+            raise OverflowError('Needed %i bytes for number, but block size '
+                'is %i' % (needed_bytes, block_size))
     
-    return string
+    # Convert the number to bytes.
+    bytes = []
+    while number > 0:
+        bytes.insert(0, chr(number & 0xFF))
+        number >>= 8
+
+    # Pad with zeroes to fill the block
+    if block_size is not None:
+        padding = (block_size - needed_bytes) * '\x00'
+    else:
+        padding = ''
+
+    return padding + ''.join(bytes)
 
 
 def to64(number):
@@ -158,3 +192,9 @@ def str642int(string):
         integer += from64(byte)
 
     return integer
+
+
+if __name__ == '__main__':
+    import doctest
+    doctest.testmod()
+