def testStartScheduledFlowsContinuesNextOnFailure(self): client_id = self.SetupClient(0) username = self.SetupUser("u0") self.ScheduleFlow( client_id=client_id, creator=username, flow_name=file.CollectSingleFile.__name__, flow_args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60)) self.ScheduleFlow( client_id=client_id, creator=username, flow_name=file.CollectSingleFile.__name__, flow_args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60)) with mock.patch.object( rdf_file_finder.CollectSingleFileArgs, "Validate", side_effect=[ValueError("foobazzle"), mock.DEFAULT]): flow.StartScheduledFlows(client_id, username) self.assertLen(data_store.REL_DB.ReadAllFlowObjects(client_id), 1) self.assertLen(flow.ListScheduledFlows(client_id, username), 1)
def testUnscheduleFlowRemovesScheduledFlow(self, db: abstract_db.Database): token = _CreateToken(db) client_id = db_test_utils.InitializeClient(db) handler = flow_plugin.ApiScheduleFlowHandler() sf1 = handler.Handle(flow_plugin.ApiCreateFlowArgs( client_id=client_id, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))), token=token) sf2 = handler.Handle(flow_plugin.ApiCreateFlowArgs( client_id=client_id, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))), token=token) handler = flow_plugin.ApiUnscheduleFlowHandler() args = flow_plugin.ApiUnscheduleFlowArgs( client_id=client_id, scheduled_flow_id=sf1.scheduled_flow_id) handler.Handle(args, token=token) handler = flow_plugin.ApiListScheduledFlowsHandler() args = flow_plugin.ApiListScheduledFlowsArgs(client_id=client_id, creator=token.username) results = handler.Handle(args, token=token) self.assertEqual(results.scheduled_flows, [sf2])
def testListScheduledFlows(self, db: abstract_db.Database): context = _CreateContext(db) client_id1 = db_test_utils.InitializeClient(db) client_id2 = db_test_utils.InitializeClient(db) handler = flow_plugin.ApiScheduleFlowHandler() sf1 = handler.Handle(flow_plugin.ApiCreateFlowArgs( client_id=client_id1, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))), context=context) sf2 = handler.Handle(flow_plugin.ApiCreateFlowArgs( client_id=client_id1, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))), context=context) handler.Handle(flow_plugin.ApiCreateFlowArgs( client_id=client_id2, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))), context=context) handler = flow_plugin.ApiListScheduledFlowsHandler() args = flow_plugin.ApiListScheduledFlowsArgs(client_id=client_id1, creator=context.username) results = handler.Handle(args, context=context) self.assertEqual(results.scheduled_flows, [sf1, sf2])
def testCorrectlyDisplaysCollectedResult(self): flow_args = rdf_file_finder.CollectSingleFileArgs(path="/etc/hosts") flow_test_lib.StartFlow(file.CollectSingleFile, creator=self.test_username, client_id=self.client_id, flow_args=flow_args) with flow_test_lib.FlowProgressOverride( file.CollectSingleFile, rdf_file_finder.CollectSingleFileProgress( status=rdf_file_finder.CollectSingleFileProgress.Status. COLLECTED, result=self._GenSampleResult())): self.Open(f"/v2/clients/{self.client_id}") self.WaitUntil( self.IsElementPresent, "css=collect-single-file-details .collected-result:contains('/etc/hosts')" ) self.WaitUntil( self.IsElementPresent, "css=collect-single-file-details .collected-result:contains('4.14 KiB')" ) self.WaitUntilNot( self.IsElementPresent, "css=collect-single-file-details .requested-path")
def testErrorDuringStartFlowDoesNotBubbleUpToApprovalApiCall(self): flow.ScheduleFlow( client_id=self.client_id, creator=self.context.username, flow_name=file.CollectSingleFile.__name__, flow_args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs()) with mock.patch.object( flow, "StartFlow", side_effect=ValueError("foobazzle")) as start_flow_mock: approval_id = self.RequestAndGrantClientApproval( self.client_id, reason=u"blah", approver=u"approver", requestor=self.context.username) args = user_plugin.ApiGetClientApprovalArgs( client_id=self.client_id, approval_id=approval_id, username=self.context.username) handler = user_plugin.ApiGetClientApprovalHandler() approval = handler.Handle(args, context=self.context) self.assertTrue(approval.is_valid) self.assertTrue(start_flow_mock.called)
def ScheduleFlow(self, **kwargs): merged_kwargs = { "flow_name": file.CollectSingleFile.__name__, "flow_args": rdf_file_finder.CollectSingleFileArgs( path="/foo{}".format(random.randint(0, 1000))), "runner_args": rdf_flow_runner.FlowRunnerArgs(cpu_limit=random.randint(0, 60)), **kwargs } return flow.ScheduleFlow(**merged_kwargs)
def testCorrectlyDisplaysNotFoundResult(self): flow_args = rdf_file_finder.CollectSingleFileArgs(path="/etc/hosts") flow_test_lib.StartFlow( file.CollectSingleFile, creator=self.token.username, client_id=self.client_id, flow_args=flow_args) with flow_test_lib.FlowProgressOverride( file.CollectSingleFile, rdf_file_finder.CollectSingleFileProgress( status=rdf_file_finder.CollectSingleFileProgress.Status.NOT_FOUND)): self.Open(f"/v2/clients/{self.client_id}") self.WaitUntil( self.IsElementPresent, "css=collect-single-file-details .error:contains('Not found')")
def testCorrectlyDisplaysError(self): flow_args = rdf_file_finder.CollectSingleFileArgs(path="/etc/hosts") flow_id = flow_test_lib.StartFlow(file.CollectSingleFile, creator=self.test_username, client_id=self.client_id, flow_args=flow_args) flow_test_lib.MarkFlowAsFailed(self.client_id, flow_id) with flow_test_lib.FlowProgressOverride( file.CollectSingleFile, rdf_file_finder.CollectSingleFileProgress( status=rdf_file_finder.CollectSingleFileProgress.Status. FAILED, error_description="Something went wrong")): self.Open(f"/v2/clients/{self.client_id}") self.WaitUntil( self.IsElementPresent, "css=flow-details :contains('Something went wrong')")
def testScheduleFlow(self, db: abstract_db.Database): token = _CreateToken(db) client_id = db_test_utils.InitializeClient(db) handler = flow_plugin.ApiScheduleFlowHandler() args = flow_plugin.ApiCreateFlowArgs( client_id=client_id, flow=flow_plugin.ApiFlow( name=file.CollectSingleFile.__name__, args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60))) sf = handler.Handle(args, token=token) self.assertEqual(sf.client_id, client_id) self.assertEqual(sf.creator, token.username) self.assertNotEmpty(sf.scheduled_flow_id) self.assertEqual(sf.flow_name, file.CollectSingleFile.__name__) self.assertEqual(sf.flow_args.path, "/foo") self.assertEqual(sf.runner_args.cpu_limit, 60)
def testStartScheduledFlowsHandlesErrorInFlowArgsValidation(self): client_id = self.SetupClient(0) username = self.SetupUser("u0") self.ScheduleFlow( client_id=client_id, creator=username, flow_name=file.CollectSingleFile.__name__, flow_args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60)) with mock.patch.object(rdf_file_finder.CollectSingleFileArgs, "Validate", side_effect=ValueError("foobazzle")): flow.StartScheduledFlows(client_id, username) self.assertEmpty(data_store.REL_DB.ReadAllFlowObjects(client_id)) scheduled_flows = flow.ListScheduledFlows(client_id, username) self.assertLen(scheduled_flows, 1) self.assertIn("foobazzle", scheduled_flows[0].error)
def testCorrectlyDisplaysInProgressStatus(self): flow_args = rdf_file_finder.CollectSingleFileArgs(path="/etc/hosts") flow_test_lib.StartFlow(file.CollectSingleFile, creator=self.test_username, client_id=self.client_id, flow_args=flow_args) with flow_test_lib.FlowProgressOverride( file.CollectSingleFile, rdf_file_finder.CollectSingleFileProgress( status=rdf_file_finder.CollectSingleFileProgress.Status. IN_PROGRESS)): self.Open(f"/v2/clients/{self.client_id}") self.WaitUntil(self.IsElementPresent, "css=.flow-title:contains('File content')") self.WaitUntil( self.IsElementPresent, "css=collect-single-file-details .requested-path:contains('/etc/hosts')" ) self.WaitUntilNot( self.IsElementPresent, "css=collect-single-file-details .collected-result")
def testCorrectlyDisplaysNonStandardPathTypeNote(self): flow_args = rdf_file_finder.CollectSingleFileArgs(path="/etc/hosts") flow_test_lib.StartFlow( file.CollectSingleFile, creator=self.token.username, client_id=self.client_id, flow_args=flow_args) with flow_test_lib.FlowProgressOverride( file.CollectSingleFile, rdf_file_finder.CollectSingleFileProgress( status=rdf_file_finder.CollectSingleFileProgress.Status.COLLECTED, result=self._GenSampleResult(use_ntfs=True))): self.Open(f"/v2/clients/{self.client_id}") self.WaitUntil( self.IsElementPresent, "css=collect-single-file-details .collected-result:contains('/etc/hosts')" ) self.WaitUntil( self.IsElementPresent, "css=collect-single-file-details .path-type-note:contains(" "'File was fetched by parsing the raw disk image with libfsntfs')")
def testCorrectlyDisplaysDownloadButtonOnSuccess(self): flow_args = rdf_file_finder.CollectSingleFileArgs(path="/etc/hosts") flow_id = flow_test_lib.StartFlow(file.CollectSingleFile, creator=self.test_username, client_id=self.client_id, flow_args=flow_args) self.Open(f"/v2/clients/{self.client_id}") self.WaitUntil( self.IsElementPresent, "css=collect-single-file-details .requested-path:contains('/etc/hosts')" ) self.WaitUntilNot(self.IsElementPresent, "css=a[mat-stroked-button]:contains('Download')") flow_test_lib.MarkFlowAsFinished(self.client_id, flow_id) with flow_test_lib.FlowProgressOverride( file.CollectSingleFile, rdf_file_finder.CollectSingleFileProgress( status=rdf_file_finder.CollectSingleFileProgress.Status. COLLECTED, result=self._GenSampleResult())): with flow_test_lib.FlowResultMetadataOverride( file.CollectSingleFile, rdf_flow_objects.FlowResultMetadata( is_metadata_set=True, num_results_per_type_tag=[ rdf_flow_objects.FlowResultCount( type=rdf_file_finder.CollectSingleFileResult. __name__, count=1) ])): self.WaitUntil( self.IsElementPresent, "css=a[mat-stroked-button]:contains('Download')")
def testStartScheduledFlowsCreatesFlow(self): client_id = self.SetupClient(0) username = self.SetupUser("u0") self.ScheduleFlow( client_id=client_id, creator=username, flow_name=file.CollectSingleFile.__name__, flow_args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs(cpu_limit=60)) flow.StartScheduledFlows(client_id, username) flows = data_store.REL_DB.ReadAllFlowObjects(client_id) self.assertLen(flows, 1) self.assertEqual(flows[0].client_id, client_id) self.assertEqual(flows[0].creator, username) self.assertEqual(flows[0].flow_class_name, file.CollectSingleFile.__name__) self.assertEqual(flows[0].args.path, "/foo") self.assertEqual(flows[0].flow_state, rdf_flow_objects.Flow.FlowState.RUNNING) self.assertEqual(flows[0].cpu_limit, 60)
def testScheduledFlowsAreShown(self): client_id = self.SetupClient(0) self.CreateUser("requestrick") self.CreateUser("approveannie") flow.ScheduleFlow( client_id=client_id, creator="requestrick", flow_name=file.CollectSingleFile.__name__, flow_args=rdf_file_finder.CollectSingleFileArgs(path="/foo"), runner_args=rdf_flow_runner.FlowRunnerArgs()) approval_id = self.RequestClientApproval(client_id, reason="t/1234", requestor="requestrick", approver="approveannie") self.Open( f"/v2/clients/{client_id}/users/requestrick/approvals/{approval_id}" ) # TODO(user): Change to pretty display name as soon as ScheduledFlowList # uses these. self.WaitUntil(self.IsTextPresent, "CollectSingleFile")
def GetDefaultArgs(cls, username=None): """See base class.""" del username # Unused. return rdf_file_finder.CollectSingleFileArgs(path="", max_size_bytes="1 GiB")