Much cleaner implementation of int2bytes. No loss in speed.

2 files changed, 71 insertions, 58 deletions

diff --git a/rsa/_compat.py b/rsa/_compat.py
index b01fd3e..39a967b 100644
--- a/rsa/_compat.py
+++ b/rsa/_compat.py
@@ -121,8 +121,8 @@ def byte(num):
     return pack("B", num)
 
 
-def get_machine_alignment(num, force_arch=64,
-                          _machine_word_size=MACHINE_WORD_SIZE):
+def get_word_alignment(num, force_arch=64,
+                       _machine_word_size=MACHINE_WORD_SIZE):
     """
     Returns alignment details for the given number based on the platform
     Python is running on.
diff --git a/rsa/transform.py b/rsa/transform.py
index 40f7ca6..7dba2e0 100644
--- a/rsa/transform.py
+++ b/rsa/transform.py
@@ -33,7 +33,7 @@ except ImportError:
 import binascii
 from struct import pack
 from rsa import common
-from rsa._compat import is_integer, b, byte, get_machine_alignment
+from rsa._compat import is_integer, b, byte, get_word_alignment
 
 
 ZERO_BYTE = b('\x00')
@@ -111,89 +111,102 @@ def _int2bytes(number, block_size=0):
     return padding + b('').join(raw_bytes)
 
 
-
-def int2bytes(number, chunk_size=0,
-                     _zero_byte=ZERO_BYTE,
-                     _get_machine_alignment=get_machine_alignment):
+def bytes_leading(raw_bytes, needle=ZERO_BYTE):
     """
-    Convert a integer to bytes (base-256 representation)::
+    Finds the number of prefixed byte occurrences in the haystack.
 
-        int2bytes(n:int, chunk_size:int) : string
+    Useful when you want to deal with padding.
 
-    .. WARNING:
-        Does not preserve leading zeros if you don't specify a chunk size.
+    :param raw_bytes:
+        Raw bytes.
+    :param needle:
+        The byte to count. Default \000.
+    :returns:
+        The number of leading needle bytes.
+    """
+    leading = 0
+    # Indexing keeps compatibility between Python 2.x and Python 3.x
+    _byte = needle[0]
+    for x in raw_bytes:
+        if x == _byte:
+            leading += 1
+        else:
+            break
+    return leading
 
-    Usage::
-    
-        >>> int2bytes(123456789)
-        b'\x07[\xcd\x15'
-        >>> bytes2int(int2bytes(123456789))
-        123456789
 
-        >>> int2bytes(123456789, 6)
-        b'\x00\x00\x07[\xcd\x15'
-        >>> bytes2int(int2bytes(123456789, 128))
-        123456789
+def int2bytes(number, fill_size=0, chunk_size=0, overflow=False):
+    """
+    Convert an unsigned integer to bytes (base-256 representation)::
+
+    Does not preserve leading zeros if you don't specify a chunk size or
+    fill size.
 
-        >>> int2bytes(123456789, 3)
-        Traceback (most recent call last):
-        ...
-        OverflowError: Need 4 bytes for number, but chunk size is 3
+    .. NOTE:
+        You must not specify both fill_size and chunk_size. Only one
+        of them is allowed.
 
     :param number:
         Integer value
+    :param fill_size:
+        If the optional fill size is given the length of the resulting
+        byte string is expected to be the fill size and will be padded
+        with prefix zero bytes to satisfy that length.
     :param chunk_size:
         If optional chunk size is given and greater than zero, pad the front of
         the byte string with binary zeros so that the length is a multiple of
-        ``chunk_size``. Raises an OverflowError if the chunk_size is not
-        sufficient to represent the integer.
+        ``chunk_size``.
+    :param overflow:
+        ``False`` (default). If this is ``True``, no ``OverflowError``
+        will be raised when the fill_size is shorter than the length
+        of the generated byte sequence. Instead the byte sequence will
+        be returned as is.
     :returns:
         Raw bytes (base-256 representation).
     :raises:
-        ``OverflowError`` when block_size is given and the number takes up more
-        bytes than fit into the block.
+        ``OverflowError`` when fill_size is given and the number takes up more
+        bytes than fit into the block. This requires the ``overflow``
+        argument to this function to be set to ``False`` otherwise, no
+        error will be raised.
     """
     if number < 0:
-        raise ValueError('Number must be unsigned integer: %d' % number)
+        raise ValueError("Number must be an unsigned integer: %d" % number)
+
+    if fill_size and chunk_size:
+        raise ValueError("You can either fill or pad chunks, but not both")
+
+    # Ensure these are integers.
+    number & 1 and chunk_size & 1 and fill_size & 1
 
     raw_bytes = b('')
-    if not number:
-        # 0 == '\x00'
-        raw_bytes = _zero_byte
 
-    # Align packing to machine word size.
+    # Pack the integer one machine word at a time into bytes.
     num = number
-    word_bits, word_bytes, max_uint, pack_type = _get_machine_alignment(num)
-    pack_format = ">" + pack_type
+    word_bits, _, max_uint, pack_type = get_word_alignment(num)
+    pack_format = ">%s" % pack_type
     while num > 0:
         raw_bytes = pack(pack_format, num & max_uint) + raw_bytes
         num >>= word_bits
-
-    # Count the number of zero prefix bytes.
-    zero_leading = 0
-    for zero_leading, x in enumerate(raw_bytes):
-        if x != _zero_byte[0]:
-            break
-
-    if chunk_size > 0:
-        # Bounds checking. We're not doing this up-front because the
-        # most common use case is not specifying a chunk size. In the worst
-        # case, the number will already have been converted to bytes above.
-        length = len(raw_bytes) - zero_leading
-        if length > chunk_size:
+    # Obtain the index of the first non-zero byte.
+    zero_leading = bytes_leading(raw_bytes)
+    if number == 0:
+        raw_bytes = ZERO_BYTE
+    # De-padding.
+    raw_bytes = raw_bytes[zero_leading:]
+
+    length = len(raw_bytes)
+    if fill_size > 0:
+        if not overflow and length > fill_size:
             raise OverflowError(
-                "Need %d bytes for number, but chunk size is %d" %
-                (length, chunk_size)
+                "Need %d bytes for number, but fill size is %d" %
+                (length, fill_size)
             )
+        raw_bytes = raw_bytes.rjust(fill_size, ZERO_BYTE)
+    elif chunk_size > 0:
         remainder = length % chunk_size
         if remainder:
-            padding_size = (chunk_size - remainder)
-            if zero_leading > 0:
-                raw_bytes = raw_bytes[zero_leading-padding_size:]
-            else:
-                raw_bytes = (padding_size * _zero_byte) + raw_bytes
-    else:
-        raw_bytes = raw_bytes[zero_leading:]
+            padding_size = chunk_size - remainder
+            raw_bytes = raw_bytes.rjust(length + padding_size, ZERO_BYTE)
     return raw_bytes