padding.py 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221
  1. # This file is dual licensed under the terms of the Apache License, Version
  2. # 2.0, and the BSD License. See the LICENSE file in the root of this repository
  3. # for complete details.
  4. import abc
  5. import typing
  6. from cryptography import utils
  7. from cryptography.exceptions import AlreadyFinalized
  8. from cryptography.hazmat.bindings._padding import lib
  9. class PaddingContext(metaclass=abc.ABCMeta):
  10. @abc.abstractmethod
  11. def update(self, data: bytes) -> bytes:
  12. """
  13. Pads the provided bytes and returns any available data as bytes.
  14. """
  15. @abc.abstractmethod
  16. def finalize(self) -> bytes:
  17. """
  18. Finalize the padding, returns bytes.
  19. """
  20. def _byte_padding_check(block_size: int) -> None:
  21. if not (0 <= block_size <= 2040):
  22. raise ValueError("block_size must be in range(0, 2041).")
  23. if block_size % 8 != 0:
  24. raise ValueError("block_size must be a multiple of 8.")
  25. def _byte_padding_update(
  26. buffer_: typing.Optional[bytes], data: bytes, block_size: int
  27. ) -> typing.Tuple[bytes, bytes]:
  28. if buffer_ is None:
  29. raise AlreadyFinalized("Context was already finalized.")
  30. utils._check_byteslike("data", data)
  31. buffer_ += bytes(data)
  32. finished_blocks = len(buffer_) // (block_size // 8)
  33. result = buffer_[: finished_blocks * (block_size // 8)]
  34. buffer_ = buffer_[finished_blocks * (block_size // 8) :]
  35. return buffer_, result
  36. def _byte_padding_pad(
  37. buffer_: typing.Optional[bytes],
  38. block_size: int,
  39. paddingfn: typing.Callable[[int], bytes],
  40. ) -> bytes:
  41. if buffer_ is None:
  42. raise AlreadyFinalized("Context was already finalized.")
  43. pad_size = block_size // 8 - len(buffer_)
  44. return buffer_ + paddingfn(pad_size)
  45. def _byte_unpadding_update(
  46. buffer_: typing.Optional[bytes], data: bytes, block_size: int
  47. ) -> typing.Tuple[bytes, bytes]:
  48. if buffer_ is None:
  49. raise AlreadyFinalized("Context was already finalized.")
  50. utils._check_byteslike("data", data)
  51. buffer_ += bytes(data)
  52. finished_blocks = max(len(buffer_) // (block_size // 8) - 1, 0)
  53. result = buffer_[: finished_blocks * (block_size // 8)]
  54. buffer_ = buffer_[finished_blocks * (block_size // 8) :]
  55. return buffer_, result
  56. def _byte_unpadding_check(
  57. buffer_: typing.Optional[bytes],
  58. block_size: int,
  59. checkfn: typing.Callable[[bytes, int], int],
  60. ) -> bytes:
  61. if buffer_ is None:
  62. raise AlreadyFinalized("Context was already finalized.")
  63. if len(buffer_) != block_size // 8:
  64. raise ValueError("Invalid padding bytes.")
  65. valid = checkfn(buffer_, block_size // 8)
  66. if not valid:
  67. raise ValueError("Invalid padding bytes.")
  68. pad_size = buffer_[-1]
  69. return buffer_[:-pad_size]
  70. class PKCS7(object):
  71. def __init__(self, block_size: int):
  72. _byte_padding_check(block_size)
  73. self.block_size = block_size
  74. def padder(self):
  75. return _PKCS7PaddingContext(self.block_size)
  76. def unpadder(self):
  77. return _PKCS7UnpaddingContext(self.block_size)
  78. class _PKCS7PaddingContext(PaddingContext):
  79. _buffer: typing.Optional[bytes]
  80. def __init__(self, block_size: int):
  81. self.block_size = block_size
  82. # TODO: more copies than necessary, we should use zero-buffer (#193)
  83. self._buffer = b""
  84. def update(self, data: bytes) -> bytes:
  85. self._buffer, result = _byte_padding_update(
  86. self._buffer, data, self.block_size
  87. )
  88. return result
  89. def _padding(self, size: int) -> bytes:
  90. return bytes([size]) * size
  91. def finalize(self) -> bytes:
  92. result = _byte_padding_pad(
  93. self._buffer, self.block_size, self._padding
  94. )
  95. self._buffer = None
  96. return result
  97. class _PKCS7UnpaddingContext(PaddingContext):
  98. _buffer: typing.Optional[bytes]
  99. def __init__(self, block_size: int):
  100. self.block_size = block_size
  101. # TODO: more copies than necessary, we should use zero-buffer (#193)
  102. self._buffer = b""
  103. def update(self, data: bytes) -> bytes:
  104. self._buffer, result = _byte_unpadding_update(
  105. self._buffer, data, self.block_size
  106. )
  107. return result
  108. def finalize(self) -> bytes:
  109. result = _byte_unpadding_check(
  110. self._buffer, self.block_size, lib.Cryptography_check_pkcs7_padding
  111. )
  112. self._buffer = None
  113. return result
  114. class ANSIX923(object):
  115. def __init__(self, block_size: int):
  116. _byte_padding_check(block_size)
  117. self.block_size = block_size
  118. def padder(self) -> PaddingContext:
  119. return _ANSIX923PaddingContext(self.block_size)
  120. def unpadder(self) -> PaddingContext:
  121. return _ANSIX923UnpaddingContext(self.block_size)
  122. class _ANSIX923PaddingContext(PaddingContext):
  123. _buffer: typing.Optional[bytes]
  124. def __init__(self, block_size: int):
  125. self.block_size = block_size
  126. # TODO: more copies than necessary, we should use zero-buffer (#193)
  127. self._buffer = b""
  128. def update(self, data: bytes) -> bytes:
  129. self._buffer, result = _byte_padding_update(
  130. self._buffer, data, self.block_size
  131. )
  132. return result
  133. def _padding(self, size: int) -> bytes:
  134. return bytes([0]) * (size - 1) + bytes([size])
  135. def finalize(self) -> bytes:
  136. result = _byte_padding_pad(
  137. self._buffer, self.block_size, self._padding
  138. )
  139. self._buffer = None
  140. return result
  141. class _ANSIX923UnpaddingContext(PaddingContext):
  142. _buffer: typing.Optional[bytes]
  143. def __init__(self, block_size: int):
  144. self.block_size = block_size
  145. # TODO: more copies than necessary, we should use zero-buffer (#193)
  146. self._buffer = b""
  147. def update(self, data: bytes) -> bytes:
  148. self._buffer, result = _byte_unpadding_update(
  149. self._buffer, data, self.block_size
  150. )
  151. return result
  152. def finalize(self) -> bytes:
  153. result = _byte_unpadding_check(
  154. self._buffer,
  155. self.block_size,
  156. lib.Cryptography_check_ansix923_padding,
  157. )
  158. self._buffer = None
  159. return result