def testProcessListingOnlyFleetspeak(self):
"""Test that the ListProcesses flow works with Fleetspeak."""
client_mock = action_mocks.ListProcessesMock([
rdf_client.Process(pid=2,
ppid=1,
cmdline=["cmd.exe"],
exe=r"c:\windows\cmd.exe",
ctime=1333718907167083)
])
client_mock.mock_task_queue = []
def SendCallback(fs_msg):
pb_msg = jobs_pb2.GrrMessage()
fs_msg.data.Unpack(pb_msg)
msg = rdf_flows.GrrMessage.FromSerializedString(
pb_msg.SerializeToString())
client_mock.mock_task_queue.append(msg)
service_name = "GRR"
fake_service_client = _FakeGRPCServiceClient(
service_name, send_callback=SendCallback)
fleetspeak_connector.Reset()
fleetspeak_connector.Init(service_client=fake_service_client)
with mock.patch.object(
fake_service_client.outgoing,
"InsertMessage",
wraps=fake_service_client.outgoing.InsertMessage):
flow_urn = flow.StartFlow(
client_id=self.client_id,
flow_name=flow_processes.ListProcesses.__name__,
token=self.token)
session_id = flow_test_lib.TestFlowHelper(flow_urn,
client_mock,
client_id=self.client_id,
token=self.token)
fleetspeak_connector.CONN.outgoing.InsertMessage.assert_called()
# Check the output collection
processes = flow.GRRFlow.ResultCollectionForFID(session_id)
self.assertEqual(len(processes), 1)
process, = processes
self.assertEqual(process.ctime, 1333718907167083)
self.assertEqual(process.cmdline, ["cmd.exe"])
def testReceiveMessagesFleetspeak(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_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 = "\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()))
fs_messages = [
fs_common_pb2.Message(message_type="GrrMessage",
source=fs_common_pb2.Address(
client_id=fs_client_id,
service_name=service_name))
for _ in range(num_msgs)
]
for fs_message, message in itertools.izip(fs_messages, messages):
fs_message.data.Pack(message.AsPrimitiveProto())
for msg in fs_messages:
fsd.Process(msg, 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)