1"""
2Functions to work with MessagePack
3"""
4import logging
5
6log = logging.getLogger(__name__)
7
8HAS_MSGPACK = False
9try:
10    import msgpack
11
12    # There is a serialization issue on ARM and potentially other platforms for some msgpack bindings, check for it
13    if (
14        msgpack.version >= (0, 4, 0)
15        and msgpack.loads(msgpack.dumps([1, 2, 3], use_bin_type=False), use_list=True)
16        is None
17    ):
18        raise ImportError
19    elif msgpack.loads(msgpack.dumps([1, 2, 3]), use_list=True) is None:
20        raise ImportError
21    HAS_MSGPACK = True
22except ImportError:
23    try:
24        import msgpack_pure as msgpack  # pylint: disable=import-error
25
26        HAS_MSGPACK = True
27    except ImportError:
28        pass
29        # Don't exit if msgpack is not available, this is to make local mode work without msgpack
30        # sys.exit(salt.defaults.exitcodes.EX_GENERIC)
31
32if HAS_MSGPACK and hasattr(msgpack, "exceptions"):
33    exceptions = msgpack.exceptions
34else:
35
36    class PackValueError(Exception):
37        """
38        older versions of msgpack do not have PackValueError
39        """
40
41    class _exceptions:
42        """
43        older versions of msgpack do not have an exceptions module
44        """
45
46        PackValueError = PackValueError()
47
48    exceptions = _exceptions()
49
50# One-to-one mappings
51Packer = msgpack.Packer
52ExtType = msgpack.ExtType
53version = (0, 0, 0) if not HAS_MSGPACK else msgpack.version
54
55
56def _sanitize_msgpack_kwargs(kwargs):
57    """
58    Clean up msgpack keyword arguments based on the version
59    https://github.com/msgpack/msgpack-python/blob/master/ChangeLog.rst
60    """
61    assert isinstance(kwargs, dict)
62    if version < (0, 6, 0) and kwargs.pop("strict_map_key", None) is not None:
63        log.info("removing unsupported `strict_map_key` argument from msgpack call")
64    if version < (0, 5, 2) and kwargs.pop("raw", None) is not None:
65        log.info("removing unsupported `raw` argument from msgpack call")
66    if version < (0, 4, 0) and kwargs.pop("use_bin_type", None) is not None:
67        log.info("removing unsupported `use_bin_type` argument from msgpack call")
68    if version >= (1, 0, 0) and kwargs.pop("encoding", None) is not None:
69        log.debug("removing unsupported `encoding` argument from msgpack call")
70
71    return kwargs
72
73
74def _sanitize_msgpack_unpack_kwargs(kwargs):
75    """
76    Clean up msgpack keyword arguments for unpack operations, based on
77    the version
78    https://github.com/msgpack/msgpack-python/blob/master/ChangeLog.rst
79    """
80    assert isinstance(kwargs, dict)
81    if version >= (1, 0, 0):
82        kwargs.setdefault("raw", True)
83        kwargs.setdefault("strict_map_key", False)
84    return _sanitize_msgpack_kwargs(kwargs)
85
86
87def _add_msgpack_unpack_kwargs(kwargs):
88    """
89    Add any msgpack unpack kwargs here.
90
91    max_buffer_size: will make sure the buffer is set to a minimum
92    of 100MiB in versions >=6 and <1.0
93    """
94    assert isinstance(kwargs, dict)
95    if version >= (0, 6, 0) and version < (1, 0, 0):
96        kwargs["max_buffer_size"] = 100 * 1024 * 1024
97    return _sanitize_msgpack_unpack_kwargs(kwargs)
98
99
100class Unpacker(msgpack.Unpacker):
101    """
102    Wraps the msgpack.Unpacker and removes non-relevant arguments
103    """
104
105    def __init__(self, *args, **kwargs):
106        msgpack.Unpacker.__init__(self, *args, **_add_msgpack_unpack_kwargs(kwargs))
107
108
109def pack(o, stream, **kwargs):
110    """
111    .. versionadded:: 2018.3.4
112
113    Wraps msgpack.pack and ensures that the passed object is unwrapped if it is
114    a proxy.
115
116    By default, this function uses the msgpack module and falls back to
117    msgpack_pure, if the msgpack is not available.
118    """
119    # Writes to a stream, there is no return
120    msgpack.pack(o, stream, **_sanitize_msgpack_kwargs(kwargs))
121
122
123def packb(o, **kwargs):
124    """
125    .. versionadded:: 2018.3.4
126
127    Wraps msgpack.packb and ensures that the passed object is unwrapped if it
128    is a proxy.
129
130    By default, this function uses the msgpack module and falls back to
131    msgpack_pure, if the msgpack is not available.
132    """
133    return msgpack.packb(o, **_sanitize_msgpack_kwargs(kwargs))
134
135
136def unpack(stream, **kwargs):
137    """
138    .. versionadded:: 2018.3.4
139
140    Wraps msgpack.unpack.
141
142    By default, this function uses the msgpack module and falls back to
143    msgpack_pure, if the msgpack is not available.
144    """
145    return msgpack.unpack(stream, **_sanitize_msgpack_unpack_kwargs(kwargs))
146
147
148def unpackb(packed, **kwargs):
149    """
150    .. versionadded:: 2018.3.4
151
152    Wraps msgpack.unpack.
153
154    By default, this function uses the msgpack module and falls back to
155    msgpack_pure.
156    """
157    return msgpack.unpackb(packed, **_sanitize_msgpack_unpack_kwargs(kwargs))
158
159
160# alias for compatibility to simplejson/marshal/pickle.
161load = unpack
162loads = unpackb
163
164dump = pack
165dumps = packb
166