Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
 
 
 
 

235 řádky
6.2 KiB

  1. #
  2. # This file is part of pyasn1 software.
  3. #
  4. # Copyright (c) 2005-2019, Ilya Etingof <etingof@gmail.com>
  5. # License: https://pyasn1.readthedocs.io/en/latest/license.html
  6. #
  7. import io
  8. import os
  9. from pyasn1 import error
  10. from pyasn1.type import univ
  11. class CachingStreamWrapper(io.IOBase):
  12. """Wrapper around non-seekable streams.
  13. Note that the implementation is tied to the decoder,
  14. not checking for dangerous arguments for the sake
  15. of performance.
  16. The read bytes are kept in an internal cache until
  17. setting _markedPosition which may reset the cache.
  18. """
  19. def __init__(self, raw):
  20. self._raw = raw
  21. self._cache = io.BytesIO()
  22. self._markedPosition = 0
  23. def peek(self, n):
  24. result = self.read(n)
  25. self._cache.seek(-len(result), os.SEEK_CUR)
  26. return result
  27. def seekable(self):
  28. return True
  29. def seek(self, n=-1, whence=os.SEEK_SET):
  30. # Note that this not safe for seeking forward.
  31. return self._cache.seek(n, whence)
  32. def read(self, n=-1):
  33. read_from_cache = self._cache.read(n)
  34. if n != -1:
  35. n -= len(read_from_cache)
  36. if not n: # 0 bytes left to read
  37. return read_from_cache
  38. read_from_raw = self._raw.read(n)
  39. self._cache.write(read_from_raw)
  40. return read_from_cache + read_from_raw
  41. @property
  42. def markedPosition(self):
  43. """Position where the currently processed element starts.
  44. This is used for back-tracking in SingleItemDecoder.__call__
  45. and (indefLen)ValueDecoder and should not be used for other purposes.
  46. The client is not supposed to ever seek before this position.
  47. """
  48. return self._markedPosition
  49. @markedPosition.setter
  50. def markedPosition(self, value):
  51. # By setting the value, we ensure we won't seek back before it.
  52. # `value` should be the same as the current position
  53. # We don't check for this for performance reasons.
  54. self._markedPosition = value
  55. # Whenever we set _marked_position, we know for sure
  56. # that we will not return back, and thus it is
  57. # safe to drop all cached data.
  58. if self._cache.tell() > io.DEFAULT_BUFFER_SIZE:
  59. self._cache = io.BytesIO(self._cache.read())
  60. self._markedPosition = 0
  61. def tell(self):
  62. return self._cache.tell()
  63. def asSeekableStream(substrate):
  64. """Convert object to seekable byte-stream.
  65. Parameters
  66. ----------
  67. substrate: :py:class:`bytes` or :py:class:`io.IOBase` or :py:class:`univ.OctetString`
  68. Returns
  69. -------
  70. : :py:class:`io.IOBase`
  71. Raises
  72. ------
  73. : :py:class:`~pyasn1.error.PyAsn1Error`
  74. If the supplied substrate cannot be converted to a seekable stream.
  75. """
  76. if isinstance(substrate, io.BytesIO):
  77. return substrate
  78. elif isinstance(substrate, bytes):
  79. return io.BytesIO(substrate)
  80. elif isinstance(substrate, univ.OctetString):
  81. return io.BytesIO(substrate.asOctets())
  82. try:
  83. if substrate.seekable(): # Will fail for most invalid types
  84. return substrate
  85. else:
  86. return CachingStreamWrapper(substrate)
  87. except AttributeError:
  88. raise error.UnsupportedSubstrateError(
  89. "Cannot convert " + substrate.__class__.__name__ +
  90. " to a seekable bit stream.")
  91. def isEndOfStream(substrate):
  92. """Check whether we have reached the end of a stream.
  93. Although it is more effective to read and catch exceptions, this
  94. function
  95. Parameters
  96. ----------
  97. substrate: :py:class:`IOBase`
  98. Stream to check
  99. Returns
  100. -------
  101. : :py:class:`bool`
  102. """
  103. if isinstance(substrate, io.BytesIO):
  104. cp = substrate.tell()
  105. substrate.seek(0, os.SEEK_END)
  106. result = substrate.tell() == cp
  107. substrate.seek(cp, os.SEEK_SET)
  108. yield result
  109. else:
  110. received = substrate.read(1)
  111. if received is None:
  112. yield
  113. if received:
  114. substrate.seek(-1, os.SEEK_CUR)
  115. yield not received
  116. def peekIntoStream(substrate, size=-1):
  117. """Peek into stream.
  118. Parameters
  119. ----------
  120. substrate: :py:class:`IOBase`
  121. Stream to read from.
  122. size: :py:class:`int`
  123. How many bytes to peek (-1 = all available)
  124. Returns
  125. -------
  126. : :py:class:`bytes` or :py:class:`str`
  127. The return type depends on Python major version
  128. """
  129. if hasattr(substrate, "peek"):
  130. received = substrate.peek(size)
  131. if received is None:
  132. yield
  133. while len(received) < size:
  134. yield
  135. yield received
  136. else:
  137. current_position = substrate.tell()
  138. try:
  139. for chunk in readFromStream(substrate, size):
  140. yield chunk
  141. finally:
  142. substrate.seek(current_position)
  143. def readFromStream(substrate, size=-1, context=None):
  144. """Read from the stream.
  145. Parameters
  146. ----------
  147. substrate: :py:class:`IOBase`
  148. Stream to read from.
  149. Keyword parameters
  150. ------------------
  151. size: :py:class:`int`
  152. How many bytes to read (-1 = all available)
  153. context: :py:class:`dict`
  154. Opaque caller context will be attached to exception objects created
  155. by this function.
  156. Yields
  157. ------
  158. : :py:class:`bytes` or :py:class:`str` or :py:class:`SubstrateUnderrunError`
  159. Read data or :py:class:`~pyasn1.error.SubstrateUnderrunError`
  160. object if no `size` bytes is readily available in the stream. The
  161. data type depends on Python major version
  162. Raises
  163. ------
  164. : :py:class:`~pyasn1.error.EndOfStreamError`
  165. Input stream is exhausted
  166. """
  167. while True:
  168. # this will block unless stream is non-blocking
  169. received = substrate.read(size)
  170. if received is None: # non-blocking stream can do this
  171. yield error.SubstrateUnderrunError(context=context)
  172. elif not received and size != 0: # end-of-stream
  173. raise error.EndOfStreamError(context=context)
  174. elif len(received) < size:
  175. substrate.seek(-len(received), os.SEEK_CUR)
  176. # behave like a non-blocking stream
  177. yield error.SubstrateUnderrunError(context=context)
  178. else:
  179. break
  180. yield received