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 main(argv):
del argv # Unused.
config.CONFIG.AddContext("FleetspeakFrontend Context")
server_startup.Init()
server_startup.DropPrivileges()
fleetspeak_connector.Init()
fsd = fleetspeak_frontend_server.GRRFSServer()
fleetspeak_connector.CONN.Listen(fsd.Process)
logging.info("Serving through Fleetspeak ...")
try:
while True:
time.sleep(600)
except KeyboardInterrupt:
print("Caught keyboard interrupt, stopping")
def testWriteLastPingForNewClients(self):
fs_server = fleetspeak_frontend_server.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)
