def testReceiveMessageListFleetspeak(self):
fsd = fs_frontend_tool.GRRFSServer()
grr_client_nr = 0xab
grr_client_id_urn = self.SetupClient(grr_client_nr)
flow_obj = self.FlowSetup(flow_test_lib.FlowOrderTest.__name__,
grr_client_id_urn)
num_msgs = 9
session_id = flow_obj.session_id
messages = [
rdf_flows.GrrMessage(request_id=1,
response_id=i,
session_id=session_id,
payload=rdfvalue.RDFInteger(i))
for i in range(1, num_msgs + 1)
]
fs_client_id = b"\x10\x00\x00\x00\x00\x00\x00\xab"
# fs_client_id should be equivalent to grr_client_id_urn
self.assertEqual(
fs_client_id,
fleetspeak_utils.GRRIDToFleetspeakID(grr_client_id_urn.Basename()))
message_list = rdf_flows.PackedMessageList()
communicator.Communicator.EncodeMessageList(
rdf_flows.MessageList(job=messages), message_list)
fs_message = fs_common_pb2.Message(message_type="MessageList",
source=fs_common_pb2.Address(
client_id=fs_client_id,
service_name=FS_SERVICE_NAME))
fs_message.data.Pack(message_list.AsPrimitiveProto())
fsd.Process(fs_message, None)
# Make sure the task is still on the client queue
manager = queue_manager.QueueManager(token=self.token)
tasks_on_client_queue = manager.Query(grr_client_id_urn, 100)
self.assertEqual(len(tasks_on_client_queue), 1)
want_messages = [message.Copy() for message in messages]
for want_message in want_messages:
# This is filled in by the frontend as soon as it gets the message.
want_message.auth_state = (
rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED)
want_message.source = grr_client_id_urn
stored_messages = data_store.DB.ReadResponsesForRequestId(
session_id, 1)
self.assertEqual(len(stored_messages), len(want_messages))
stored_messages.sort(key=lambda m: m.response_id)
# Check that messages were stored correctly
for stored_message, want_message in itertools.izip(
stored_messages, want_messages):
stored_message.timestamp = None
self.assertRDFValuesEqual(stored_message, want_message)
def _SendMessages(self, grr_msgs, background=False):
"""Sends a block of messages through Fleetspeak."""
message_list = rdf_flows.PackedMessageList()
communicator.Communicator.EncodeMessageList(
rdf_flows.MessageList(job=grr_msgs), message_list)
fs_msg = fs_common_pb2.Message(
message_type="MessageList",
destination=fs_common_pb2.Address(service_name="GRR"),
background=background)
fs_msg.data.Pack(message_list.AsPrimitiveProto())
for grr_msg in grr_msgs:
if (grr_msg.session_id is None or grr_msg.request_id is None or
grr_msg.response_id is None):
continue
# Place all ids in a single annotation, instead of having separate
# annotations for the flow-id, request-id and response-id. This reduces
# overall size of the annotations by half (~60 bytes to ~30 bytes).
annotation = fs_msg.annotations.entries.add()
annotation.key = _DATA_IDS_ANNOTATION_KEY
annotation.value = "%s:%d:%d" % (grr_msg.session_id.Basename(),
grr_msg.request_id, grr_msg.response_id)
if fs_msg.annotations.ByteSize() >= _MAX_ANNOTATIONS_BYTES:
break
try:
sent_bytes = self._fs.Send(fs_msg)
except (IOError, struct.error):
logging.critical("Broken local Fleetspeak connection (write end).")
raise
communicator.GRR_CLIENT_SENT_BYTES.Increment(sent_bytes)
def SendGrrMessageThroughFleetspeak(grr_id: str,
grr_msg: rdf_flows.GrrMessage) -> None:
"""Sends the given GrrMessage through FS with retrying.
The send operation is retried if a `grpc.RpcError` occurs.
The maximum number of retries corresponds to the config value
`Server.fleetspeak_send_retry_attempts`.
A retry is delayed by the number of seconds specified in the config value
`Server.fleetspeak_send_retry_sleep_time_secs`.
Args:
grr_id: ID of grr client to send message to.
grr_msg: GRR message to send.
"""
fs_msg = fs_common_pb2.Message(message_type="GrrMessage",
destination=fs_common_pb2.Address(
client_id=GRRIDToFleetspeakID(grr_id),
service_name="GRR"))
fs_msg.data.Pack(grr_msg.AsPrimitiveProto())
if grr_msg.session_id is not None:
annotation = fs_msg.annotations.entries.add()
annotation.key, annotation.value = "flow_id", grr_msg.session_id.Basename(
)
if grr_msg.request_id is not None:
annotation = fs_msg.annotations.entries.add()
annotation.key, annotation.value = "request_id", str(
grr_msg.request_id)
fleetspeak_connector.CONN.outgoing.InsertMessage(
fs_msg,
single_try_timeout=WRITE_SINGLE_TRY_TIMEOUT,
timeout=WRITE_TOTAL_TIMEOUT)
def testInsertMessage(self):
t = FakeStub()
s = client.OutgoingConnection(None, 'test', t)
s.InsertMessage(common_pb2.Message())
self.assertTrue(t.insert_done)
self.assertFalse(t.insert_errors)
self.assertEqual(t.message.source.service_name, 'test')
def _SendMessages(self, grr_msgs, priority=fs_common_pb2.Message.MEDIUM):
"""Sends a block of messages through Fleetspeak."""
message_list = rdf_flows.PackedMessageList()
communicator.Communicator.EncodeMessageList(
rdf_flows.MessageList(job=grr_msgs), message_list)
fs_msg = fs_common_pb2.Message(
message_type="MessageList",
destination=fs_common_pb2.Address(service_name="GRR"),
priority=priority)
fs_msg.data.Pack(message_list.AsPrimitiveProto())
try:
sent_bytes = self._fs.Send(fs_msg)
except (IOError, struct.error) as e:
logging.fatal(
"Broken local Fleetspeak connection (write end): %r",
e,
exc_info=True)
# The fatal call above doesn't terminate the program. The reasons for
# this might include Python threads persistency, or Python logging
# mechanisms' inconsistency.
os._exit(1) # pylint: disable=protected-access
stats.STATS.IncrementCounter("grr_client_sent_bytes", sent_bytes)
def _WriteStartupData(self, version):
startup_msg = common_pb2.Message(
message_type="StartupData",
destination=common_pb2.Address(service_name="system"))
startup_msg.data.Pack(
channel_pb2.StartupData(pid=os.getpid(), version=version))
self._SendImpl(startup_msg)
def SendGrrMessageThroughFleetspeak(grr_id, msg):
"""Sends the given GrrMessage through FS."""
fs_msg = fs_common_pb2.Message(message_type="GrrMessage",
destination=fs_common_pb2.Address(
client_id=GRRIDToFleetspeakID(grr_id),
service_name="GRR"))
fs_msg.data.Pack(msg.AsPrimitiveProto())
fleetspeak_connector.CONN.outgoing.InsertMessage(fs_msg)
def RestartFleetspeakGrrService(grr_id: Text) -> None: """Restarts the GRR service on the given client.""" restart_req = fs_system_pb2.RestartServiceRequest(name="GRR") fs_msg = fs_common_pb2.Message() fs_msg.message_type = "RestartService" fs_msg.destination.client_id = GRRIDToFleetspeakID(grr_id) fs_msg.destination.service_name = "system" fs_msg.data.Pack(restart_req) fleetspeak_connector.CONN.outgoing.InsertMessage(fs_msg)
def KillFleetspeak(grr_id: Text, force: bool) -> None: """Kills Fleespeak on the given client.""" die_req = fs_system_pb2.DieRequest(force=force) fs_msg = fs_common_pb2.Message() fs_msg.message_type = "Die" fs_msg.destination.client_id = GRRIDToFleetspeakID(grr_id) fs_msg.destination.service_name = "system" fs_msg.data.Pack(die_req) fleetspeak_connector.CONN.outgoing.InsertMessage(fs_msg)
def testReceiveMessageList(self):
fs_server = fleetspeak_frontend_server.GRRFSServer()
client_id = "C.1234567890123456"
flow_id = "12345678"
data_store.REL_DB.WriteClientMetadata(client_id,
fleetspeak_enabled=True)
rdf_flow = rdf_flow_objects.Flow(
client_id=client_id,
flow_id=flow_id,
create_time=rdfvalue.RDFDatetime.Now())
data_store.REL_DB.WriteFlowObject(rdf_flow)
flow_request = rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=1)
data_store.REL_DB.WriteFlowRequests([flow_request])
session_id = "%s/%s" % (client_id, flow_id)
fs_client_id = fleetspeak_utils.GRRIDToFleetspeakID(client_id)
grr_messages = []
for i in range(1, 10):
grr_message = rdf_flows.GrrMessage(request_id=1,
response_id=i + 1,
session_id=session_id,
payload=rdfvalue.RDFInteger(i))
grr_messages.append(grr_message)
packed_messages = rdf_flows.PackedMessageList()
communicator.Communicator.EncodeMessageList(
rdf_flows.MessageList(job=grr_messages), packed_messages)
fs_message = fs_common_pb2.Message(message_type="MessageList",
source=fs_common_pb2.Address(
client_id=fs_client_id,
service_name=FS_SERVICE_NAME))
fs_message.data.Pack(packed_messages.AsPrimitiveProto())
fs_message.validation_info.tags["foo"] = "bar"
with test_lib.FakeTime(
rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)):
fs_server.Process(fs_message, None)
# Ensure the last-ping timestamp gets updated.
client_data = data_store.REL_DB.MultiReadClientMetadata([client_id])
self.assertEqual(client_data[client_id].ping,
rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123))
self.assertEqual(
client_data[client_id].last_fleetspeak_validation_info.
ToStringDict(), {"foo": "bar"})
flow_data = data_store.REL_DB.ReadAllFlowRequestsAndResponses(
client_id, flow_id)
self.assertLen(flow_data, 1)
stored_flow_request, flow_responses = flow_data[0]
self.assertEqual(stored_flow_request, flow_request)
self.assertLen(flow_responses, 9)
def Heartbeat(self):
"""Sends a heartbeat to the Fleetspeak client.
If this daemonservice is configured to use heartbeats, clients that don't
call this method often enough are considered faulty and are restarted by
Fleetspeak.
"""
heartbeat_msg = common_pb2.Message(
message_type="Heartbeat",
destination=common_pb2.Address(service_name="system"))
self._SendImpl(heartbeat_msg)
def testReceiveMessages_Relational(self):
if not data_store.RelationalDBEnabled():
self.skipTest("Rel-db-only test.")
fs_server = fs_frontend_tool.GRRFSServer()
client_id = "C.1234567890123456"
flow_id = "12345678"
data_store.REL_DB.WriteClientMetadata(client_id,
fleetspeak_enabled=True)
rdf_flow = rdf_flow_objects.Flow(
client_id=client_id,
flow_id=flow_id,
create_time=rdfvalue.RDFDatetime.Now())
data_store.REL_DB.WriteFlowObject(rdf_flow)
flow_request = rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=1)
data_store.REL_DB.WriteFlowRequests([flow_request])
session_id = "%s/%s" % (client_id, flow_id)
fs_client_id = fleetspeak_utils.GRRIDToFleetspeakID(client_id)
fs_messages = []
for i in range(1, 10):
grr_message = rdf_flows.GrrMessage(request_id=1,
response_id=i + 1,
session_id=session_id,
payload=rdfvalue.RDFInteger(i))
fs_message = fs_common_pb2.Message(
message_type="GrrMessage",
source=fs_common_pb2.Address(client_id=fs_client_id,
service_name=FS_SERVICE_NAME))
fs_message.data.Pack(grr_message.AsPrimitiveProto())
fs_messages.append(fs_message)
with test_lib.FakeTime(
rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)):
for fs_message in fs_messages:
fs_server.Process(fs_message, None)
# Ensure the last-ping timestamp gets updated.
client_data = data_store.REL_DB.MultiReadClientMetadata([client_id])
self.assertEqual(client_data[client_id].ping,
rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123))
flow_data = data_store.REL_DB.ReadAllFlowRequestsAndResponses(
client_id, flow_id)
self.assertLen(flow_data, 1)
stored_flow_request, flow_responses = flow_data[0]
self.assertEqual(stored_flow_request, flow_request)
self.assertLen(flow_responses, 9)
def SendGrrMessageThroughFleetspeak(grr_id, grr_msg):
"""Sends the given GrrMessage through FS."""
fs_msg = fs_common_pb2.Message(
message_type="GrrMessage",
destination=fs_common_pb2.Address(
client_id=GRRIDToFleetspeakID(grr_id), service_name="GRR"))
fs_msg.data.Pack(grr_msg.AsPrimitiveProto())
if grr_msg.session_id is not None:
annotation = fs_msg.annotations.entries.add()
annotation.key, annotation.value = "flow_id", grr_msg.session_id.Basename()
if grr_msg.request_id is not None:
annotation = fs_msg.annotations.entries.add()
annotation.key, annotation.value = "request_id", str(grr_msg.request_id)
fleetspeak_connector.CONN.outgoing.InsertMessage(fs_msg)
def testGetFleetspeakPendingMessages(self):
conn = mock.MagicMock()
proto_response = admin_pb2.GetPendingMessagesResponse(
messages=[common_pb2.Message(message_id=b"foo")])
conn.outgoing.GetPendingMessages.return_value = proto_response
with mock.patch.object(fleetspeak_connector, "CONN", conn):
result = fleetspeak_utils.GetFleetspeakPendingMessages(
"C.1000000000000000", 1, 2, True)
self.assertEqual(result, proto_response)
get_args, _ = conn.outgoing.GetPendingMessages.call_args
get_req = get_args[0]
self.assertSameElements(get_req.client_ids,
[b"\x10\x00\x00\x00\x00\x00\x00\x00"])
self.assertEqual(get_req.offset, 1)
self.assertEqual(get_req.limit, 2)
self.assertEqual(get_req.want_data, True)
def testPingIsRecorded(self):
service_name = "GRR"
fake_service_client = _FakeGRPCServiceClient(service_name)
fleetspeak_connector.Reset()
fleetspeak_connector.Init(service_client=fake_service_client)
fsd = fs_frontend_tool.GRRFSServer()
grr_client_nr = 0xab
grr_client = self.SetupTestClientObject(grr_client_nr)
self.SetupClient(grr_client_nr)
messages = [
rdf_flows.GrrMessage(request_id=1,
response_id=1,
session_id="F:123456",
payload=rdfvalue.RDFInteger(1))
]
fs_client_id = "\x10\x00\x00\x00\x00\x00\x00\xab"
# fs_client_id should be equivalent to grr_client_id_urn
self.assertEqual(
fs_client_id,
fleetspeak_utils.GRRIDToFleetspeakID(grr_client.client_id))
message_list = rdf_flows.PackedMessageList()
communicator.Communicator.EncodeMessageList(
rdf_flows.MessageList(job=messages), message_list)
fs_message = fs_common_pb2.Message(message_type="MessageList",
source=fs_common_pb2.Address(
client_id=fs_client_id,
service_name=service_name))
fs_message.data.Pack(message_list.AsPrimitiveProto())
fake_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(42)
with test_lib.FakeTime(fake_time):
fsd.Process(fs_message, None)
md = data_store.REL_DB.ReadClientMetadata(grr_client.client_id)
self.assertEqual(md.ping, fake_time)
with aff4.FACTORY.Open(grr_client.client_id) as client:
self.assertEqual(client.Get(client.Schema.PING), fake_time)
def Recv(self):
"""Accept a message from Fleetspeak.
Returns:
A tuple (common_pb2.Message, size of the message in bytes).
Raises:
ProtocolError: If we receive unexpected data from Fleetspeak.
"""
size = struct.unpack(_STRUCT_FMT, self._ReadN(_STRUCT_LEN))[0]
if size > MAX_SIZE:
raise ProtocolError("Expected size to be at most %d, got %d" %
(MAX_SIZE, size))
buf = self._ReadN(size)
self._ReadMagic()
res = common_pb2.Message()
res.ParseFromString(buf)
return res, len(buf)
def _SendMessages(self, grr_msgs, background=False):
"""Sends a block of messages through Fleetspeak."""
message_list = rdf_flows.PackedMessageList()
communicator.Communicator.EncodeMessageList(
rdf_flows.MessageList(job=grr_msgs), message_list)
fs_msg = fs_common_pb2.Message(
message_type="MessageList",
destination=fs_common_pb2.Address(service_name="GRR"),
background=background)
fs_msg.data.Pack(message_list.AsPrimitiveProto())
try:
sent_bytes = self._fs.Send(fs_msg)
except (IOError, struct.error) as e:
logging.critical("Broken local Fleetspeak connection (write end).")
raise e
stats.STATS.IncrementCounter("grr_client_sent_bytes", sent_bytes)
def testWriteLastPingForNewClients(self):
if not data_store.RelationalDBEnabled():
self.skipTest("Rel-db-only test.")
fs_server = fs_frontend_tool.GRRFSServer()
client_id = "C.1234567890123456"
flow_id = "12345678"
session_id = "%s/%s" % (client_id, flow_id)
fs_client_id = fleetspeak_utils.GRRIDToFleetspeakID(client_id)
grr_message = rdf_flows.GrrMessage(request_id=1,
response_id=1,
session_id=session_id,
payload=rdfvalue.RDFInteger(1))
fs_message = fs_common_pb2.Message(message_type="GrrMessage",
source=fs_common_pb2.Address(
client_id=fs_client_id,
service_name=FS_SERVICE_NAME))
fs_message.data.Pack(grr_message.AsPrimitiveProto())
fake_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(123)
with mock.patch.object(
events.Events, "PublishEvent",
wraps=events.Events.PublishEvent) as publish_event_fn:
with mock.patch.object(data_store.REL_DB,
"WriteClientMetadata",
wraps=data_store.REL_DB.WriteClientMetadata
) as write_metadata_fn:
with test_lib.FakeTime(fake_time):
fs_server.Process(fs_message, None)
self.assertEqual(write_metadata_fn.call_count, 1)
client_data = data_store.REL_DB.MultiReadClientMetadata(
[client_id])
self.assertEqual(client_data[client_id].ping, fake_time)
# TODO(user): publish_event_fn.assert_any_call(
# "ClientEnrollment", mock.ANY, token=mock.ANY) doesn't work here
# for some reason.
triggered_events = []
for call_args, _ in publish_event_fn.call_args_list:
if call_args:
triggered_events.append(call_args[0])
self.assertIn("ClientEnrollment", triggered_events)