Much cleaner implementation of int2bytes. No loss in speed.
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 @@
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 @@
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 @@
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)::
- >>> int2bytes(123456789, 3)
- Traceback (most recent call last):
- ...
- OverflowError: Need 4 bytes for number, but chunk size is 3
+ Does not preserve leading zeros if you don't specify a chunk size or
+ fill size.
+
+ .. 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
+ # 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:]
- # 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:
+ 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