def testCountHuntFlowsIgnoresSubflows(self):
hunt_obj = rdf_hunt_objects.Hunt(description="foo")
self.db.WriteHuntObject(hunt_obj)
hunt_id = hunt_obj.hunt_id
_, flow_id = self._SetupHuntClientAndFlow(
hunt_id=hunt_id, flow_state=rdf_flow_objects.Flow.FlowState.RUNNING)
# Whatever state the subflow is in, it should be ignored.
self._SetupHuntClientAndFlow(
hunt_id=hunt_id,
flow_id=flow.RandomFlowId(),
parent_flow_id=flow_id,
flow_state=rdf_flow_objects.Flow.FlowState.ERROR)
self._SetupHuntClientAndFlow(
hunt_id=hunt_id,
flow_id=flow.RandomFlowId(),
parent_flow_id=flow_id,
flow_state=rdf_flow_objects.Flow.FlowState.FINISHED)
self._SetupHuntClientAndFlow(
hunt_id=hunt_id,
flow_id=flow.RandomFlowId(),
parent_flow_id=flow_id,
flow_state=rdf_flow_objects.Flow.FlowState.RUNNING)
self.assertEqual(self.db.CountHuntFlows(hunt_id), 1)
def testReadHuntFlowsIgnoresSubflows(self):
hunt_obj = rdf_hunt_objects.Hunt(description="foo")
self.db.WriteHuntObject(hunt_obj)
hunt_id = hunt_obj.hunt_id
_, flow_id = self._SetupHuntClientAndFlow(
hunt_id=hunt_id, flow_state=rdf_flow_objects.Flow.FlowState.RUNNING)
# Whatever state the subflow is in, it should be ignored.
self._SetupHuntClientAndFlow(
hunt_id=hunt_id,
flow_id=flow.RandomFlowId(),
parent_flow_id=flow_id,
flow_state=rdf_flow_objects.Flow.FlowState.ERROR)
self._SetupHuntClientAndFlow(
hunt_id=hunt_id,
flow_id=flow.RandomFlowId(),
parent_flow_id=flow_id,
flow_state=rdf_flow_objects.Flow.FlowState.FINISHED)
self._SetupHuntClientAndFlow(
hunt_id=hunt_id,
flow_id=flow.RandomFlowId(),
parent_flow_id=flow_id,
flow_state=rdf_flow_objects.Flow.FlowState.RUNNING)
for state, expceted_results in [
(db.HuntFlowsCondition.COMPLETED_FLOWS_ONLY, 0),
(db.HuntFlowsCondition.SUCCEEDED_FLOWS_ONLY, 0),
(db.HuntFlowsCondition.FLOWS_IN_PROGRESS_ONLY, 1)
]:
results = self.db.ReadHuntFlows(hunt_id, 0, 10, filter_condition=state)
self.assertLen(results, expceted_results)
def testHuntOverviewGetsUpdatedWhenHuntChanges(self):
hunt_urn = self.CreateSampleHunt()
hunt_id = hunt_urn.Basename()
if data_store.RelationalDBEnabled():
client_id = self.SetupClient(0).Basename()
rdf_flow = rdf_flow_objects.Flow(
client_id=client_id,
flow_id=flow.RandomFlowId(),
parent_hunt_id=hunt_id,
create_time=rdfvalue.RDFDatetime.Now())
rdf_flow.cpu_time_used.user_cpu_time = 5000
rdf_flow.network_bytes_sent = 1000000
data_store.REL_DB.WriteFlowObject(rdf_flow)
else:
with aff4.FACTORY.Open(hunt_urn, mode="rw") as hunt_obj:
hunt_stats = hunt_obj.context.usage_stats
hunt_stats.user_cpu_stats.sum = 5000
hunt_stats.network_bytes_sent_stats.sum = 1000000
self.Open("/")
# Ensure auto-refresh updates happen every second.
self.GetJavaScriptValue(
"grrUi.hunt.huntOverviewDirective.setAutoRefreshInterval(1000);")
self.Click("css=a[grrtarget=hunts]")
self.Click("css=td:contains('%s')" % hunt_id)
self.WaitUntil(self.IsTextPresent, "1h 23m 20s")
self.WaitUntil(self.IsTextPresent, "976.6KiB")
if data_store.RelationalDBEnabled():
client_id = self.SetupClient(1).Basename()
rdf_flow = rdf_flow_objects.Flow(
client_id=client_id,
flow_id=flow.RandomFlowId(),
parent_hunt_id=hunt_id,
create_time=rdfvalue.RDFDatetime.Now())
rdf_flow.cpu_time_used.user_cpu_time = 1000
rdf_flow.network_bytes_sent = 10000000
data_store.REL_DB.WriteFlowObject(rdf_flow)
else:
with aff4.FACTORY.Open(hunt_urn, mode="rw") as hunt_obj:
hunt_stats = hunt_obj.context.usage_stats
hunt_stats.user_cpu_stats.sum = 6000
hunt_stats.network_bytes_sent_stats.sum = 11000000
self.WaitUntil(self.IsTextPresent, "1h 40m")
self.WaitUntil(self.IsTextPresent, "10.5MiB")
def testHuntOverviewShowsStats(self):
"""Test the detailed client view works."""
hunt_urn = self.CreateSampleHunt()
hunt_id = hunt_urn.Basename()
if data_store.RelationalDBEnabled():
client_id = self.SetupClient(0).Basename()
rdf_flow = rdf_flow_objects.Flow(
client_id=client_id,
flow_id=flow.RandomFlowId(),
parent_hunt_id=hunt_id,
create_time=rdfvalue.RDFDatetime.Now())
rdf_flow.cpu_time_used.user_cpu_time = 5000
rdf_flow.network_bytes_sent = 1000000
data_store.REL_DB.WriteFlowObject(rdf_flow)
else:
with aff4.FACTORY.Open(hunt_urn, mode="rw") as hunt_obj:
hunt_stats = hunt_obj.context.usage_stats
hunt_stats.user_cpu_stats.sum = 5000
hunt_stats.network_bytes_sent_stats.sum = 1000000
# Open up and click on View Hunts then the first Hunt.
self.Open("/")
self.WaitUntil(self.IsElementPresent, "client_query")
self.Click("css=a[grrtarget=hunts]")
self.WaitUntil(self.IsTextPresent, hunt_id)
self.Click("css=td:contains('%s')" % hunt_id)
# Click the Overview Tab and check that the stats are present.
self.Click("css=li[heading=Overview]")
self.WaitUntil(self.IsTextPresent, "1h 23m 20s")
self.WaitUntil(self.IsTextPresent, "976.6KiB")
def testHuntOverviewShowsStats(self):
"""Test the detailed client view works."""
hunt_id = self.CreateSampleHunt()
client_id = self.SetupClient(0)
rdf_flow = rdf_flow_objects.Flow(
client_id=client_id,
flow_id=flow.RandomFlowId(),
parent_hunt_id=hunt_id,
create_time=rdfvalue.RDFDatetime.Now())
rdf_flow.cpu_time_used.user_cpu_time = 5000
rdf_flow.network_bytes_sent = 1000000
data_store.REL_DB.WriteFlowObject(rdf_flow)
# Open up and click on View Hunts then the first Hunt.
self.Open("/")
self.WaitUntil(self.IsElementPresent, "client_query")
self.Click("css=a[grrtarget=hunts]")
self.WaitUntil(self.IsTextPresent, hunt_id)
self.Click("css=td:contains('%s')" % hunt_id)
# Click the Overview Tab and check that the stats are present.
self.Click("css=li[heading=Overview]")
self.WaitUntil(self.IsTextPresent, "1h 23m 20s")
self.WaitUntil(self.IsTextPresent, "976.6KiB")
def testStartupTriggersInterrogateWhenPreviousInterrogateIsDone(self):
with test_lib.ConfigOverrider({"Source.version_numeric": 3000}):
client_id = self.SetupClient(0)
self._RunSendStartupInfo(client_id)
data_store.REL_DB.WriteFlowObject(
rdf_flow_objects.Flow(
flow_id=flow.RandomFlowId(),
client_id=client_id,
flow_class_name=discovery.Interrogate.__name__,
flow_state=rdf_flow_objects.Flow.FlowState.FINISHED))
flows = data_store.REL_DB.ReadAllFlowObjects(client_id,
include_child_flows=False)
orig_count = len([
f for f in flows
if f.flow_class_name == discovery.Interrogate.__name__
])
with mock.patch.multiple(discovery.Interrogate,
Start=mock.DEFAULT,
End=mock.DEFAULT):
with test_lib.ConfigOverrider({"Source.version_numeric": 3001}):
self._RunSendStartupInfo(client_id)
flows = data_store.REL_DB.ReadAllFlowObjects(client_id,
include_child_flows=False)
interrogates = [
f for f in flows
if f.flow_class_name == discovery.Interrogate.__name__
]
self.assertLen(interrogates, orig_count + 1)
def testStartupDoesNotTriggerInterrogateIfRecentInterrogateIsRunning(self):
with test_lib.ConfigOverrider({"Source.version_numeric": 3000}):
client_id = self.SetupClient(0)
self._RunSendStartupInfo(client_id)
data_store.REL_DB.WriteFlowObject(
rdf_flow_objects.Flow(
flow_id=flow.RandomFlowId(),
client_id=client_id,
flow_class_name=discovery.Interrogate.__name__,
flow_state=rdf_flow_objects.Flow.FlowState.RUNNING))
flows = data_store.REL_DB.ReadAllFlowObjects(
client_id, include_child_flows=False)
orig_count = len([
f for f in flows
if f.flow_class_name == discovery.Interrogate.__name__
])
self._RunSendStartupInfo(client_id)
flows = data_store.REL_DB.ReadAllFlowObjects(
client_id, include_child_flows=False)
same_ver_count = len([
f for f in flows
if f.flow_class_name == discovery.Interrogate.__name__
])
self.assertEqual(same_ver_count, orig_count)
def testFlowIdGeneration(self):
self.assertLen(flow.RandomFlowId(), 16)
with mock.patch.object(flow.random,
"Id64",
return_value=0xF0F1F2F3F4F5F6F7):
self.assertEqual(flow.RandomFlowId(), "F0F1F2F3F4F5F6F7")
with mock.patch.object(flow.random, "Id64", return_value=0):
self.assertEqual(flow.RandomFlowId(), "0000000000000000")
with mock.patch.object(flow.random, "Id64", return_value=1):
self.assertEqual(flow.RandomFlowId(), "0000000000000001")
with mock.patch.object(flow.random,
"Id64",
return_value=0x0000000100000000):
self.assertEqual(flow.RandomFlowId(), "0000000100000000")
def testHuntOverviewGetsUpdatedWhenHuntChanges(self):
hunt_id = self.CreateSampleHunt()
client_id = self.SetupClient(0)
rdf_flow = rdf_flow_objects.Flow(
client_id=client_id,
flow_id=flow.RandomFlowId(),
parent_hunt_id=hunt_id,
create_time=rdfvalue.RDFDatetime.Now())
rdf_flow.cpu_time_used.user_cpu_time = 5000
rdf_flow.network_bytes_sent = 1000000
data_store.REL_DB.WriteFlowObject(rdf_flow)
self.Open("/")
# Ensure auto-refresh updates happen every second.
self.GetJavaScriptValue(
"grrUi.hunt.huntOverviewDirective.setAutoRefreshInterval(1000);")
self.Click("css=a[grrtarget=hunts]")
self.Click("css=td:contains('%s')" % hunt_id)
self.WaitUntil(self.IsTextPresent, "1h 23m 20s")
self.WaitUntil(self.IsTextPresent, "976.6KiB")
client_id = self.SetupClient(1)
rdf_flow = rdf_flow_objects.Flow(
client_id=client_id,
flow_id=flow.RandomFlowId(),
parent_hunt_id=hunt_id,
create_time=rdfvalue.RDFDatetime.Now())
rdf_flow.cpu_time_used.user_cpu_time = 1000
rdf_flow.network_bytes_sent = 10000000
data_store.REL_DB.WriteFlowObject(rdf_flow)
self.WaitUntil(self.IsTextPresent, "1h 40m")
self.WaitUntil(self.IsTextPresent, "10.5MiB")
def _SetupHuntClientAndFlow(self,
client_id=None,
hunt_id=None,
**additional_flow_args):
client_id = self.InitializeClient(client_id=client_id)
flow_id = flow.RandomFlowId()
self.db.WriteClientMetadata(client_id, fleetspeak_enabled=False)
rdf_flow = rdf_flow_objects.Flow(
client_id=client_id,
flow_id=flow_id,
parent_hunt_id=hunt_id,
create_time=rdfvalue.RDFDatetime.Now(),
**additional_flow_args)
self.db.WriteFlowObject(rdf_flow)
return client_id, flow_id
def testDrainTaskSchedulerQueue(self):
client_id = u"C.1234567890123456"
flow_id = flow.RandomFlowId()
data_store.REL_DB.WriteClientMetadata(client_id,
fleetspeak_enabled=False)
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)
action_requests = []
for i in range(3):
data_store.REL_DB.WriteFlowRequests([
rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=i)
])
action_requests.append(
rdf_flows.ClientActionRequest(client_id=client_id,
flow_id=flow_id,
request_id=i,
action_identifier="WmiQuery"))
data_store.REL_DB.WriteClientActionRequests(action_requests)
server = TestServer()
res = server.DrainTaskSchedulerQueueForClient(
rdfvalue.RDFURN(client_id))
msgs = [
rdf_flow_objects.GRRMessageFromClientActionRequest(r)
for r in action_requests
]
for r in res:
r.task_id = 0
for m in msgs:
m.task_id = 0
self.assertItemsEqual(res, msgs)
def _WriteNestedAndNonNestedLogEntries(self, hunt_obj):
client_id, flow_id = self._SetupHuntClientAndFlow(hunt_id=hunt_obj.hunt_id)
self.db.WriteFlowLogEntries([
rdf_flow_objects.FlowLogEntry(
client_id=client_id,
# Top-level hunt-induced flows should have hunt's ids.
flow_id=flow_id,
hunt_id=hunt_obj.hunt_id,
message="blah")
])
for i in range(10):
_, nested_flow_id = self._SetupHuntClientAndFlow(
client_id=client_id,
parent_flow_id=flow_id,
hunt_id=hunt_obj.hunt_id,
flow_id=flow.RandomFlowId())
self.db.WriteFlowLogEntries([
rdf_flow_objects.FlowLogEntry(
client_id=client_id,
flow_id=nested_flow_id,
hunt_id=hunt_obj.hunt_id,
message="blah_%d" % i)
])
