async def send_message(self, msg: message.Message) -> None:
"""Send a protobuf message to the remote."""
if not self._is_socket_open:
raise APIConnectionError("Connection isn't established yet")
for message_type, klass in MESSAGE_TYPE_TO_PROTO.items():
if isinstance(msg, klass):
break
else:
raise ValueError(f"Message type id not found for type {type(msg)}")
encoded = msg.SerializeToString()
_LOGGER.debug("%s: Sending %s: %s", self._params.address, type(msg),
str(msg))
try:
assert self._frame_helper is not None
# pylint: disable=undefined-loop-variable
await self._frame_helper.write_packet(
Packet(
type=message_type,
data=encoded,
))
except SocketAPIError as err: # pylint: disable=broad-except
# If writing packet fails, we don't know what state the frames
# are in anymore and we have to close the connection
await self._report_fatal_error(err)
raise
def checker(proto: Message,
ctx: C.CheckerContext = DEFAULT_CONTEXT) -> Any:
if not isinstance(proto, proto_type):
raise RuntimeError(
'You cannot pass an object that is not of type {}'.format(
proto_type.__name__))
return getattr(C, py_func.__name__)(proto.SerializeToString(), ctx)
def send_message(self, message: Message) -> None:
'''
send protobuf message to server
:param message: protobuf message
'''
packed_len = struct.pack('>I', message.ByteSize())
self.sendall(packed_len)
self.sendall(message.SerializeToString())
def send(self, method: str, request: Message,
waiter: Callable[[RPCMessage], None]) -> None:
"""Send a method invocation message to the peer."""
message = RPCMessage(id=self.__get_next_id(),
method=method,
request=request.SerializeToString())
self.__waiters[message.id] = waiter
self.serialize(message)
def write_packet(self, name: str, message: Message):
ms = current_time_millis()
pack = Packet()
pack.stamp.sec = ms // 1000
pack.stamp.nsec = (ms % 1000) * 1000
pack.type = name
pack.serialized_data = message.SerializeToString()
self.write(pack)
def encode_response(rpc: tuple, response: message.Message) -> bytes:
channel, service, method = _ids(rpc)
return packet_pb2.RpcPacket(
type=packet_pb2.PacketType.RESPONSE,
channel_id=channel,
service_id=service,
method_id=method,
payload=response.SerializeToString()).SerializeToString()
def serialize_protobuf(pb_object: Message) -> bytes:
"""
Serialize protobuf object to bytes.
:param pb_object: object to serialize.
:return: the serialized bytes.
"""
return pb_object.SerializeToString()
def send_to_client(self, message: ProtobufMessage):
"""Send protobuf message to client."""
data = message.SerializeToString()
_LOGGER.info("<<(DECRYPTED): %s", message)
if self.chacha:
data = self.chacha.encrypt(data)
log_binary(_LOGGER, "<<(ENCRYPTED)", Message=message)
length = variant.write_variant(len(data))
self.transport.write(length + data)
def write_tfrecord_file(file_name: Text, proto: Message) -> None:
"""Writes a serialized tfrecord to file."""
try:
import tensorflow as tf # pylint: disable=g-import-not-at-top
except ModuleNotFoundError as e:
raise Exception(
'TensorFlow must be installed to use this functionality.') from e
fileio.makedirs(os.path.dirname(file_name))
with tf.io.TFRecordWriter(file_name) as writer:
writer.write(proto.SerializeToString())
async def send_proto_msg(writer: StreamWriter,
msg: Message,
aes_key: Optional[bytes] = None):
serialized = msg.SerializeToString()
if aes_key is not None:
serialized = aes_encrypt(serialized, aes_key)
writer.write(len(serialized).to_bytes(4, 'big'))
await writer.drain()
writer.write(serialized)
await writer.drain()
def send_to_client(self, message: ProtobufMessage) -> None:
if not self.transport:
return
data = message.SerializeToString()
if self.chacha:
data = self.chacha.encrypt(data)
length = variant.write_variant(len(data))
self.transport.write(length + data)
log_protobuf(_LOGGER, ">> Send: Protobuf", message)
def update(self, updated_msg: Message, call_cb: bool = True) -> None:
serialized_msg = updated_msg.SerializeToString()
query = "INSERT OR REPLACE INTO {} (id, serialized_pb) VALUES(?, ?)".format(
self.table_name)
id_value = getattr(updated_msg, self.id_field.name)
with self.conn:
self.conn.execute(query, (id_value, serialized_msg))
assert id_value is not None
if call_cb:
for cb in self.callbacks:
cb(id_value, serialized_msg)
def _post_proto(url: str, data: message.Message,
out_type: Type[_ProtoOut]) -> _ProtoOut:
response = requests.post(
url,
data=base64.encodebytes(data.SerializeToString()).decode('ascii'),
headers={
'Accept': 'application/x-protobuf-base64',
'Content-type': 'application/x-protobuf-base64',
})
response.raise_for_status()
output = out_type()
output.ParseFromString(base64.decodebytes(response.content))
return output
async def send_message(self, msg: message.Message) -> None:
for message_type, klass in MESSAGE_TYPE_TO_PROTO.items():
if isinstance(msg, klass):
break
else:
raise ValueError
encoded = msg.SerializeToString()
_LOGGER.debug("%s: Sending %s: %s", self._params.address, type(msg), str(msg))
req = bytes([0])
req += _varuint_to_bytes(len(encoded))
req += _varuint_to_bytes(message_type)
req += encoded
await self._write(req)
def process_response(self, response: Message):
protocol = self.get_output_protocol_type()
if protocol == ProtocolType.MSGPACK:
content_type = "application/x-msgpack"
content = msgpack.dumps(MessageToDict(response))
elif protocol == ProtocolType.JSON:
content_type = "application/json"
content = json.dumps(MessageToDict(response))
elif protocol == ProtocolType.PROTOBUF:
content_type = "application/protobuf"
content = response.SerializeToString()
else:
raise ValueError("Unexpected protocol type {}".format(protocol))
return Response(content, content_type=content_type)
def write_proto(message: Message, filename: str):
if not os.path.exists(os.path.dirname(filename)):
try:
os.makedirs(os.path.dirname(filename))
except OSError as exc: # Guard against race condition
if exc.errno != errno.EEXIST:
raise
binary_filename = filename + '.binarypb'
with open(binary_filename, "wb") as f:
f.write(message.SerializeToString())
text_filename = filename + '.textproto'
with open(text_filename, "w") as f:
# These comments point to the proto to use as a schema for IDEs.
f.write("# proto-file: proto/npc_infos.proto\n")
f.write("# proto-message: NpcInfos\n\n")
PrintMessage(message, f)
def _post_to_supervisor(self,
endpoint: str,
message: Message,
private: bool = True) -> List[str]:
responses: List[str] = list()
data = message.SerializeToString()
for client in self._clients:
supervisor_api = (client.service_private_api("supervisor")
if private else
client.service_public_api("supervisor"))
response = supervisor_api.post_service(endpoint,
data,
data_format="binary")
responses.append(response.json())
return responses
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)
def create_proto(proto_message: message.Message) -> packet.Packet:
"""Create a MediaPipe protobuf message packet.
Args:
proto_message: A Python protobuf message.
Returns:
A MediaPipe protobuf message Packet.
Raises:
RuntimeError: If the protobuf message type is not registered in MediaPipe.
Examples:
detection = detection_pb2.Detection()
text_format.Parse('score: 0.5', detection)
packet = mp.packet_creator.create_proto(detection)
output_detection = mp.packet_getter.get_proto(packet)
"""
# pylint:disable=protected-access
return _packet_creator._create_proto(proto_message.DESCRIPTOR.full_name,
proto_message.SerializeToString())
def CreateHunt(
flow_name: str,
flow_args: message.Message,
hunt_runner_args: flows_pb2.HuntRunnerArgs,
context: context_lib.GrrApiContext,
) -> Hunt:
"""Creates a new hunt.
Args:
flow_name: String with a name of a flow that will run on all the clients
in the hunt.
flow_args: Flow arguments to be used. A proto, that depends on a flow.
hunt_runner_args: flows_pb2.HuntRunnerArgs instance. Used to specify
description, client_rule_set, output_plugins and other useful
hunt attributes.
context: API context.
Raises:
ValueError: if flow_name is empty.
Returns:
Hunt object corresponding to the created hunt.
"""
if not flow_name:
raise ValueError("flow_name can't be empty")
request = hunt_pb2.ApiCreateHuntArgs(flow_name=flow_name)
if flow_args:
request.flow_args.value = flow_args.SerializeToString()
request.flow_args.type_url = utils.GetTypeUrl(flow_args)
if hunt_runner_args:
request.hunt_runner_args.CopyFrom(hunt_runner_args)
data = context.SendRequest("CreateHunt", request)
if not isinstance(data, hunt_pb2.ApiHunt):
raise TypeError(f"Unexpected response type: '{type(data)}'")
return Hunt(data=data, context=context)
def send_rpc(route: str, in_msg: message.Message, res_class: message.Message, server_id: str='') -> message.Message:
""" sends a rpc to other pitaya server """
if not issubclass(type(in_msg), message.Message) or not issubclass(res_class, message.Message):
raise TypeError
msg_bytes = in_msg.SerializeToString()
msg_len = len(msg_bytes)
c_bytes = (c_char * msg_len)(*msg_bytes)
ret_ptr = POINTER(MemoryBuffer)()
err = PitayaError()
res = LIB.tfg_pitc_RPC(server_id.encode(
'utf-8'), route.encode('utf-8'), addressof(c_bytes), msg_len, byref(ret_ptr), byref(err))
if not res:
exception_msg = "code: {} msg: {}".format(err.code, err.msg)
LIB.tfg_pitc_FreePitayaError(err)
raise Exception(exception_msg)
ret_bytes = (
c_char * ret_ptr.contents.size).from_address(ret_ptr.contents.data)
response = Response()
response.MergeFromString(ret_bytes.value)
res = res_class()
res.MergeFromString(response.data)
LIB.tfg_pitc_FreeMemoryBuffer(ret_ptr)
return res
def make_command(self, message: Message) -> Command:
"""Make a command instance from the given message."""
return Command(name=type(message).__name__,
data=message.SerializeToString())
def serialize(stream: io.RawIOBase, message: Message):
stream.write(message.SerializeToString())
async def write_pbmsg(stream: asyncio.StreamWriter, pbmsg: PBMessage) -> None:
size = pbmsg.ByteSize()
await write_unsigned_varint(stream, size)
msg_bytes: bytes = pbmsg.SerializeToString()
stream.write(msg_bytes)
def _serialize(message: Message) -> bytes:
out = BytesIO()
with gzip.GzipFile(fileobj=out, mode="w") as f:
f.write(message.SerializeToString())
return out.getvalue()
def frame_encode(msg: Message) -> bytes:
pb_data = msg.SerializeToString()
return cobs.encode(pb_data)
def write_tfrecord_file(file_name: Text, proto: Message) -> None:
"""Writes a serialized tfrecord to file."""
tf.gfile.MakeDirs(os.path.dirname(file_name))
with tf.python_io.TFRecordWriter(file_name) as writer:
writer.write(proto.SerializeToString())
async def write_pbmsg(stream: anyio.abc.SocketStream, pbmsg: PBMessage) -> None:
size = pbmsg.ByteSize()
await write_unsigned_varint(stream, size)
msg_bytes: bytes = pbmsg.SerializeToString()
await stream.send_all(msg_bytes)
def _base64_encode(msg: message.Message) -> str:
return base64.b64encode(msg.SerializeToString()).decode('utf-8')
def write(self, message: Message):
_bytes = message.SerializeToString()
logger.info("sending: %r", _bytes)
self._raw_write(_bytes)
