serialization.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458
  1. """Serialization utilities."""
  2. from __future__ import absolute_import, unicode_literals
  3. import codecs
  4. import os
  5. import sys
  6. import pickle as pypickle
  7. try:
  8. import cPickle as cpickle
  9. except ImportError: # pragma: no cover
  10. cpickle = None # noqa
  11. from collections import namedtuple
  12. from contextlib import contextmanager
  13. from io import BytesIO
  14. from .exceptions import (
  15. ContentDisallowed, DecodeError, EncodeError, SerializerNotInstalled
  16. )
  17. from .five import reraise, text_t
  18. from .utils.compat import entrypoints
  19. from .utils.encoding import bytes_to_str, str_to_bytes, bytes_t
  20. __all__ = ('pickle', 'loads', 'dumps', 'register', 'unregister')
  21. SKIP_DECODE = frozenset(['binary', 'ascii-8bit'])
  22. TRUSTED_CONTENT = frozenset(['application/data', 'application/text'])
  23. if sys.platform.startswith('java'): # pragma: no cover
  24. def _decode(t, coding):
  25. return codecs.getdecoder(coding)(t)[0]
  26. else:
  27. _decode = codecs.decode
  28. pickle = cpickle or pypickle
  29. pickle_load = pickle.load
  30. #: Kombu requires Python 2.5 or later so we use protocol 2 by default.
  31. #: There's a new protocol (3) but this is only supported by Python 3.
  32. pickle_protocol = int(os.environ.get('PICKLE_PROTOCOL', 2))
  33. codec = namedtuple('codec', ('content_type', 'content_encoding', 'encoder'))
  34. @contextmanager
  35. def _reraise_errors(wrapper,
  36. include=(Exception,), exclude=(SerializerNotInstalled,)):
  37. try:
  38. yield
  39. except exclude:
  40. raise
  41. except include as exc:
  42. reraise(wrapper, wrapper(exc), sys.exc_info()[2])
  43. def pickle_loads(s, load=pickle_load):
  44. # used to support buffer objects
  45. return load(BytesIO(s))
  46. def parenthesize_alias(first, second):
  47. return '%s (%s)' % (first, second) if first else second
  48. class SerializerRegistry(object):
  49. """The registry keeps track of serialization methods."""
  50. def __init__(self):
  51. self._encoders = {}
  52. self._decoders = {}
  53. self._default_encode = None
  54. self._default_content_type = None
  55. self._default_content_encoding = None
  56. self._disabled_content_types = set()
  57. self.type_to_name = {}
  58. self.name_to_type = {}
  59. def register(self, name, encoder, decoder, content_type,
  60. content_encoding='utf-8'):
  61. """Register a new encoder/decoder.
  62. Arguments:
  63. name (str): A convenience name for the serialization method.
  64. encoder (callable): A method that will be passed a python data
  65. structure and should return a string representing the
  66. serialized data. If :const:`None`, then only a decoder
  67. will be registered. Encoding will not be possible.
  68. decoder (Callable): A method that will be passed a string
  69. representing serialized data and should return a python
  70. data structure. If :const:`None`, then only an encoder
  71. will be registered. Decoding will not be possible.
  72. content_type (str): The mime-type describing the serialized
  73. structure.
  74. content_encoding (str): The content encoding (character set) that
  75. the `decoder` method will be returning. Will usually be
  76. `utf-8`, `us-ascii`, or `binary`.
  77. """
  78. if encoder:
  79. self._encoders[name] = codec(
  80. content_type, content_encoding, encoder,
  81. )
  82. if decoder:
  83. self._decoders[content_type] = decoder
  84. self.type_to_name[content_type] = name
  85. self.name_to_type[name] = content_type
  86. def enable(self, name):
  87. if '/' not in name:
  88. name = self.name_to_type[name]
  89. self._disabled_content_types.discard(name)
  90. def disable(self, name):
  91. if '/' not in name:
  92. name = self.name_to_type[name]
  93. self._disabled_content_types.add(name)
  94. def unregister(self, name):
  95. """Unregister registered encoder/decoder.
  96. Arguments:
  97. name (str): Registered serialization method name.
  98. Raises:
  99. SerializerNotInstalled: If a serializer by that name
  100. cannot be found.
  101. """
  102. try:
  103. content_type = self.name_to_type[name]
  104. self._decoders.pop(content_type, None)
  105. self._encoders.pop(name, None)
  106. self.type_to_name.pop(content_type, None)
  107. self.name_to_type.pop(name, None)
  108. except KeyError:
  109. raise SerializerNotInstalled(
  110. 'No encoder/decoder installed for {0}'.format(name))
  111. def _set_default_serializer(self, name):
  112. """Set the default serialization method used by this library.
  113. Arguments:
  114. name (str): The name of the registered serialization method.
  115. For example, `json` (default), `pickle`, `yaml`, `msgpack`,
  116. or any custom methods registered using :meth:`register`.
  117. Raises:
  118. SerializerNotInstalled: If the serialization method
  119. requested is not available.
  120. """
  121. try:
  122. (self._default_content_type, self._default_content_encoding,
  123. self._default_encode) = self._encoders[name]
  124. except KeyError:
  125. raise SerializerNotInstalled(
  126. 'No encoder installed for {0}'.format(name))
  127. def dumps(self, data, serializer=None):
  128. """Encode data.
  129. Serialize a data structure into a string suitable for sending
  130. as an AMQP message body.
  131. Arguments:
  132. data (List, Dict, str): The message data to send.
  133. serializer (str): An optional string representing
  134. the serialization method you want the data marshalled
  135. into. (For example, `json`, `raw`, or `pickle`).
  136. If :const:`None` (default), then json will be used, unless
  137. `data` is a :class:`str` or :class:`unicode` object. In this
  138. latter case, no serialization occurs as it would be
  139. unnecessary.
  140. Note that if `serializer` is specified, then that
  141. serialization method will be used even if a :class:`str`
  142. or :class:`unicode` object is passed in.
  143. Returns:
  144. Tuple[str, str, str]: A three-item tuple containing the
  145. content type (e.g., `application/json`), content encoding, (e.g.,
  146. `utf-8`) and a string containing the serialized data.
  147. Raises:
  148. SerializerNotInstalled: If the serialization method
  149. requested is not available.
  150. """
  151. if serializer == 'raw':
  152. return raw_encode(data)
  153. if serializer and not self._encoders.get(serializer):
  154. raise SerializerNotInstalled(
  155. 'No encoder installed for {0}'.format(serializer))
  156. # If a raw string was sent, assume binary encoding
  157. # (it's likely either ASCII or a raw binary file, and a character
  158. # set of 'binary' will encompass both, even if not ideal.
  159. if not serializer and isinstance(data, bytes_t):
  160. # In Python 3+, this would be "bytes"; allow binary data to be
  161. # sent as a message without getting encoder errors
  162. return 'application/data', 'binary', data
  163. # For Unicode objects, force it into a string
  164. if not serializer and isinstance(data, text_t):
  165. with _reraise_errors(EncodeError, exclude=()):
  166. payload = data.encode('utf-8')
  167. return 'text/plain', 'utf-8', payload
  168. if serializer:
  169. content_type, content_encoding, encoder = \
  170. self._encoders[serializer]
  171. else:
  172. encoder = self._default_encode
  173. content_type = self._default_content_type
  174. content_encoding = self._default_content_encoding
  175. with _reraise_errors(EncodeError):
  176. payload = encoder(data)
  177. return content_type, content_encoding, payload
  178. def loads(self, data, content_type, content_encoding,
  179. accept=None, force=False, _trusted_content=TRUSTED_CONTENT):
  180. """Decode serialized data.
  181. Deserialize a data stream as serialized using `dumps`
  182. based on `content_type`.
  183. Arguments:
  184. data (bytes, buffer, str): The message data to deserialize.
  185. content_type (str): The content-type of the data.
  186. (e.g., `application/json`).
  187. content_encoding (str): The content-encoding of the data.
  188. (e.g., `utf-8`, `binary`, or `us-ascii`).
  189. accept (Set): List of content-types to accept.
  190. Raises:
  191. ContentDisallowed: If the content-type is not accepted.
  192. Returns:
  193. Any: The unserialized data.
  194. """
  195. content_type = (bytes_to_str(content_type) if content_type
  196. else 'application/data')
  197. if accept is not None:
  198. if content_type not in _trusted_content \
  199. and content_type not in accept:
  200. raise self._for_untrusted_content(content_type, 'untrusted')
  201. else:
  202. if content_type in self._disabled_content_types and not force:
  203. raise self._for_untrusted_content(content_type, 'disabled')
  204. content_encoding = (content_encoding or 'utf-8').lower()
  205. if data:
  206. decode = self._decoders.get(content_type)
  207. if decode:
  208. with _reraise_errors(DecodeError):
  209. return decode(data)
  210. if content_encoding not in SKIP_DECODE and \
  211. not isinstance(data, text_t):
  212. with _reraise_errors(DecodeError):
  213. return _decode(data, content_encoding)
  214. return data
  215. def _for_untrusted_content(self, ctype, why):
  216. return ContentDisallowed(
  217. 'Refusing to deserialize {0} content of type {1}'.format(
  218. why,
  219. parenthesize_alias(self.type_to_name.get(ctype, ctype), ctype),
  220. ),
  221. )
  222. #: Global registry of serializers/deserializers.
  223. registry = SerializerRegistry()
  224. dumps = registry.dumps
  225. loads = registry.loads
  226. register = registry.register
  227. unregister = registry.unregister
  228. def raw_encode(data):
  229. """Special case serializer."""
  230. content_type = 'application/data'
  231. payload = data
  232. if isinstance(payload, text_t):
  233. content_encoding = 'utf-8'
  234. with _reraise_errors(EncodeError, exclude=()):
  235. payload = payload.encode(content_encoding)
  236. else:
  237. content_encoding = 'binary'
  238. return content_type, content_encoding, payload
  239. def register_json():
  240. """Register a encoder/decoder for JSON serialization."""
  241. from kombu.utils import json as _json
  242. registry.register('json', _json.dumps, _json.loads,
  243. content_type='application/json',
  244. content_encoding='utf-8')
  245. def register_yaml():
  246. """Register a encoder/decoder for YAML serialization.
  247. It is slower than JSON, but allows for more data types
  248. to be serialized. Useful if you need to send data such as dates
  249. """
  250. try:
  251. import yaml
  252. registry.register('yaml', yaml.safe_dump, yaml.safe_load,
  253. content_type='application/x-yaml',
  254. content_encoding='utf-8')
  255. except ImportError:
  256. def not_available(*args, **kwargs):
  257. """Raise SerializerNotInstalled.
  258. Used in case a client receives a yaml message, but yaml
  259. isn't installed.
  260. """
  261. raise SerializerNotInstalled(
  262. 'No decoder installed for YAML. Install the PyYAML library')
  263. registry.register('yaml', None, not_available, 'application/x-yaml')
  264. if sys.version_info[0] == 3: # pragma: no cover
  265. def unpickle(s):
  266. return pickle_loads(str_to_bytes(s))
  267. else:
  268. unpickle = pickle_loads # noqa
  269. def register_pickle():
  270. """Register pickle serializer.
  271. The fastest serialization method, but restricts
  272. you to python clients.
  273. """
  274. def pickle_dumps(obj, dumper=pickle.dumps):
  275. return dumper(obj, protocol=pickle_protocol)
  276. registry.register('pickle', pickle_dumps, unpickle,
  277. content_type='application/x-python-serialize',
  278. content_encoding='binary')
  279. def register_msgpack():
  280. """Register msgpack serializer.
  281. See Also:
  282. https://msgpack.org/.
  283. """
  284. pack = unpack = None
  285. try:
  286. import msgpack
  287. if msgpack.version >= (0, 4):
  288. from msgpack import packb, unpackb
  289. def pack(s):
  290. return packb(s, use_bin_type=True)
  291. def unpack(s):
  292. return unpackb(s, raw=False)
  293. else:
  294. def version_mismatch(*args, **kwargs):
  295. raise SerializerNotInstalled(
  296. 'msgpack requires msgpack-python >= 0.4.0')
  297. pack = unpack = version_mismatch
  298. except (ImportError, ValueError):
  299. def not_available(*args, **kwargs):
  300. raise SerializerNotInstalled(
  301. 'No decoder installed for msgpack. '
  302. 'Please install the msgpack-python library')
  303. pack = unpack = not_available
  304. registry.register(
  305. 'msgpack', pack, unpack,
  306. content_type='application/x-msgpack',
  307. content_encoding='binary',
  308. )
  309. # Register the base serialization methods.
  310. register_json()
  311. register_pickle()
  312. register_yaml()
  313. register_msgpack()
  314. # Default serializer is 'json'
  315. registry._set_default_serializer('json')
  316. _setupfuns = {
  317. 'json': register_json,
  318. 'pickle': register_pickle,
  319. 'yaml': register_yaml,
  320. 'msgpack': register_msgpack,
  321. 'application/json': register_json,
  322. 'application/x-yaml': register_yaml,
  323. 'application/x-python-serialize': register_pickle,
  324. 'application/x-msgpack': register_msgpack,
  325. }
  326. def enable_insecure_serializers(choices=['pickle', 'yaml', 'msgpack']):
  327. """Enable serializers that are considered to be unsafe.
  328. Note:
  329. Will enable ``pickle``, ``yaml`` and ``msgpack`` by default, but you
  330. can also specify a list of serializers (by name or content type)
  331. to enable.
  332. """
  333. for choice in choices:
  334. try:
  335. registry.enable(choice)
  336. except KeyError:
  337. pass
  338. def disable_insecure_serializers(allowed=['json']):
  339. """Disable untrusted serializers.
  340. Will disable all serializers except ``json``
  341. or you can specify a list of deserializers to allow.
  342. Note:
  343. Producers will still be able to serialize data
  344. in these formats, but consumers will not accept
  345. incoming data using the untrusted content types.
  346. """
  347. for name in registry._decoders:
  348. registry.disable(name)
  349. if allowed is not None:
  350. for name in allowed:
  351. registry.enable(name)
  352. # Insecure serializers are disabled by default since v3.0
  353. disable_insecure_serializers()
  354. # Load entrypoints from installed extensions
  355. for ep, args in entrypoints('kombu.serializers'): # pragma: no cover
  356. register(ep.name, *args)
  357. def prepare_accept_content(l, name_to_type=registry.name_to_type):
  358. if l is not None:
  359. return {n if '/' in n else name_to_type[n] for n in l}
  360. return l