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))