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)
