def testUpdateFromFleetspeak(self): client_id_1 = client_plugin.ApiClientId("C." + "1" * 16) client_id_2 = client_plugin.ApiClientId("C." + "2" * 16) client_id_3 = client_plugin.ApiClientId("C." + "3" * 16) clients = [ client_plugin.ApiClient(client_id=client_id_1, fleetspeak_enabled=True), client_plugin.ApiClient(client_id=client_id_2, fleetspeak_enabled=True), client_plugin.ApiClient(client_id=client_id_3, fleetspeak_enabled=False), ] conn = mock.MagicMock() conn.outgoing.ListClients.return_value = admin_pb2.ListClientsResponse( clients=[ admin_pb2.Client( client_id=fleetspeak_utils.GRRIDToFleetspeakID( client_id_1), last_contact_time=timestamp_pb2.Timestamp(seconds=100000, nanos=50000000), last_clock=timestamp_pb2.Timestamp(seconds=100000, nanos=60000000), ), admin_pb2.Client( client_id=fleetspeak_utils.GRRIDToFleetspeakID( client_id_2), last_contact_time=timestamp_pb2.Timestamp(seconds=200000, nanos=50000000), last_clock=timestamp_pb2.Timestamp(seconds=200000, nanos=60000000), ) ]) with mock.patch.object(fleetspeak_connector, "CONN", conn): client_plugin._UpdateFromFleetspeak(clients) self.assertEqual(clients, [ client_plugin.ApiClient( client_id=client_id_1, fleetspeak_enabled=True, last_seen_at=rdfvalue.RDFDatetime(100000050), last_clock=rdfvalue.RDFDatetime(100000060)), client_plugin.ApiClient( client_id=client_id_2, fleetspeak_enabled=True, last_seen_at=rdfvalue.RDFDatetime(200000050), last_clock=rdfvalue.RDFDatetime(200000060)), client_plugin.ApiClient(client_id=client_id_3, fleetspeak_enabled=False), ])
def testUpdateClientsFromFleetspeak(self): client_id_1 = client_plugin.ApiClientId("C." + "1" * 16) client_id_2 = client_plugin.ApiClientId("C." + "2" * 16) client_id_3 = client_plugin.ApiClientId("C." + "3" * 16) clients = [ client_plugin.ApiClient(client_id=client_id_1, fleetspeak_enabled=True), client_plugin.ApiClient(client_id=client_id_2, fleetspeak_enabled=True), client_plugin.ApiClient(client_id=client_id_3, fleetspeak_enabled=False), ] conn = mock.MagicMock() conn.outgoing.ListClients.return_value = admin_pb2.ListClientsResponse( clients=[ admin_pb2.Client( client_id=fleetspeak_utils.GRRIDToFleetspeakID( client_id_1), last_contact_time=TSProtoFromString( "2018-01-01T00:00:01Z"), last_clock=TSProtoFromString("2018-01-01T00:00:02Z")), admin_pb2.Client( client_id=fleetspeak_utils.GRRIDToFleetspeakID( client_id_2), last_contact_time=TSProtoFromString( "2018-01-02T00:00:01Z"), last_clock=TSProtoFromString("2018-01-02T00:00:02Z")) ]) with mock.patch.object(fleetspeak_connector, "CONN", conn): client_plugin.UpdateClientsFromFleetspeak(clients) self.assertEqual(clients, [ client_plugin.ApiClient( client_id=client_id_1, fleetspeak_enabled=True, last_seen_at=rdfvalue.RDFDatetime.FromHumanReadable( "2018-01-01T00:00:01Z"), last_clock=rdfvalue.RDFDatetime.FromHumanReadable( "2018-01-01T00:00:02Z")), client_plugin.ApiClient( client_id=client_id_2, fleetspeak_enabled=True, last_seen_at=rdfvalue.RDFDatetime.FromHumanReadable( "2018-01-02T00:00:01Z"), last_clock=rdfvalue.RDFDatetime.FromHumanReadable( "2018-01-02T00:00:02Z")), client_plugin.ApiClient(client_id=client_id_3, fleetspeak_enabled=False), ])
def _GetAddrFromFleetspeak(client_id): res = fleetspeak_connector.CONN.outgoing.ListClients( admin_pb2.ListClientsRequest( client_ids=[fleetspeak_utils.GRRIDToFleetspeakID(client_id)])) if not res.clients or not res.clients[0].last_contact_address: return "", None # last_contact_address typically includes a port parsed = urlparse.urlparse("//{}".format( res.clients[0].last_contact_address)) ip_str = parsed.hostname return ip_str, ipaddr.IPAddress(ip_str)
def testGetAddrFromFleetspeakMissing(self): client_id = client_plugin.ApiClientId("C." + "1" * 16) conn = mock.MagicMock() conn.outgoing.ListClients.return_value = admin_pb2.ListClientsResponse( clients=[ admin_pb2.Client( client_id=fleetspeak_utils.GRRIDToFleetspeakID(client_id), last_contact_time=TSProtoFromString( "2018-01-01T00:00:01Z"), last_clock=TSProtoFromString("2018-01-01T00:00:02Z")) ]) with mock.patch.object(fleetspeak_connector, "CONN", conn): ip_str, ipaddr_obj = client_plugin._GetAddrFromFleetspeak( client_id) self.assertEqual(ip_str, "") self.assertIsNone(ipaddr_obj)
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 UpdateClientsFromFleetspeak(clients): """Updates ApiClient records to include info from Fleetspeak.""" if not fleetspeak_connector.CONN or not fleetspeak_connector.CONN.outgoing: # FS not configured, or an outgoing connection is otherwise unavailable. return id_map = {} for client in clients: if client.fleetspeak_enabled: id_map[fleetspeak_utils.GRRIDToFleetspeakID( client.client_id)] = client if not id_map: return res = fleetspeak_connector.CONN.outgoing.ListClients( admin_pb2.ListClientsRequest(client_ids=id_map.keys())) for read in res.clients: api_client = id_map[read.client_id] api_client.last_seen_at = fleetspeak_utils.TSToRDFDatetime( read.last_contact_time) api_client.last_clock = fleetspeak_utils.TSToRDFDatetime( read.last_clock)
def testGetAddrFromFleetspeakIpV6(self): client_id = client_plugin.ApiClientId("C." + "1" * 16) conn = mock.MagicMock() conn.outgoing.ListClients.return_value = admin_pb2.ListClientsResponse( clients=[ admin_pb2.Client( client_id=fleetspeak_utils.GRRIDToFleetspeakID(client_id), last_contact_address= "[2001:0db8:85a3::8a2e:0370:7334]:50000", last_contact_time=TSProtoFromString( "2018-01-01T00:00:01Z"), last_clock=TSProtoFromString("2018-01-01T00:00:02Z")) ]) with mock.patch.object(fleetspeak_connector, "CONN", conn): ip_str, ipaddr_obj = client_plugin._GetAddrFromFleetspeak( client_id) self.assertEqual(ip_str, "2001:0db8:85a3::8a2e:0370:7334") self.assertEqual( ipaddr_obj, ipaddr.IPAddress("2001:0db8:85a3:0000:0000:8a2e:0370:7334"))
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 xrange(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 xrange(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)