data.py 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323
  1. """AMQP Table Encoding/Decoding"""
  2. import struct
  3. import decimal
  4. import calendar
  5. import warnings
  6. from datetime import datetime
  7. from pika import exceptions
  8. from pika.compat import PY2, basestring
  9. from pika.compat import unicode_type, long, as_bytes
  10. def encode_short_string(pieces, value):
  11. """Encode a string value as short string and append it to pieces list
  12. returning the size of the encoded value.
  13. :param list pieces: Already encoded values
  14. :param str value: String value to encode
  15. :rtype: int
  16. """
  17. encoded_value = as_bytes(value)
  18. length = len(encoded_value)
  19. # 4.2.5.3
  20. # Short strings, stored as an 8-bit unsigned integer length followed by zero
  21. # or more octets of data. Short strings can carry up to 255 octets of UTF-8
  22. # data, but may not contain binary zero octets.
  23. # ...
  24. # 4.2.5.5
  25. # The server SHOULD validate field names and upon receiving an invalid field
  26. # name, it SHOULD signal a connection exception with reply code 503 (syntax
  27. # error).
  28. # -> validate length (avoid truncated utf-8 / corrupted data), but skip null
  29. # byte check.
  30. if length > 255:
  31. raise exceptions.ShortStringTooLong(encoded_value)
  32. pieces.append(struct.pack('B', length))
  33. pieces.append(encoded_value)
  34. return 1 + length
  35. if PY2:
  36. def decode_short_string(encoded, offset):
  37. """Decode a short string value from ``encoded`` data at ``offset``.
  38. """
  39. length = struct.unpack_from('B', encoded, offset)[0]
  40. offset += 1
  41. # Purely for compatibility with original python2 code. No idea what
  42. # and why this does.
  43. value = encoded[offset:offset + length]
  44. try:
  45. value = bytes(value)
  46. except UnicodeEncodeError:
  47. pass
  48. offset += length
  49. return value, offset
  50. else:
  51. def decode_short_string(encoded, offset):
  52. """Decode a short string value from ``encoded`` data at ``offset``.
  53. """
  54. length = struct.unpack_from('B', encoded, offset)[0]
  55. offset += 1
  56. value = encoded[offset:offset + length]
  57. try:
  58. value = value.decode('utf8')
  59. except UnicodeDecodeError:
  60. pass
  61. offset += length
  62. return value, offset
  63. def encode_table(pieces, table):
  64. """Encode a dict as an AMQP table appending the encded table to the
  65. pieces list passed in.
  66. :param list pieces: Already encoded frame pieces
  67. :param dict table: The dict to encode
  68. :rtype: int
  69. """
  70. table = table or {}
  71. length_index = len(pieces)
  72. pieces.append(None) # placeholder
  73. tablesize = 0
  74. for (key, value) in table.items():
  75. tablesize += encode_short_string(pieces, key)
  76. tablesize += encode_value(pieces, value)
  77. pieces[length_index] = struct.pack('>I', tablesize)
  78. return tablesize + 4
  79. def encode_value(pieces, value): # pylint: disable=R0911
  80. """Encode the value passed in and append it to the pieces list returning
  81. the the size of the encoded value.
  82. :param list pieces: Already encoded values
  83. :param any value: The value to encode
  84. :rtype: int
  85. """
  86. if PY2:
  87. if isinstance(value, basestring):
  88. if isinstance(value, unicode_type):
  89. value = value.encode('utf-8')
  90. pieces.append(struct.pack('>cI', b'S', len(value)))
  91. pieces.append(value)
  92. return 5 + len(value)
  93. else:
  94. # support only str on Python 3
  95. if isinstance(value, basestring):
  96. value = value.encode('utf-8')
  97. pieces.append(struct.pack('>cI', b'S', len(value)))
  98. pieces.append(value)
  99. return 5 + len(value)
  100. if isinstance(value, bytes):
  101. pieces.append(struct.pack('>cI', b'x', len(value)))
  102. pieces.append(value)
  103. return 5 + len(value)
  104. if isinstance(value, bool):
  105. pieces.append(struct.pack('>cB', b't', int(value)))
  106. return 2
  107. if isinstance(value, long):
  108. pieces.append(struct.pack('>cq', b'l', value))
  109. return 9
  110. elif isinstance(value, int):
  111. with warnings.catch_warnings():
  112. warnings.filterwarnings('error')
  113. try:
  114. packed = struct.pack('>ci', b'I', value)
  115. pieces.append(packed)
  116. return 5
  117. except (struct.error, DeprecationWarning):
  118. packed = struct.pack('>cq', b'l', long(value))
  119. pieces.append(packed)
  120. return 9
  121. elif isinstance(value, decimal.Decimal):
  122. value = value.normalize()
  123. if value.as_tuple().exponent < 0:
  124. decimals = -value.as_tuple().exponent
  125. raw = int(value * (decimal.Decimal(10)**decimals))
  126. pieces.append(struct.pack('>cBi', b'D', decimals, raw))
  127. else:
  128. # per spec, the "decimals" octet is unsigned (!)
  129. pieces.append(struct.pack('>cBi', b'D', 0, int(value)))
  130. return 6
  131. elif isinstance(value, datetime):
  132. pieces.append(
  133. struct.pack('>cQ', b'T', calendar.timegm(value.utctimetuple())))
  134. return 9
  135. elif isinstance(value, dict):
  136. pieces.append(struct.pack('>c', b'F'))
  137. return 1 + encode_table(pieces, value)
  138. elif isinstance(value, list):
  139. list_pieces = []
  140. for val in value:
  141. encode_value(list_pieces, val)
  142. piece = b''.join(list_pieces)
  143. pieces.append(struct.pack('>cI', b'A', len(piece)))
  144. pieces.append(piece)
  145. return 5 + len(piece)
  146. elif value is None:
  147. pieces.append(struct.pack('>c', b'V'))
  148. return 1
  149. else:
  150. raise exceptions.UnsupportedAMQPFieldException(pieces, value)
  151. def decode_table(encoded, offset):
  152. """Decode the AMQP table passed in from the encoded value returning the
  153. decoded result and the number of bytes read plus the offset.
  154. :param str encoded: The binary encoded data to decode
  155. :param int offset: The starting byte offset
  156. :rtype: tuple
  157. """
  158. result = {}
  159. tablesize = struct.unpack_from('>I', encoded, offset)[0]
  160. offset += 4
  161. limit = offset + tablesize
  162. while offset < limit:
  163. key, offset = decode_short_string(encoded, offset)
  164. value, offset = decode_value(encoded, offset)
  165. result[key] = value
  166. return result, offset
  167. def decode_value(encoded, offset): # pylint: disable=R0912,R0915
  168. """Decode the value passed in returning the decoded value and the number
  169. of bytes read in addition to the starting offset.
  170. :param str encoded: The binary encoded data to decode
  171. :param int offset: The starting byte offset
  172. :rtype: tuple
  173. :raises: pika.exceptions.InvalidFieldTypeException
  174. """
  175. # slice to get bytes in Python 3 and str in Python 2
  176. kind = encoded[offset:offset + 1]
  177. offset += 1
  178. # Bool
  179. if kind == b't':
  180. value = struct.unpack_from('>B', encoded, offset)[0]
  181. value = bool(value)
  182. offset += 1
  183. # Short-Short Int
  184. elif kind == b'b':
  185. value = struct.unpack_from('>B', encoded, offset)[0]
  186. offset += 1
  187. # Short-Short Unsigned Int
  188. elif kind == b'B':
  189. value = struct.unpack_from('>b', encoded, offset)[0]
  190. offset += 1
  191. # Short Int
  192. elif kind == b'U':
  193. value = struct.unpack_from('>h', encoded, offset)[0]
  194. offset += 2
  195. # Short Unsigned Int
  196. elif kind == b'u':
  197. value = struct.unpack_from('>H', encoded, offset)[0]
  198. offset += 2
  199. # Long Int
  200. elif kind == b'I':
  201. value = struct.unpack_from('>i', encoded, offset)[0]
  202. offset += 4
  203. # Long Unsigned Int
  204. elif kind == b'i':
  205. value = struct.unpack_from('>I', encoded, offset)[0]
  206. offset += 4
  207. # Long-Long Int
  208. elif kind == b'L':
  209. value = long(struct.unpack_from('>q', encoded, offset)[0])
  210. offset += 8
  211. # Long-Long Unsigned Int
  212. elif kind == b'l':
  213. value = long(struct.unpack_from('>Q', encoded, offset)[0])
  214. offset += 8
  215. # Float
  216. elif kind == b'f':
  217. value = long(struct.unpack_from('>f', encoded, offset)[0])
  218. offset += 4
  219. # Double
  220. elif kind == b'd':
  221. value = long(struct.unpack_from('>d', encoded, offset)[0])
  222. offset += 8
  223. # Decimal
  224. elif kind == b'D':
  225. decimals = struct.unpack_from('B', encoded, offset)[0]
  226. offset += 1
  227. raw = struct.unpack_from('>i', encoded, offset)[0]
  228. offset += 4
  229. value = decimal.Decimal(raw) * (decimal.Decimal(10)**-decimals)
  230. # Short String
  231. elif kind == b's':
  232. value, offset = decode_short_string(encoded, offset)
  233. # Long String
  234. elif kind == b'S':
  235. length = struct.unpack_from('>I', encoded, offset)[0]
  236. offset += 4
  237. value = encoded[offset:offset + length]
  238. try:
  239. value = value.decode('utf8')
  240. except UnicodeDecodeError:
  241. pass
  242. offset += length
  243. elif kind == b'x':
  244. length = struct.unpack_from('>I', encoded, offset)[0]
  245. offset += 4
  246. value = encoded[offset:offset + length]
  247. offset += length
  248. # Field Array
  249. elif kind == b'A':
  250. length = struct.unpack_from('>I', encoded, offset)[0]
  251. offset += 4
  252. offset_end = offset + length
  253. value = []
  254. while offset < offset_end:
  255. val, offset = decode_value(encoded, offset)
  256. value.append(val)
  257. # Timestamp
  258. elif kind == b'T':
  259. value = datetime.utcfromtimestamp(
  260. struct.unpack_from('>Q', encoded, offset)[0])
  261. offset += 8
  262. # Field Table
  263. elif kind == b'F':
  264. (value, offset) = decode_table(encoded, offset)
  265. # Null / Void
  266. elif kind == b'V':
  267. value = None
  268. else:
  269. raise exceptions.InvalidFieldTypeException(kind)
  270. return value, offset