transform.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215
  1. # -*- coding: utf-8 -*-
  2. #
  3. # Copyright 2011 Sybren A. Stüvel <sybren@stuvel.eu>
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License");
  6. # you may not use this file except in compliance with the License.
  7. # You may obtain a copy of the License at
  8. #
  9. # https://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. """Data transformation functions.
  17. From bytes to a number, number to bytes, etc.
  18. """
  19. from __future__ import absolute_import
  20. import binascii
  21. from struct import pack
  22. from rsa._compat import byte, is_integer
  23. from rsa import common, machine_size
  24. def bytes2int(raw_bytes):
  25. r"""Converts a list of bytes or an 8-bit string to an integer.
  26. When using unicode strings, encode it to some encoding like UTF8 first.
  27. >>> (((128 * 256) + 64) * 256) + 15
  28. 8405007
  29. >>> bytes2int(b'\x80@\x0f')
  30. 8405007
  31. """
  32. return int(binascii.hexlify(raw_bytes), 16)
  33. def _int2bytes(number, block_size=None):
  34. r"""Converts a number to a string of bytes.
  35. Usage::
  36. >>> _int2bytes(123456789)
  37. b'\x07[\xcd\x15'
  38. >>> bytes2int(_int2bytes(123456789))
  39. 123456789
  40. >>> _int2bytes(123456789, 6)
  41. b'\x00\x00\x07[\xcd\x15'
  42. >>> bytes2int(_int2bytes(123456789, 128))
  43. 123456789
  44. >>> _int2bytes(123456789, 3)
  45. Traceback (most recent call last):
  46. ...
  47. OverflowError: Needed 4 bytes for number, but block size is 3
  48. @param number: the number to convert
  49. @param block_size: the number of bytes to output. If the number encoded to
  50. bytes is less than this, the block will be zero-padded. When not given,
  51. the returned block is not padded.
  52. @throws OverflowError when block_size is given and the number takes up more
  53. bytes than fit into the block.
  54. """
  55. # Type checking
  56. if not is_integer(number):
  57. raise TypeError("You must pass an integer for 'number', not %s" %
  58. number.__class__)
  59. if number < 0:
  60. raise ValueError('Negative numbers cannot be used: %i' % number)
  61. # Do some bounds checking
  62. if number == 0:
  63. needed_bytes = 1
  64. raw_bytes = [b'\x00']
  65. else:
  66. needed_bytes = common.byte_size(number)
  67. raw_bytes = []
  68. # You cannot compare None > 0 in Python 3x. It will fail with a TypeError.
  69. if block_size and block_size > 0:
  70. if needed_bytes > block_size:
  71. raise OverflowError('Needed %i bytes for number, but block size '
  72. 'is %i' % (needed_bytes, block_size))
  73. # Convert the number to bytes.
  74. while number > 0:
  75. raw_bytes.insert(0, byte(number & 0xFF))
  76. number >>= 8
  77. # Pad with zeroes to fill the block
  78. if block_size and block_size > 0:
  79. padding = (block_size - needed_bytes) * b'\x00'
  80. else:
  81. padding = b''
  82. return padding + b''.join(raw_bytes)
  83. def bytes_leading(raw_bytes, needle=b'\x00'):
  84. """
  85. Finds the number of prefixed byte occurrences in the haystack.
  86. Useful when you want to deal with padding.
  87. :param raw_bytes:
  88. Raw bytes.
  89. :param needle:
  90. The byte to count. Default \x00.
  91. :returns:
  92. The number of leading needle bytes.
  93. """
  94. leading = 0
  95. # Indexing keeps compatibility between Python 2.x and Python 3.x
  96. _byte = needle[0]
  97. for x in raw_bytes:
  98. if x == _byte:
  99. leading += 1
  100. else:
  101. break
  102. return leading
  103. def int2bytes(number, fill_size=None, chunk_size=None, overflow=False):
  104. """
  105. Convert an unsigned integer to bytes (base-256 representation)::
  106. Does not preserve leading zeros if you don't specify a chunk size or
  107. fill size.
  108. .. NOTE:
  109. You must not specify both fill_size and chunk_size. Only one
  110. of them is allowed.
  111. :param number:
  112. Integer value
  113. :param fill_size:
  114. If the optional fill size is given the length of the resulting
  115. byte string is expected to be the fill size and will be padded
  116. with prefix zero bytes to satisfy that length.
  117. :param chunk_size:
  118. If optional chunk size is given and greater than zero, pad the front of
  119. the byte string with binary zeros so that the length is a multiple of
  120. ``chunk_size``.
  121. :param overflow:
  122. ``False`` (default). If this is ``True``, no ``OverflowError``
  123. will be raised when the fill_size is shorter than the length
  124. of the generated byte sequence. Instead the byte sequence will
  125. be returned as is.
  126. :returns:
  127. Raw bytes (base-256 representation).
  128. :raises:
  129. ``OverflowError`` when fill_size is given and the number takes up more
  130. bytes than fit into the block. This requires the ``overflow``
  131. argument to this function to be set to ``False`` otherwise, no
  132. error will be raised.
  133. """
  134. if number < 0:
  135. raise ValueError("Number must be an unsigned integer: %d" % number)
  136. if fill_size and chunk_size:
  137. raise ValueError("You can either fill or pad chunks, but not both")
  138. # Ensure these are integers.
  139. number & 1
  140. raw_bytes = b''
  141. # Pack the integer one machine word at a time into bytes.
  142. num = number
  143. word_bits, _, max_uint, pack_type = machine_size.get_word_alignment(num)
  144. pack_format = ">%s" % pack_type
  145. while num > 0:
  146. raw_bytes = pack(pack_format, num & max_uint) + raw_bytes
  147. num >>= word_bits
  148. # Obtain the index of the first non-zero byte.
  149. zero_leading = bytes_leading(raw_bytes)
  150. if number == 0:
  151. raw_bytes = b'\x00'
  152. # De-padding.
  153. raw_bytes = raw_bytes[zero_leading:]
  154. length = len(raw_bytes)
  155. if fill_size and fill_size > 0:
  156. if not overflow and length > fill_size:
  157. raise OverflowError(
  158. "Need %d bytes for number, but fill size is %d" %
  159. (length, fill_size)
  160. )
  161. raw_bytes = raw_bytes.rjust(fill_size, b'\x00')
  162. elif chunk_size and chunk_size > 0:
  163. remainder = length % chunk_size
  164. if remainder:
  165. padding_size = chunk_size - remainder
  166. raw_bytes = raw_bytes.rjust(length + padding_size, b'\x00')
  167. return raw_bytes
  168. if __name__ == '__main__':
  169. import doctest
  170. doctest.testmod()