def testPendingFlowTermination(self):
client_mock = ClientMock()
flow_id = flow.StartFlow(flow_cls=ParentFlow, client_id=self.client_id)
flow_obj = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(flow_obj.flow_state, "RUNNING")
pending_termination = rdf_flow_objects.PendingFlowTermination(
reason="testing")
data_store.REL_DB.UpdateFlow(self.client_id,
flow_id,
pending_termination=pending_termination)
with flow_test_lib.TestWorker() as worker:
with test_lib.SuppressLogs():
flow_test_lib.RunFlow(self.client_id,
flow_id,
client_mock=client_mock,
worker=worker,
check_flow_errors=False)
flow_obj = data_store.REL_DB.ReadFlowObject(
self.client_id, flow_id)
self.assertEqual(flow_obj.flow_state, "ERROR")
self.assertEqual(flow_obj.error_message, "testing")
def testProgress(self):
client_id = self.client_id
with temp.AutoTempDirPath(remove_non_empty=True) as tempdir:
filesystem_test_lib.CreateFile(os.path.join(tempdir, "foo"))
filesystem_test_lib.CreateFile(os.path.join(tempdir, "bar"))
filesystem_test_lib.CreateFile(os.path.join(tempdir, "baz"))
args = rdf_timeline.TimelineArgs()
args.root = tempdir.encode("utf-8")
flow_id = flow_test_lib.StartFlow(timeline_flow.TimelineFlow,
client_id=client_id,
flow_args=args)
progress = flow_test_lib.GetFlowProgress(client_id=client_id,
flow_id=flow_id)
self.assertEqual(progress.total_entry_count, 0)
flow_test_lib.RunFlow(client_id=client_id,
flow_id=flow_id,
client_mock=action_mocks.ActionMock(
timeline_action.Timeline))
progress = flow_test_lib.GetFlowProgress(client_id=client_id,
flow_id=flow_id)
self.assertEqual(progress.total_entry_count, 4)
def testCreatorPropagation(self):
username = u"original user"
data_store.REL_DB.WriteGRRUser(username)
client_mock = ClientMock()
flow_id = flow.StartFlow(flow_cls=ParentFlow,
client_id=self.client_id,
creator=username)
worker = flow_test_lib.TestWorker(token=True)
data_store.REL_DB.RegisterFlowProcessingHandler(worker.ProcessFlow)
flow_test_lib.RunFlow(self.client_id,
flow_id,
client_mock=client_mock,
worker=worker)
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.assertEqual(len(child_flows), 1)
child_flow = child_flows[0]
self.assertEqual(child_flow.creator, username)
def testFileFinderWorkflowWorks(self):
self.InitRouterConfig(self.__class__.FILE_FINDER_ROUTER_CONFIG %
self.token.username)
client_ref = self.api.Client(client_id=self.client_id)
args = rdf_file_finder.FileFinderArgs(
paths=[
os.path.join(self.base_path, "test.plist"),
os.path.join(self.base_path, "numbers.txt"),
os.path.join(self.base_path, "numbers.txt.ver2")
],
action=rdf_file_finder.FileFinderAction.Download()).AsPrimitiveProto()
flow_obj = client_ref.CreateFlow(
name=file_finder.FileFinder.__name__, args=args)
self.assertEqual(flow_obj.data.state, flow_obj.data.RUNNING)
# Now run the flow we just started.
flow_test_lib.RunFlow(
flow_obj.client_id,
flow_obj.flow_id,
client_mock=action_mocks.FileFinderClientMock())
# Refresh flow.
flow_obj = client_ref.Flow(flow_obj.flow_id).Get()
self.assertEqual(flow_obj.data.state, flow_obj.data.TERMINATED)
# Check that we got 3 results (we downloaded 3 files).
results = list(flow_obj.ListResults())
self.assertLen(results, 3)
# We expect results to be FileFinderResult.
self.assertCountEqual(
[os.path.basename(r.payload.stat_entry.pathspec.path) for r in results],
["test.plist", "numbers.txt", "numbers.txt.ver2"])
# Now downloads the files archive.
zip_stream = io.BytesIO()
flow_obj.GetFilesArchive().WriteToStream(zip_stream)
zip_fd = zipfile.ZipFile(zip_stream)
# Now check that the archive has only "test.plist" file, as it's the
# only file that matches the whitelist (see FILE_FINDER_ROUTER_CONFIG).
# There should be 3 items in the archive: the hash of the "test.plist"
# file, the symlink to this hash and the MANIFEST file.
namelist = zip_fd.namelist()
self.assertLen(namelist, 3)
# First component of every path in the archive is the containing folder,
# we should strip it.
namelist = [os.path.join(*n.split(os.sep)[1:]) for n in namelist]
self.assertEqual(
sorted([
# pyformat: disable
os.path.join(self.client_id, "fs", "os", self.base_path.strip("/"),
"test.plist"),
os.path.join(self.client_id, "client_info.yaml"),
"MANIFEST"
# pyformat: enable
]),
sorted(namelist))
def testFailure(self):
stderr = "Error: query syntax error"
with osquery_test_lib.FakeOsqueryiOutput(stdout="", stderr=stderr):
flow_id = flow_test_lib.StartFlow(
flow_cls=osquery_flow.OsqueryFlow,
client_id=self.client_id,
creator=self.test_username,
query="<<<FAKE OSQUERY FLOW QUERY PLACEHOLDER>>>")
with self.assertRaises(RuntimeError):
flow_test_lib.RunFlow(
client_id=self.client_id,
flow_id=flow_id,
client_mock=action_mocks.OsqueryClientMock())
flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(flow.flow_state,
rdf_flow_objects.Flow.FlowState.ERROR)
self.assertIn(stderr, flow.error_message)
def testNotificationIsSent(self):
fixture_test_lib.ClientFixture(self.client_id)
args = vfs_plugin.ApiCreateVfsRefreshOperationArgs(
client_id=self.client_id,
file_path=self.file_path,
max_depth=0,
notify_user=True)
result = self.handler.Handle(args, context=self.context)
flow_test_lib.RunFlow(
self.client_id, result.operation_id, check_flow_errors=False)
pending_notifications = self.GetUserNotifications(self.context.username)
self.assertIn("Recursive Directory Listing complete",
pending_notifications[0].message)
self.assertEqual(
pending_notifications[0].reference.vfs_file.path_components,
["Users", "Shared"])
def testNotificationIsSent(self):
fixture_test_lib.ClientFixture(self.client_id, token=self.token)
args = vfs_plugin.ApiCreateVfsRefreshOperationArgs(
client_id=self.client_id,
file_path=self.file_path,
max_depth=0,
notify_user=True)
result = self.handler.Handle(args, token=self.token)
if data_store.RelationalDBFlowsEnabled():
flow_test_lib.RunFlow(
self.client_id, result.operation_id, check_flow_errors=False)
else:
# Finish flow and check if there are any new notifications.
flow_urn = rdfvalue.RDFURN(result.operation_id)
client_mock = action_mocks.ActionMock()
flow_test_lib.TestFlowHelper(
flow_urn,
client_mock,
client_id=self.client_id,
token=self.token,
check_flow_errors=False)
pending_notifications = self.GetUserNotifications(self.token.username)
self.assertIn("Recursive Directory Listing complete",
pending_notifications[0].message)
if data_store.RelationalDBReadEnabled():
self.assertEqual(
pending_notifications[0].reference.vfs_file.path_components,
["Users", "Shared"])
else:
self.assertEqual(pending_notifications[0].subject,
self.client_id.Add(self.file_path))
