def testGetLabelsFromFleetspeak_UnknownClient(self):
conn = mock.MagicMock()
conn.outgoing.ListClients.return_value = admin_pb2.ListClientsResponse(
)
with fleetspeak_test_lib.ConnectionOverrider(conn):
self.assertEmpty(
fleetspeak_utils.GetLabelsFromFleetspeak(_TEST_CLIENT_ID))
def testGetLabelFromFleetspeakKnown(self):
client_id = "C.0000000000000001"
conn = MockConnReturningClient(
admin_pb2.Client(
client_id=fleetspeak_utils.GRRIDToFleetspeakID(client_id),
labels=[common_pb2.Label(service_name="client",
label="division1")]))
with test_lib.ConfigOverrider({
"Server.fleetspeak_label_map": ["division1:fleetspeak-division1"],
}):
with fleetspeak_test_lib.ConnectionOverrider(conn):
self.assertEqual("fleetspeak-division1",
fleetspeak_utils.GetLabelFromFleetspeak(client_id))
def testGetLabelsFromFleetspeak_NoPrefix(self):
conn = _MockConnReturningClient(_TEST_CLIENT_ID, [
("client", "foo-1"),
("client", "bar-2"),
("service-1", "foo-3"),
("service-1", "foo-4"),
("client", "foo-5"),
])
with test_lib.ConfigOverrider({
"Server.fleetspeak_label_map": ["foo-5:baz-5"],
}):
with fleetspeak_test_lib.ConnectionOverrider(conn):
self.assertListEqual(
fleetspeak_utils.GetLabelsFromFleetspeak(_TEST_CLIENT_ID),
["foo-1", "bar-2", "baz-5"])
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)
fake_conn = _FakeGRPCServiceClient(FS_SERVICE_NAME,
send_callback=SendCallback)
with fleetspeak_test_lib.ConnectionOverrider(fake_conn):
with mock.patch.object(fake_conn.outgoing,
"InsertMessage",
wraps=fake_conn.outgoing.InsertMessage):
flow_urn = flow.StartAFF4Flow(
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 setUp(self):
super(FleetspeakGRRFEServerTest, self).setUp()
fake_conn = _FakeGRPCServiceClient(FS_SERVICE_NAME)
self._conn_overrider = fleetspeak_test_lib.ConnectionOverrider(
fake_conn)
self._conn_overrider.Start()
def testGetLabelsFromFleetspeak_NoLabels(self):
conn = _MockConnReturningClient(_TEST_CLIENT_ID,
[("service-1", "foo-3")])
with fleetspeak_test_lib.ConnectionOverrider(conn):
self.assertEmpty(
fleetspeak_utils.GetLabelsFromFleetspeak(_TEST_CLIENT_ID))
