1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3
4from . import default_id_generator
5from .. import (
6    RPCError,
7    RPCErrorResponse,
8    RPCProtocol,
9    RPCRequest,
10    RPCResponse,
11    InvalidRequestError,
12    MethodNotFoundError,
13    InvalidReplyError,
14)
15
16import msgpack
17import six
18
19from typing import Any, Dict, List, Optional, Tuple, Union, Generator
20
21
22class FixedErrorMessageMixin(object):
23    def __init__(self, *args, **kwargs):
24        if not args:
25            args = [self.message]
26
27        self.request_id = kwargs.pop("request_id", None)
28        super(FixedErrorMessageMixin, self).__init__(*args, **kwargs)
29
30    def error_respond(self):
31        response = MSGPACKRPCErrorResponse()
32
33        response.error = self.message
34        response.unique_id = self.request_id
35        response._msgpackrpc_error_code = self.msgpackrpc_error_code
36        return response
37
38
39class MSGPACKRPCParseError(FixedErrorMessageMixin, InvalidRequestError):
40    msgpackrpc_error_code = -32700
41    message = "Parse error"
42
43
44class MSGPACKRPCInvalidRequestError(FixedErrorMessageMixin, InvalidRequestError):
45    msgpackrpc_error_code = -32600
46    message = "Invalid request"
47
48
49class MSGPACKRPCMethodNotFoundError(FixedErrorMessageMixin, MethodNotFoundError):
50    msgpackrpc_error_code = -32601
51    message = "Method not found"
52
53
54class MSGPACKRPCInvalidParamsError(FixedErrorMessageMixin, InvalidRequestError):
55    msgpackrpc_error_code = -32602
56    message = "Invalid params"
57
58
59class MSGPACKRPCInternalError(FixedErrorMessageMixin, InvalidRequestError):
60    msgpackrpc_error_code = -32603
61    message = "Internal error"
62
63
64class MSGPACKRPCServerError(FixedErrorMessageMixin, InvalidRequestError):
65    msgpackrpc_error_code = -32000
66    message = ""
67
68
69class MSGPACKRPCError(FixedErrorMessageMixin, RPCError):
70    """Reconstructs (to some extend) the server-side exception.
71
72    The client creates this exception by providing it with the ``error``
73    attribute of the MSGPACK error response object returned by the server.
74
75    :param error: This tuple contains the error specification: the numeric error
76        code and the error description.
77    """
78
79    def __init__(
80        self, error: Union["MSGPACKRPCErrorResponse", Tuple[int, str]]
81    ) -> None:
82        if isinstance(error, MSGPACKRPCErrorResponse):
83            super().__init__(error.error)
84            self._msgpackrpc_error_code = error._msgpackrpc_error_code
85        else:
86            super().__init__()
87            self._msgpackrpc_error_code, self.message = error
88
89
90class MSGPACKRPCSuccessResponse(RPCResponse):
91    def _to_list(self):
92        return [1, self.unique_id, None, self.result]
93
94    def serialize(self):
95        return msgpack.packb(self._to_list(), use_bin_type=True)
96
97
98class MSGPACKRPCErrorResponse(RPCErrorResponse):
99    def _to_list(self):
100        return [1, self.unique_id, [self._msgpackrpc_error_code, str(self.error)], None]
101
102    def serialize(self):
103        return msgpack.packb(self._to_list(), use_bin_type=True)
104
105
106def _get_code_and_message(error):
107    assert isinstance(error, (Exception, six.string_types))
108    if isinstance(error, Exception):
109        if hasattr(error, "msgpackrpc_error_code"):
110            code = error.msgpackrpc_error_code
111            msg = str(error)
112        elif isinstance(error, InvalidRequestError):
113            code = MSGPACKRPCInvalidRequestError.msgpackrpc_error_code
114            msg = MSGPACKRPCInvalidRequestError.message
115        elif isinstance(error, MethodNotFoundError):
116            code = MSGPACKRPCMethodNotFoundError.msgpackrpc_error_code
117            msg = MSGPACKRPCMethodNotFoundError.message
118        else:
119            # allow exception message to propagate
120            code = MSGPACKRPCServerError.msgpackrpc_error_code
121            msg = str(error)
122    else:
123        code = -32000
124        msg = error
125
126    return code, msg
127
128
129class MSGPACKRPCRequest(RPCRequest):
130    """Defines a MSGPACK-RPC request."""
131
132    def __init__(self):
133        super().__init__()
134        self.one_way = False
135        """Request or Notification.
136
137        :type: bool
138
139        This flag indicates if the client expects to receive a reply (request: ``one_way = False``)
140        or not (notification: ``one_way = True``).
141
142        Note that it is possible for the server to return an error response.
143        For example if the request becomes unreadable and the server is not able to determine that it is
144        in fact a notification an error should be returned. However, once the server had verified that the
145        request is a notification no reply (not even an error) should be returned.
146        """
147
148        self.unique_id = None
149        """Correlation ID used to match request and response.
150
151        :type: int
152
153        Generated by the client, the server copies it from request to corresponding response.
154        """
155
156        self.method = None
157        """The name of the RPC function to be called.
158
159        :type: str
160
161        The :py:attr:`method` attribute uses the name of the function as it is known by the public.
162        The :py:class:`~tinyrpc.dispatch.RPCDispatcher` allows the use of public aliases in the
163        ``@public`` decorators.
164        These are the names used in the :py:attr:`method` attribute.
165        """
166
167        self.args = []
168        """The positional arguments of the method call.
169
170        :type: list
171
172        The contents of this list are the positional parameters for the :py:attr:`method` called.
173        It is eventually called as ``method(*args)``.
174        """
175
176    def error_respond(
177        self, error: Union[Exception, str]
178    ) -> Optional["MSGPACKRPCErrorResponse"]:
179        """Create an error response to this request.
180
181        When processing the request produces an error condition this method can be used to
182        create the error response object.
183
184        :param error: Specifies what error occurred.
185        :type error: Exception or str
186        :returns: An error response object that can be serialized and sent to the client.
187        :rtype: ;py:class:`MSGPACKRPCErrorResponse`
188        """
189        if not self.unique_id:
190            return None
191
192        response = MSGPACKRPCErrorResponse()
193        response.unique_id = None if self.one_way else self.unique_id
194
195        code, msg = _get_code_and_message(error)
196
197        response.error = msg
198        response._msgpackrpc_error_code = code
199        return response
200
201    def respond(self, result: Any) -> Optional["MSGPACKRPCSuccessResponse"]:
202        """Create a response to this request.
203
204        When processing the request completed successfully this method can be used to
205        create a response object.
206
207        :param result: The result of the invoked method.
208        :type result: Anything that can be encoded by MSGPACK.
209        :returns: A response object that can be serialized and sent to the client.
210        :rtype: :py:class:`MSGPACKRPCSuccessResponse`
211        """
212        if self.one_way or self.unique_id is None:
213            return None
214
215        response = MSGPACKRPCSuccessResponse()
216
217        response.result = result
218        response.unique_id = self.unique_id
219
220        return response
221
222    def _to_list(self):
223        if self.one_way or self.unique_id is None:
224            return [2, self.method, self.args if self.args is not None else []]
225        else:
226            return [
227                0,
228                self.unique_id,
229                self.method,
230                self.args if self.args is not None else [],
231            ]
232
233    def serialize(self) -> bytes:
234        return msgpack.packb(self._to_list(), use_bin_type=True)
235
236
237class MSGPACKRPCProtocol(RPCProtocol):
238    """MSGPACKRPC protocol implementation."""
239
240    def __init__(
241            self,
242            id_generator: Generator[object, None, None] = default_id_generator(),
243            *args,
244            **kwargs
245    ) -> None:
246        super(MSGPACKRPCProtocol, self).__init__(*args, **kwargs)
247        self._id_generator = id_generator
248
249    def _get_unique_id(self):
250        return next(self._id_generator)
251
252    def request_factory(self) -> "MSGPACKRPCRequest":
253        """Factory for request objects.
254
255        Allows derived classes to use requests derived from :py:class:`MSGPACKRPCRequest`.
256
257        :rtype: :py:class:`MSGPACKRPCRequest`
258        """
259        return MSGPACKRPCRequest()
260
261    def create_request(
262        self,
263        method: str,
264        args: List[Any] = None,
265        kwargs: Dict[str, Any] = None,
266        one_way: bool = False,
267    ) -> "MSGPACKRPCRequest":
268        """Creates a new :py:class:`MSGPACKRPCRequest` object.
269
270        Called by the client when constructing a request.
271        MSGPACK-RPC allows only the ``args`` argument to be set; keyword
272        arguments are not supported.
273
274        :param str method: The method name to invoke.
275        :param list args: The positional arguments to call the method with.
276        :param dict kwargs: The keyword arguments to call the method with; must
277            be ``None`` as the protocol does not support keyword arguments.
278        :param bool one_way: The request is an update, i.e. it does not expect a reply.
279        :return: A new request instance
280        :rtype: :py:class:`MSGPACKRPCRequest`
281        :raises InvalidRequestError: when ``kwargs`` is defined.
282        """
283        if kwargs:
284            raise MSGPACKRPCInvalidRequestError("Does not support kwargs")
285
286        request = self.request_factory()
287        request.one_way = one_way
288
289        if not one_way:
290            request.unique_id = self._get_unique_id()
291
292        request.method = method
293        request.args = list(args) if args is not None else []
294        request.kwargs = None
295
296        return request
297
298    def parse_reply(
299        self, data: bytes
300    ) -> Union["MSGPACKRPCSuccessResponse", "MSGPACKRPCErrorResponse"]:
301        """De-serializes and validates a response.
302
303        Called by the client to reconstruct the serialized :py:class:`MSGPACKRPCResponse`.
304
305        :param bytes data: The data stream received by the transport layer containing the
306            serialized response.
307        :return: A reconstructed response.
308        :rtype: :py:class:`MSGPACKRPCSuccessResponse` or :py:class:`MSGPACKRPCErrorResponse`
309        :raises InvalidReplyError: if the response is not valid MSGPACK or does not conform
310            to the standard.
311        """
312        try:
313            rep = msgpack.unpackb(data, raw=False)
314        except Exception as e:
315            raise InvalidReplyError(e)
316
317        if len(rep) != 4:
318            raise InvalidReplyError("MSGPACKRPC spec requires reply of length 4")
319
320        if rep[0] != 1:
321            raise InvalidReplyError("Invalid MSGPACK message type")
322
323        if not isinstance(rep[1], int):
324            raise InvalidReplyError("Invalid or missing message ID in response")
325
326        if rep[2] is not None and rep[3] is not None:
327            raise InvalidReplyError("Reply must contain only one of result and error.")
328
329        if rep[2] is not None:
330            response = MSGPACKRPCErrorResponse()
331            if isinstance(rep[2], list) and len(rep[2]) == 2:
332                code, message = rep[2]
333                if isinstance(code, int) and isinstance(message, str):
334                    response.error = str(message)
335                    response._msgpackrpc_error_code = int(code)
336                else:
337                    response.error = rep[2]
338                    response._msgpackrpc_error_code = None
339            else:
340                response.error = rep[2]
341                response._msgpackrpc_error_code = None
342        else:
343            response = MSGPACKRPCSuccessResponse()
344            response.result = rep[3]
345
346        response.unique_id = rep[1]
347
348        return response
349
350    def parse_request(self, data: bytes) -> "MSGPACKRPCRequest":
351        """De-serializes and validates a request.
352
353        Called by the server to reconstruct the serialized :py:class:`MSGPACKRPCRequest`.
354
355        :param bytes data: The data stream received by the transport layer containing the
356            serialized request.
357        :return: A reconstructed request.
358        :rtype: :py:class:`MSGPACKRPCRequest`
359        :raises MSGPACKRPCParseError: if the ``data`` cannot be parsed as valid MSGPACK.
360        :raises MSGPACKRPCInvalidRequestError: if the request does not comply with the standard.
361        """
362        try:
363            req = msgpack.unpackb(data, raw=False)
364        except Exception:
365            raise MSGPACKRPCParseError()
366
367        if not isinstance(req, list):
368            raise MSGPACKRPCInvalidRequestError()
369
370        if len(req) < 2:
371            raise MSGPACKRPCInvalidRequestError()
372
373        if req[0] == 0:
374            # MSGPACK request
375            request_id = req[1]
376            if not isinstance(request_id, int):
377                raise MSGPACKRPCInvalidRequestError()
378
379            if len(req) == 4:
380                return self._parse_request(req)
381            else:
382                raise MSGPACKRPCInvalidRequestError(request_id=request_id)
383        elif req[0] == 2:
384            # MSGPACK notification
385            if len(req) == 3:
386                return self._parse_notification(req)
387            else:
388                raise MSGPACKRPCInvalidRequestError()
389        else:
390            raise MSGPACKRPCInvalidRequestError()
391
392    def _parse_notification(self, req):
393        if not isinstance(req[1], six.string_types):
394            raise MSGPACKRPCInvalidRequestError()
395
396        request = MSGPACKRPCRequest()
397        request.one_way = True
398        request.method = req[1]
399
400        params = req[2]
401        # params should not be None according to the spec; if there are
402        # no params, an empty array must be used
403        if isinstance(params, list):
404            request.args = params
405        else:
406            raise MSGPACKRPCInvalidParamsError(request_id=req[1])
407
408        return request
409
410    def _parse_request(self, req):
411        if not isinstance(req[2], six.string_types):
412            raise MSGPACKRPCInvalidRequestError(request_id=req[1])
413
414        request = MSGPACKRPCRequest()
415        request.one_way = False
416        request.method = req[2]
417        request.unique_id = req[1]
418
419        params = req[3]
420        # params should not be None according to the spec; if there are
421        # no params, an empty array must be used
422        if isinstance(params, list):
423            request.args = params
424        else:
425            raise MSGPACKRPCInvalidParamsError(request_id=req[1])
426
427        return request
428
429    def raise_error(
430        self, error: Union["MSGPACKRPCErrorResponse", Dict[str, Any]]
431    ) -> "MSGPACKRPCError":
432        """Recreates the exception.
433
434        Creates a :py:class:`~tinyrpc.protocols.msgpackrpc.MSGPACKRPCError`
435        instance and raises it.
436
437        This allows the error code and the message of the original exception to
438        propagate into the client code.
439
440        The :py:attr:`~tinyrpc.protocols.MSGPACKProtocol.raises_error` flag
441        controls if the exception object is raised or returned.
442
443        :returns: the exception object if it is not allowed to raise it.
444        :raises MSGPACKRPCError: when the exception can be raised.
445            The exception object will contain ``message`` and ``code``.
446        """
447        exc = MSGPACKRPCError(error)
448        if self.raises_errors:
449            raise exc
450        return exc
451