def testRaisesWhenInitializedFromInvalidValues(self):
with self.assertRaises(ValueError):
client_plugin.ApiClientId("blah")
with self.assertRaises(ValueError):
client_plugin.ApiClientId("C.0")
with self.assertRaises(ValueError):
client_plugin.ApiClientId("C." + "0" * 15)
with self.assertRaises(ValueError):
client_plugin.ApiClientId("C." + "1" * 16 + "/foo")
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 setUp(self):
super(CheckClientLabelsTest, self).setUp()
self.client_urn = self.SetupClient(0)
self.client_id = api_client.ApiClientId(self.client_urn.Basename())
self.labels_whitelist = ["foo"]
self.labels_owners_whitelist = ["GRR"]
def testInitializesClientIdForClientBasedFlows(self):
client_id = self.SetupClient(0)
flow_id = flow.StartFlow(client_id=client_id,
flow_cls=processes.ListProcesses)
flow_obj = data_store.REL_DB.ReadFlowObject(client_id, flow_id)
flow_api_obj = flow_plugin.ApiFlow().InitFromFlowObject(flow_obj)
self.assertEqual(flow_api_obj.client_id,
client_plugin.ApiClientId(client_id))
def testResolvesSimpleFlowURN(self):
flow_urn = flow.StartAFF4Flow(
flow_name=flow_test_lib.FlowWithOneNestedFlow.__name__,
client_id=self.client_urn,
token=self.token)
flow_id = flow_plugin.ApiFlowId(flow_urn.Basename())
self.assertEqual(
flow_id.ResolveClientFlowURN(client_plugin.ApiClientId(
self.client_urn),
token=self.token), flow_urn)
def testLeavesClientIdEmptyForNonClientBasedFlows(self):
client_id = self.SetupClient(0)
flow_urn = flow.StartFlow(client_id=client_id,
flow_name=processes.ListProcesses.__name__,
token=self.token)
flow_obj = aff4.FACTORY.Open(flow_urn, token=self.token)
flow_api_obj = flow_plugin.ApiFlow().InitFromAff4Object(
flow_obj, flow_id=flow_urn.Basename())
self.assertEquals(flow_api_obj.client_id,
client_plugin.ApiClientId(client_id))
def testResolvesHuntFlowURN(self):
self._StartHunt()
client_flows_urns = list(
aff4.FACTORY.ListChildren(self.client_urn.Add("flows")))
self.assertEqual(len(client_flows_urns), 1)
flow_id = flow_plugin.ApiFlowId(client_flows_urns[0].Basename())
self.assertEqual(
flow_id.ResolveClientFlowURN(
client_plugin.ApiClientId(self.client_urn), token=self.token),
client_flows_urns[0])
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 testResolvesNestedFlowURN(self):
flow_urn = flow.StartAFF4Flow(
flow_name=flow_test_lib.FlowWithOneNestedFlow.__name__,
client_id=self.client_urn,
token=self.token)
children = list(
aff4.FACTORY.MultiOpen(list(aff4.FACTORY.ListChildren(flow_urn)),
aff4_type=flow.GRRFlow,
token=self.token))
self.assertLen(children, 1)
flow_id = flow_plugin.ApiFlowId(flow_urn.Basename() + "/" +
children[0].urn.Basename())
self.assertEqual(
flow_id.ResolveClientFlowURN(client_plugin.ApiClientId(
self.client_urn),
token=self.token), children[0].urn)
def testResolvesNestedHuntFlowURN(self):
self._StartHunt()
client_flows_urns = list(
aff4.FACTORY.ListChildren(self.client_urn.Add("flows")))
self.assertEqual(len(client_flows_urns), 1)
flow_fd = aff4.FACTORY.Open(client_flows_urns[0], token=self.token)
nested_flows_urns = list(flow_fd.ListChildren())
nested_flows = list(
aff4.FACTORY.MultiOpen(
nested_flows_urns, aff4_type=flow.GRRFlow, token=self.token))
self.assertEqual(len(nested_flows), 1)
flow_id = flow_plugin.ApiFlowId(client_flows_urns[0].Basename() + "/" +
nested_flows[0].urn.Basename())
self.assertEqual(
flow_id.ResolveClientFlowURN(
client_plugin.ApiClientId(self.client_urn), token=self.token),
nested_flows[0].urn)
def testConvertsToString(self):
client_id = client_plugin.ApiClientId("C." + "1" * 16)
client_id_str = client_id.ToString()
self.assertEqual(client_id_str, "C." + "1" * 16)
def testRaisesWhenToStringCalledOnUninitializedValue(self):
client_id = client_plugin.ApiClientId()
with self.assertRaises(ValueError):
client_id.ToString()
def GenerateSample(self, number=0):
return client_plugin.ApiClientId("C.%016d" % number)
def testConvertsToClientURN(self):
client_id = client_plugin.ApiClientId("C." + "1" * 16)
client_urn = client_id.ToClientURN()
self.assertEqual(client_urn.Basename(), client_id)
self.assertEqual(client_urn, "aff4:/C." + "1" * 16)
