From 7790c1a15383b8b82ed2746fe0b54c554c430829 Mon Sep 17 00:00:00 2001 From: Yesudeep Mangalapilly Date: Sun, 14 Aug 2011 23:03:59 +0530 Subject: Much cleaner implementation of int2bytes. No loss in speed. --- rsa/_compat.py | 4 +- rsa/transform.py | 125 ++++++++++++++++++++++++++++++------------------------- 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 -- cgit v1.2.1