import json import msgpack from .base_hub_protocol import BaseHubProtocol from ..messages.handshake.request import HandshakeRequestMessage from ..messages.handshake.response import HandshakeResponseMessage from ..messages.invocation_message\ import InvocationMessage, InvocationClientStreamMessage # 1 from ..messages.stream_item_message import StreamItemMessage # 2 from ..messages.completion_message import CompletionMessage # 3 from ..messages.stream_invocation_message import StreamInvocationMessage # 4 from ..messages.cancel_invocation_message import CancelInvocationMessage # 5 from ..messages.ping_message import PingMessage # 6 from ..messages.close_message import CloseMessage # 7 from ..helpers import Helpers class MessagePackHubProtocol(BaseHubProtocol): _priority = [ "type", "headers", "invocation_id", "target", "arguments", "item", "result_kind", "result", "stream_ids" ] def __init__(self): super(MessagePackHubProtocol, self).__init__( "messagepack", 1, "Text", chr(0x1E)) self.logger = Helpers.get_logger() def parse_messages(self, raw): try: messages = [] offset = 0 while offset < len(raw): length = msgpack.unpackb(raw[offset: offset + 1]) values = msgpack.unpackb(raw[offset + 1: offset + length + 1]) offset = offset + length + 1 message = self._decode_message(values) messages.append(message) except Exception as ex: Helpers.get_logger().error("Parse messages Error {0}".format(ex)) Helpers.get_logger().error("raw msg '{0}'".format(raw)) return messages def decode_handshake(self, raw_message): try: has_various_messages = 0x1E in raw_message handshake_data = raw_message[0: raw_message.index(0x1E)] if has_various_messages else raw_message messages = self.parse_messages(raw_message[raw_message.index(0x1E) + 1:]) if has_various_messages else [] data = json.loads(handshake_data) return HandshakeResponseMessage(data.get("error", None)), messages except Exception as ex: if type(raw_message) is str: data = json.loads(raw_message[0: raw_message.index("}") + 1]) return HandshakeResponseMessage(data.get("error", None)), [] Helpers.get_logger().error(raw_message) Helpers.get_logger().error(ex) raise ex def encode(self, message): if type(message) is HandshakeRequestMessage: content = json.dumps(message.__dict__) return content + self.record_separator msg = self._encode_message(message) encoded_message = msgpack.packb(msg) varint_length = self._to_varint(len(encoded_message)) return varint_length + encoded_message def _encode_message(self, message): result = [] # sort attributes for attribute in self._priority: if hasattr(message, attribute): if (attribute == "type"): result.append(getattr(message, attribute).value) else: result.append(getattr(message, attribute)) return result def _decode_message(self, raw): # {} {"error"} # [1, Headers, InvocationId, Target, [Arguments], [StreamIds]] # [2, Headers, InvocationId, Item] # [3, Headers, InvocationId, ResultKind, Result] # [4, Headers, InvocationId, Target, [Arguments], [StreamIds]] # [5, Headers, InvocationId] # [6] # [7, Error, AllowReconnect?] if raw[0] == 1: # InvocationMessage if len(raw[5]) > 0: return InvocationClientStreamMessage( headers=raw[1], stream_ids=raw[5], target=raw[3], arguments=raw[4]) else: return InvocationMessage( headers=raw[1], invocation_id=raw[2], target=raw[3], arguments=raw[4]) elif raw[0] == 2: # StreamItemMessage return StreamItemMessage( headers=raw[1], invocation_id=raw[2], item=raw[3]) elif raw[0] == 3: # CompletionMessage result_kind = raw[3] if result_kind == 1: return CompletionMessage( headers=raw[1], invocation_id=raw[2], result=None, error=raw[4]) elif result_kind == 2: return CompletionMessage( headers=raw[1], invocation_id=raw[2], result=None, error=None) elif result_kind == 3: return CompletionMessage( headers=raw[1], invocation_id=raw[2], result=raw[4], error=None) else: raise Exception("Unknown result kind.") elif raw[0] == 4: # StreamInvocationMessage return StreamInvocationMessage( headers=raw[1], invocation_id=raw[2], target=raw[3], arguments=raw[4]) # stream_id missing? elif raw[0] == 5: # CancelInvocationMessage return CancelInvocationMessage( headers=raw[1], invocation_id=raw[2]) elif raw[0] == 6: # PingMessageEncoding return PingMessage() elif raw[0] == 7: # CloseMessageEncoding return CloseMessage(error=raw[1]) # AllowReconnect is missing print(".......................................") print(raw) print("---------------------------------------") raise Exception("Unknown message type.") def _to_varint(self, value): buffer = b'' while True: byte = value & 0x7f value >>= 7 if value: buffer += bytes((byte | 0x80, )) else: buffer += bytes((byte, )) break return buffer