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))