padding.py 5.4 KB

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