def parse_protobuf(pb_object: Message, data: bytes) -> Message:
"""
Given protobuf object, parse it from the given data.
:param pb_object: the object to load data to.
:param data: data to load to the object.
:return: the protobuf object itself.
"""
pb_object.ParseFromString(data)
return pb_object
def read_message(self, message: Message, message_len_struct: str = 'I') -> Message:
'''
read response from server and serialize it to message
:param message: empty message to serialize to
:param message_len: message len struct, default I
:return: serialized protobuf message
'''
message_len_buffer_size = struct.Struct(message_len_struct).size
message_len_buffer = self.__recv(message_len_buffer_size)
(message_len,) = struct.unpack('>' + message_len_struct, message_len_buffer)
message.ParseFromString(self.__recv(message_len))
return message
async def recv_proto_msg(reader: StreamReader,
to_merge: Message,
max_message_length=None,
aes_key: Optional[bytes] = None):
length_bytes = await reader.readexactly(4)
length = int.from_bytes(length_bytes, 'big')
if max_message_length is not None and length > max_message_length:
raise ValueError('Received message is too large')
message_bytes = await reader.readexactly(length)
if aes_key is not None:
message_bytes = aes_decrypt(message_bytes, aes_key)
to_merge.ParseFromString(message_bytes)
def unpack(self, message: GrpcMessage) -> None:
"""Deserializes the serialized protocol buffer message.
Args:
message (:class:`GrpcMessage`): the protocol buffer message object
to which the response data is deserialized.
Raises:
ValueError: message is not protocol buffer message object or message's type is not
matched with the response data type
"""
if self.content_type is not None and self.content_type.lower() == "application/x-protobuf":
message.ParseFromString(self.data)
return
unpack(self.proto, message)
async def call(
self, cmd: str, request: protobuf.Message = None, response: protobuf.Message = None
) -> None:
if request is not None:
payload = request.SerializeToString()
else:
payload = b''
logger.debug("%s: sending %s to %s...", self.id, cmd, self.__server_address)
start_time = time.time()
serialized_response = await self.__call_internal(
b'C', [cmd.encode('ascii'), payload])
logger.debug(
"%s: %s to %s finished in %.2fmsec",
self.id, cmd, self.__server_address, 1000 * (time.time() - start_time))
if response is not None:
response.ParseFromString(serialized_response)
async def read_pbmsg_safe(stream: asyncio.StreamReader,
pbmsg: PBMessage) -> None:
len_msg_bytes = await read_unsigned_varint(stream)
msg_bytes = await stream.readexactly(len_msg_bytes)
pbmsg.ParseFromString(msg_bytes)
async def read_pbmsg_safe(stream: anyio.abc.SocketStream, pbmsg: PBMessage) -> None:
len_msg_bytes = await read_unsigned_varint(stream)
msg_bytes = await stream.receive_exactly(len_msg_bytes)
pbmsg.ParseFromString(msg_bytes)
def _attempt_parse(obj: ProtobufMessage, data: bytes) -> None:
try:
obj.ParseFromString(data)
except ProtobufDecodeError:
raise ParseError("Incorrect protobuf message")