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, ipaddress.ip_address(ip_str)
def testReceiveMessageList_Relational(self):
if not data_store.RelationalDBFlowsEnabled():
self.skipTest("Rel-db-only test.")
fs_server = fs_frontend_tool.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())
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))
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 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 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 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,
ipaddress.ip_address("2001:0db8:85a3:0000:0000:8a2e:0370:7334"))
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=list(iterkeys(id_map))))
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 testWriteLastPingForNewClients(self):
if not data_store.RelationalDBEnabled():
self.skipTest("Rel-db-only test.")
fs_server = fs_frontend_tool.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)
def Run(self):
if not fleetspeak_connector.CONN or not fleetspeak_connector.CONN.outgoing:
# Nothing to do if Fleetspeak is not enabled.
self.Log("Fleetspeak has not been initialized. Will do nothing.")
return
if not data_store.RelationalDBWriteEnabled():
raise NotImplementedError(
"Cronjob does not support the legacy datastore.")
age_threshold = config.CONFIG["Server.fleetspeak_last_ping_threshold"]
max_last_ping = rdfvalue.RDFDatetime.Now() - age_threshold
last_pings = data_store.REL_DB.ReadClientLastPings(
max_last_ping=max_last_ping, fleetspeak_enabled=True)
num_clients_updated = 0
batch_size = config.CONFIG["Server.fleetspeak_list_clients_batch_size"]
for client_ids in collection.Batch(iterkeys(last_pings), batch_size):
fs_ids = [
fleetspeak_utils.GRRIDToFleetspeakID(i) for i in client_ids
]
request_start = rdfvalue.RDFDatetime.Now()
fs_result = fleetspeak_connector.CONN.outgoing.ListClients(
admin_pb2.ListClientsRequest(client_ids=fs_ids))
latency = rdfvalue.RDFDatetime.Now() - request_start
logging.info("Fleetspeak ListClients() took %s.", latency)
stats_collector_instance.Get().RecordEvent(
"fleetspeak_last_ping_latency_millis", latency.milliseconds)
for fs_client in fs_result.clients:
grr_id = fleetspeak_utils.FleetspeakIDToGRRID(
fs_client.client_id)
new_last_ping = fleetspeak_utils.TSToRDFDatetime(
fs_client.last_contact_time)
if last_pings[grr_id] is None or last_pings[
grr_id] < new_last_ping:
data_store.REL_DB.WriteClientMetadata(
grr_id, last_ping=new_last_ping)
num_clients_updated += 1
self.Log("Updated timestamps for %d clients.", num_clients_updated)
def testReceiveMessagesFleetspeak(self):
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 = b"\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=FS_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)
def testCronJob(self, fs_conn_mock):
if not data_store.RelationalDBReadEnabled():
self.skipTest("Test is only for the relational DB. Skipping...")
client_id1 = "C.0000000000000001"
client_id2 = "C.0000000000000002"
client_id3 = "C.0000000000000003"
client_id4 = "C.0000000000000004"
client_id5 = "C.0000000000000005"
client_id6 = "C.0000000000000006"
client_id7 = "C.0000000000000007"
data_store.REL_DB.WriteClientMetadata(client_id1,
fleetspeak_enabled=False)
data_store.REL_DB.WriteClientMetadata(client_id2,
fleetspeak_enabled=True)
data_store.REL_DB.WriteClientMetadata(
client_id3,
last_ping=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(3),
fleetspeak_enabled=True)
data_store.REL_DB.WriteClientMetadata(
client_id4,
last_ping=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(41),
fleetspeak_enabled=True)
data_store.REL_DB.WriteClientMetadata(
client_id5,
last_ping=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(5),
fleetspeak_enabled=True)
data_store.REL_DB.WriteClientMetadata(
client_id6,
last_ping=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(61),
fleetspeak_enabled=True)
data_store.REL_DB.WriteClientMetadata(
client_id7,
last_ping=rdfvalue.RDFDatetime.FromSecondsSinceEpoch(68),
fleetspeak_enabled=True)
fs_enabled_ids = [
client_id2, client_id3, client_id4, client_id5, client_id6,
client_id7
]
fs_clients = {}
for i, client_id in enumerate(fs_enabled_ids):
client_number = i + 2
fs_client_id = fleetspeak_utils.GRRIDToFleetspeakID(client_id)
fs_client = admin_pb2.Client(client_id=fs_client_id)
fs_client.last_contact_time.FromSeconds(client_number * 10)
fs_clients[fs_client_id] = fs_client
def FakeListClients(list_request):
clients = []
for fs_client_id in list_request.client_ids:
clients.append(fs_clients[fs_client_id])
return admin_pb2.ListClientsResponse(clients=clients)
fs_conn_mock.outgoing.ListClients = FakeListClients
cron_run = rdf_cronjobs.CronJobRun()
job_data = rdf_cronjobs.CronJob()
cron = system.UpdateFSLastPingTimestamps(cron_run, job_data)
with test_lib.FakeTime(
rdfvalue.RDFDatetime.FromSecondsSinceEpoch(100)):
with test_lib.ConfigOverrider({
"Server.fleetspeak_last_ping_threshold":
"35s",
"Server.fleetspeak_list_clients_batch_size":
2,
}):
cron.Run()
actual_timestamps = data_store.REL_DB.ReadClientLastPings()
expected_timestamps = {
# Skipped because not a Fleetspeak client.
client_id1: None,
client_id2: rdfvalue.RDFDatetime.FromSecondsSinceEpoch(20),
client_id3: rdfvalue.RDFDatetime.FromSecondsSinceEpoch(30),
# Skipped because FS timestamp is old.
client_id4: rdfvalue.RDFDatetime.FromSecondsSinceEpoch(41),
client_id5: rdfvalue.RDFDatetime.FromSecondsSinceEpoch(50),
# Skipped because FS timestamp is old.
client_id6: rdfvalue.RDFDatetime.FromSecondsSinceEpoch(61),
# Skipped because existing GRR timestamp is too recent.
client_id7: rdfvalue.RDFDatetime.FromSecondsSinceEpoch(68),
}
self.assertEqual(actual_timestamps, expected_timestamps)
self.assertMultiLineEqual(cron._log_messages.popleft(),
"Updated timestamps for 3 clients.")