def testLimitPropagation(self):
"""This tests that client actions are limited properly."""
client_mock = action_mocks.CPULimitClientMock(
user_cpu_usage=[10],
system_cpu_usage=[10],
network_usage=[1000],
runtime_us=[rdfvalue.Duration.From(1, rdfvalue.SECONDS)])
flow_test_lib.StartAndRunFlow(flow_test_lib.CPULimitFlow,
client_mock=client_mock,
client_id=self.client_id,
cpu_limit=1000,
network_bytes_limit=10000,
runtime_limit=rdfvalue.Duration.From(
5, rdfvalue.SECONDS))
self.assertEqual(client_mock.storage["cpulimit"], [1000, 980, 960])
self.assertEqual(client_mock.storage["networklimit"],
[10000, 9000, 8000])
self.assertEqual(client_mock.storage["networklimit"],
[10000, 9000, 8000])
self.assertEqual(client_mock.storage["runtimelimit"], [
rdfvalue.Duration.From(5, rdfvalue.SECONDS),
rdfvalue.Duration.From(4, rdfvalue.SECONDS),
rdfvalue.Duration.From(3, rdfvalue.SECONDS),
])
def _RunFlow(self, client_id):
flow_args = transfer.GetFileArgs(
pathspec=rdf_paths.PathSpec(
path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS))
client_mock = hunt_test_lib.SampleHuntMock()
if data_store.RelationalDBFlowsEnabled():
with test_lib.FakeTime(42):
return flow_test_lib.StartAndRunFlow(
transfer.GetFile,
client_id=client_id,
client_mock=client_mock,
flow_args=flow_args)
else:
runner_args = rdf_flow_runner.FlowRunnerArgs(
flow_name=transfer.GetFile.__name__)
with test_lib.FakeTime(42):
flow_urn = flow.StartAFF4Flow(
client_id=client_id,
args=flow_args,
runner_args=runner_args,
token=self.token)
flow_test_lib.TestFlowHelper(
flow_urn,
client_mock=client_mock,
client_id=client_id,
token=self.token)
return flow_urn.Basename()
def Run(self):
client_id = self.SetupClient(0)
failing_descriptor = rdf_output_plugin.OutputPluginDescriptor(
plugin_name=hunt_test_lib.FailingDummyHuntOutputPlugin.__name__)
with test_lib.FakeTime(42):
if data_store.RelationalDBFlowsEnabled():
flow_id = flow_test_lib.StartAndRunFlow(
flow_cls=flow_test_lib.DummyFlowWithSingleReply,
client_id=client_id.Basename(),
output_plugins=[failing_descriptor])
else:
flow_urn = flow.StartAFF4Flow(
flow_name=flow_test_lib.DummyFlowWithSingleReply.__name__,
client_id=client_id,
output_plugins=[failing_descriptor],
token=self.token)
flow_id = flow_urn.Basename()
flow_test_lib.TestFlowHelper(flow_urn, token=self.token)
self.Check(
"ListFlowOutputPluginErrors",
args=flow_plugin.ApiListFlowOutputPluginErrorsArgs(
client_id=client_id.Basename(),
flow_id=flow_id,
plugin_id="FailingDummyHuntOutputPlugin_0"),
replace={flow_id: "W:ABCDEF"})
def Run(self):
client_id = self.SetupClient(0)
email_descriptor = rdf_output_plugin.OutputPluginDescriptor(
plugin_name=email_plugin.EmailOutputPlugin.__name__,
plugin_args=email_plugin.EmailOutputPluginArgs(
email_address="test@localhost", emails_limit=42))
with test_lib.FakeTime(42):
if data_store.RelationalDBFlowsEnabled():
flow_id = flow_test_lib.StartAndRunFlow(
flow_cls=flow_test_lib.DummyFlowWithSingleReply,
client_id=client_id.Basename(),
output_plugins=[email_descriptor])
else:
flow_urn = flow.StartAFF4Flow(
flow_name=flow_test_lib.DummyFlowWithSingleReply.__name__,
client_id=client_id,
output_plugins=[email_descriptor],
token=self.token)
flow_id = flow_urn.Basename()
flow_test_lib.TestFlowHelper(flow_urn, token=self.token)
self.Check(
"ListFlowOutputPluginLogs",
args=flow_plugin.ApiListFlowOutputPluginLogsArgs(
client_id=client_id.Basename(),
flow_id=flow_id,
plugin_id="EmailOutputPlugin_0"),
replace={flow_id: "W:ABCDEF"})
def testKeepAliveRunsSuccessfully(self):
client_ids = self.SetupClients(1)
client_mock = action_mocks.ActionMock(admin.Echo)
flow_test_lib.StartAndRunFlow(administrative.KeepAlive,
duration=rdfvalue.Duration("1s"),
client_id=client_ids[0].Basename(),
client_mock=client_mock)
def testCallState(self):
"""Test the ability to chain flows."""
CallStateFlow.success = False
# Run the flow in the simulated way
flow_test_lib.StartAndRunFlow(CallStateFlow, client_id=self.client_id)
self.assertEqual(CallStateFlow.success, True)
def setUp(self):
super().setUp()
self.handler = flow_plugin.ApiListFlowResultsHandler()
self.client_id = self.SetupClient(0)
self.flow_id = flow_test_lib.StartAndRunFlow(
DummyFlowWithTwoTaggedReplies, client_id=self.client_id)
def testChainedFlow(self):
"""Test the ability to chain flows."""
ParentFlow.success = False
# Run the flow in the simulated way
flow_test_lib.StartAndRunFlow(
ParentFlow, client_mock=ClientMock(), client_id=self.client_id)
self.assertEqual(ParentFlow.success, True)
def _RunBrowserHistoryFlow(self, **kwargs):
flow_args = webhistory.BrowserHistoryFlowArgs(**kwargs)
flow_id = flow_test_lib.StartAndRunFlow(
webhistory.BrowserHistoryFlow,
creator=self.token.username,
client_mock=action_mocks.ActionMock(),
client_id=self.client_id,
flow_args=flow_args)
return flow_test_lib.GetFlowResults(self.client_id, flow_id)
def testIncrementalCallbackIsCalledWhenAllResponsesArriveAtOnce(self):
flow_id = flow_test_lib.StartAndRunFlow(FlowWithIncrementalCallback,
client_mock=ClientMock(),
client_id=self.client_id)
results = flow_test_lib.GetFlowResults(self.client_id, flow_id)
self.assertListEqual(results, [
rdfvalue.RDFString("Hello World"),
rdfvalue.RDFString("Final: Hello World")
])
def _RunFlow(self, client_id):
flow_args = transfer.GetFileArgs(pathspec=rdf_paths.PathSpec(
path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS))
client_mock = hunt_test_lib.SampleHuntMock(failrate=2)
with test_lib.FakeTime(42):
return flow_test_lib.StartAndRunFlow(transfer.GetFile,
client_id=client_id,
client_mock=client_mock,
flow_args=flow_args)
def testBrokenChainedFlow(self):
BrokenParentFlow.success = False
# Run the flow in the simulated way
with test_lib.SuppressLogs():
flow_test_lib.StartAndRunFlow(BrokenParentFlow,
client_mock=ClientMock(),
client_id=self.client_id,
check_flow_errors=False)
self.assertEqual(BrokenParentFlow.success, True)
def testUserGetsNotificationWithNumberOfResults(self):
username = "******"
self.CreateUser(username)
flow_test_lib.StartAndRunFlow(
FlowWithMultipleResultTypes, client_id=self.client_id, creator=username)
notifications = self.GetUserNotifications(username)
self.assertIn("FlowWithMultipleResultTypes completed with 6 results",
notifications[0].message)
def testProgressCorrectlyIndicatesNotFoundStatus(self):
flow_id = flow_test_lib.StartAndRunFlow(file.CollectSingleFile,
self.client_mock,
client_id=self.client_id,
check_flow_errors=False,
path="/nonexistent")
progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id)
self.assertEqual(
progress.status,
rdf_file_finder.CollectSingleFileProgress.Status.NOT_FOUND)
def testNestedFlowsHaveTheirResultsSaved(self):
# Run the flow in the simulated way
parent_flow_id = flow_test_lib.StartAndRunFlow(
ParentFlow, client_mock=ClientMock(), client_id=self.client_id)
child_flows = data_store.REL_DB.ReadChildFlowObjects(
self.client_id, parent_flow_id)
self.assertLen(child_flows, 1)
child_flow_results = flow_test_lib.GetFlowResults(
self.client_id, child_flows[0].flow_id)
self.assertNotEmpty(child_flow_results)
def _RunCollectBrowserHistory(self, **kwargs):
flow_args = webhistory.CollectBrowserHistoryArgs(**kwargs)
flow_id = flow_test_lib.StartAndRunFlow(
webhistory.CollectBrowserHistory,
creator=self.token.username,
client_mock=action_mocks.ActionMock(),
client_id=self.client_id,
flow_args=flow_args)
results = flow_test_lib.GetFlowResults(self.client_id, flow_id)
progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id)
return flow_id, results, progress
def testProgressContainsResultOnSuccess(self):
flow_id = flow_test_lib.StartAndRunFlow(
file.CollectSingleFile,
self.client_mock,
client_id=self.client_id,
path=self.files["bar"].path)
progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id)
self.assertEqual(progress.status,
rdf_file_finder.CollectSingleFileProgress.Status.COLLECTED)
self.assertEqual(progress.result.stat.pathspec.path, self.files["bar"].path)
self.assertEqual(progress.result.stat.pathspec.pathtype,
rdf_paths.PathSpec.PathType.OS)
self.assertEqual(str(progress.result.hash.sha1), self.files["bar"].sha1)
def testLimitPropagation(self):
"""This tests that client actions are limited properly."""
client_mock = action_mocks.CPULimitClientMock(user_cpu_usage=[10],
system_cpu_usage=[10],
network_usage=[1000])
flow_test_lib.StartAndRunFlow(flow_test_lib.CPULimitFlow,
client_mock=client_mock,
client_id=self.client_id,
cpu_limit=1000,
network_bytes_limit=10000)
self.assertEqual(client_mock.storage["cpulimit"], [1000, 980, 960])
self.assertEqual(client_mock.storage["networklimit"],
[10000, 9000, 8000])
def testSendReplyWorksCorrectlyInIncrementalCallback(self):
flow_id = flow_test_lib.StartAndRunFlow(
FlowWithIncrementalCallback,
client_mock=NotSendingStatusClientMock(),
client_id=self.client_id,
# Set check_flow_errors to False, otherwise test runner will complain
# that the flow has finished in the RUNNING state.
check_flow_errors=False)
flow_obj = flow_test_lib.GetFlowObj(self.client_id, flow_id)
self.assertEqual(flow_obj.flow_state, flow_obj.FlowState.RUNNING)
results = flow_test_lib.GetFlowResults(self.client_id, flow_id)
self.assertListEqual(results, [
rdfvalue.RDFString(f"Hello World {i}") for i in range(
NotSendingStatusClientMock.NUM_INCREMENTAL_RESPONSES)
])
def testNetworkLimitExceeded(self):
"""This tests that the network limit for flows is working."""
client_mock = action_mocks.CPULimitClientMock(user_cpu_usage=[10],
system_cpu_usage=[10],
network_usage=[1000])
with test_lib.SuppressLogs():
flow_id = flow_test_lib.StartAndRunFlow(flow_test_lib.CPULimitFlow,
client_mock=client_mock,
client_id=self.client_id,
cpu_limit=1000,
network_bytes_limit=1500,
check_flow_errors=False)
rdf_flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(rdf_flow.flow_state, "ERROR")
self.assertIn("bytes limit exceeded", rdf_flow.error_message)
def testProgressCorrectlyIndicatesErrorStatus(self):
with mock.patch.object(flow_base.FlowBase, "client_os", "Windows"):
with mock.patch.object(vfs, "VFSOpen", side_effect=IOError("mock err")):
flow_id = flow_test_lib.StartAndRunFlow(
file.CollectSingleFile,
self.client_mock,
client_id=self.client_id,
check_flow_errors=False,
path="/nonexistent")
progress = flow_test_lib.GetFlowProgress(self.client_id, flow_id)
self.assertEqual(progress.status,
rdf_file_finder.CollectSingleFileProgress.Status.FAILED)
self.assertEqual(
progress.error_description,
f"mock err when fetching /nonexistent with {config.CONFIG['Server.raw_filesystem_access_pathtype']}"
)
def Run(self):
client_id = self.SetupClient(0)
failing_descriptor = rdf_output_plugin.OutputPluginDescriptor(
plugin_name=hunt_test_lib.FailingDummyHuntOutputPlugin.__name__)
with test_lib.FakeTime(42):
flow_id = flow_test_lib.StartAndRunFlow(
flow_cls=flow_test_lib.DummyFlowWithSingleReply,
client_id=client_id,
output_plugins=[failing_descriptor])
self.Check("ListFlowOutputPluginErrors",
args=flow_plugin.ApiListFlowOutputPluginErrorsArgs(
client_id=client_id,
flow_id=flow_id,
plugin_id="FailingDummyHuntOutputPlugin_0"),
replace={flow_id: "W:ABCDEF"})
def testIncrementalCallbackIsNotCalledWhenStatusMessageArrivesEarly(
self, m):
# Mocks don't have names by default.
m.__name__ = "ReceiveHelloCallback"
flow_id = flow_test_lib.StartAndRunFlow(
FlowWithIncrementalCallback,
client_mock=StatusOnlyClientMock(),
client_id=self.client_id,
# Set check_flow_errors to False, otherwise test runner will complain
# that the flow has finished in the RUNNING state.
check_flow_errors=False)
flow_obj = flow_test_lib.GetFlowObj(self.client_id, flow_id)
self.assertEqual(flow_obj.flow_state, flow_obj.FlowState.RUNNING)
self.assertEqual(m.call_count, 0)
def testRuntimeLimitExceeded(self):
client_mock = action_mocks.CPULimitClientMock(
user_cpu_usage=[1],
system_cpu_usage=[1],
network_usage=[1],
runtime_us=[rdfvalue.Duration.From(4, rdfvalue.SECONDS)])
with test_lib.SuppressLogs():
flow_id = flow_test_lib.StartAndRunFlow(
flow_test_lib.CPULimitFlow,
client_mock=client_mock,
client_id=self.client_id,
runtime_limit=rdfvalue.Duration.From(9, rdfvalue.SECONDS),
check_flow_errors=False)
rdf_flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(rdf_flow.flow_state, "ERROR")
self.assertIn("Runtime limit exceeded", rdf_flow.error_message)
def Run(self):
client_id = self.SetupClient(0)
email_descriptor = rdf_output_plugin.OutputPluginDescriptor(
plugin_name=email_plugin.EmailOutputPlugin.__name__,
plugin_args=email_plugin.EmailOutputPluginArgs(
email_address="test@localhost", emails_limit=42))
with test_lib.FakeTime(42):
flow_id = flow_test_lib.StartAndRunFlow(
flow_cls=flow_test_lib.DummyFlowWithSingleReply,
client_id=client_id,
output_plugins=[email_descriptor])
self.Check("ListFlowOutputPluginLogs",
args=flow_plugin.ApiListFlowOutputPluginLogsArgs(
client_id=client_id,
flow_id=flow_id,
plugin_id="EmailOutputPlugin_0"),
replace={flow_id: "W:ABCDEF"})
def RunFlow(self,
flow_cls=None,
output_plugins=None,
flow_args=None,
client_mock=None):
if flow_args is None:
flow_args = transfer.GetFileArgs(pathspec=rdf_paths.PathSpec(
path="/tmp/evil.txt", pathtype=rdf_paths.PathSpec.PathType.OS))
if client_mock is None:
client_mock = hunt_test_lib.SampleHuntMock(failrate=2)
flow_urn = flow_test_lib.StartAndRunFlow(flow_cls or transfer.GetFile,
client_mock=client_mock,
client_id=self.client_id,
flow_args=flow_args,
output_plugins=output_plugins)
return flow_urn
def testCreatorPropagation(self):
username = u"original user"
data_store.REL_DB.WriteGRRUser(username)
client_mock = ClientMock()
flow_id = flow_test_lib.StartAndRunFlow(flow_cls=ParentFlow,
client_id=self.client_id,
creator=username,
client_mock=client_mock)
flow_obj = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(flow_obj.creator, username)
child_flows = data_store.REL_DB.ReadChildFlowObjects(
self.client_id, flow_id)
self.assertLen(child_flows, 1)
child_flow = child_flows[0]
self.assertEqual(child_flow.creator, username)
def testIncrementalCallbackReturnsResultsBeforeStatus(self, m):
# Mocks don't have names by default.
m.__name__ = "ReceiveHelloCallback"
flow_id = flow_test_lib.StartAndRunFlow(
FlowWithIncrementalCallback,
client_mock=NotSendingStatusClientMock(),
client_id=self.client_id,
# Set check_flow_errors to False, otherwise test runner will complain
# that the flow has finished in the RUNNING state.
check_flow_errors=False)
flow_obj = flow_test_lib.GetFlowObj(self.client_id, flow_id)
self.assertEqual(flow_obj.flow_state, flow_obj.FlowState.RUNNING)
self.assertEqual(m.call_count,
NotSendingStatusClientMock.NUM_INCREMENTAL_RESPONSES)
for i in range(NotSendingStatusClientMock.NUM_INCREMENTAL_RESPONSES):
# Get the positional arguments of each call.
args = m.mock_calls[i][1]
# Compare the first positional argument ('responses') to the responses
# list that we expect to have been passed to the callback.
self.assertEqual(list(args[0]),
[rdfvalue.RDFString(f"Hello World {i}")])
