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_urn = flow.StartAFF4Flow( flow_name=processes.ListProcesses.__name__, client_id=client_id, output_plugins=[email_descriptor], token=self.token) self.Check("ListFlowOutputPlugins", args=flow_plugin.ApiListFlowOutputPluginsArgs( client_id=client_id.Basename(), flow_id=flow_urn.Basename()), replace={flow_urn.Basename(): "W:ABCDEF"})
def testNotificationRacesAreResolved(self): # We need a random flow object for this test. session_id = flow.StartAFF4Flow(client_id=self.client_id, flow_name="WorkerSendingTestFlow", token=self.token) worker_obj = self._TestWorker() manager = queue_manager.QueueManager(token=self.token) manager.DeleteNotification(session_id) manager.Flush() # We simulate a race condition here - the notification for request #1 is # there but the actual request #1 is not. The worker should pick up the # notification, notice that the request #1 is not there yet and reschedule # the notification. notification = rdf_flows.GrrNotification(session_id=session_id, last_status=1) with data_store.DB.GetMutationPool() as pool: manager.NotifyQueue(notification, mutation_pool=pool) notifications = manager.GetNotifications(queues.FLOWS) # Check the notification is there. notifications = [ n for n in notifications if n.session_id == session_id ] self.assertLen(notifications, 1) # Process all messages worker_obj.RunOnce() worker_obj.thread_pool.Join() delay = flow_runner.FlowRunner.notification_retry_interval with test_lib.FakeTime(time.time() + 10 + delay): requeued_notifications = manager.GetNotifications(queues.FLOWS) # Check that there is a new notification. notifications = [ n for n in notifications if n.session_id == session_id ] self.assertLen(requeued_notifications, 1) self.assertEqual(requeued_notifications[0].first_queued, notifications[0].first_queued) self.assertNotEqual(requeued_notifications[0].timestamp, notifications[0].timestamp)
def testResolvesNestedFlowURN(self): flow_urn = flow.StartAFF4Flow( flow_name=flow_test_lib.FlowWithOneNestedFlow.__name__, client_id=self.client_urn, token=self.token) children = list( aff4.FACTORY.MultiOpen( list(aff4.FACTORY.ListChildren(flow_urn)), aff4_type=flow.GRRFlow, token=self.token)) self.assertEqual(len(children), 1) flow_id = flow_plugin.ApiFlowId(flow_urn.Basename() + "/" + children[0] .urn.Basename()) self.assertEqual( flow_id.ResolveClientFlowURN( client_plugin.ApiClientId(self.client_urn), token=self.token), children[0].urn)
def testStuckNotificationGetsDeletedAfterTheFlowIsTerminated(self): worker_obj = self._TestWorker() initial_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(100) stuck_flows_timeout = flow_runner.FlowRunner.stuck_flows_timeout try: with test_lib.FakeTime(initial_time.AsSecondsSinceEpoch()): session_id = flow.StartAFF4Flow( flow_name=WorkerStuckableTestFlow.__name__, client_id=self.client_id, token=self.token, sync=False) # Process all messages worker_obj.RunOnce() # Wait until worker thread starts processing the flow. WorkerStuckableTestFlow.WaitUntilWorkerStartsProcessing() # Set the time to max worker flow duration + 1 minute. The flow is # currently blocked because of the way how semaphores are set up. # Worker should consider the flow to be stuck and terminate it. future_time = (initial_time + rdfvalue.Duration("1m") + stuck_flows_timeout) with test_lib.FakeTime(future_time.AsSecondsSinceEpoch()): worker_obj.RunOnce() killed_flow = aff4.FACTORY.Open(session_id, token=self.token) self.assertEqual(killed_flow.context.state, rdf_flow_runner.FlowContext.State.ERROR) self.assertEqual( killed_flow.context.status, "Terminated by user test. Reason: Stuck in the worker") # Check that stuck notification has been removed. qm = queue_manager.QueueManager(token=self.token) notifications = qm.GetNotifications(queues.FLOWS) for n in notifications: self.assertFalse(n.in_progress) finally: # Release the semaphore so that worker thread unblocks and finishes # processing the flow. WorkerStuckableTestFlow.LetWorkerFinishProcessing() worker_obj.thread_pool.Join()
def testWaitUntilDoneRaisesWhenFlowFails(self): client_urn = self.SetupClient(0) flow_urn = flow.StartAFF4Flow( client_id=client_urn, flow_name=processes.ListProcesses.__name__, token=self.token) result_flow = self.api.Client(client_id=client_urn.Basename()).Flow( flow_urn.Basename()).Get() def ProcessFlow(): time.sleep(1) with aff4.FACTORY.Open(flow_urn, mode="rw", token=self.token) as fd: fd.GetRunner().Error("") threading.Thread(target=ProcessFlow).start() with self.assertRaises(grr_api_errors.FlowFailedError): result_flow.WaitUntilDone()
def testStuckFlowGetsTerminated(self): worker_obj = self._TestWorker() initial_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(100) with test_lib.SuppressLogs(): try: with test_lib.FakeTime(initial_time.AsSecondsSinceEpoch()): session_id = flow.StartAFF4Flow( flow_name=WorkerStuckableTestFlow.__name__, client_id=self.client_id, token=self.token, sync=False) # Process all messages while worker_obj.RunOnce(): pass # Wait until worker thread starts processing the flow. WorkerStuckableTestFlow.WaitUntilWorkerStartsProcessing() # Set the time to max worker flow duration + 1 minute. The flow is # currently blocked because of the way semaphores are set up. # Worker should consider the flow to be stuck and terminate it. stuck_flows_timeout = flow_runner.FlowRunner.stuck_flows_timeout future_time = (initial_time + rdfvalue.Duration("1m") + stuck_flows_timeout) with test_lib.FakeTime(future_time.AsSecondsSinceEpoch()): worker_obj.RunOnce() finally: # Release the semaphore so that worker thread unblocks and finishes # processing the flow. WorkerStuckableTestFlow.StopFlow() WorkerStuckableTestFlow.LetWorkerFinishProcessing() worker_obj.thread_pool.Join() killed_flow = aff4.FACTORY.Open(session_id, token=self.token) self.assertEqual(killed_flow.context.state, rdf_flow_runner.FlowContext.State.ERROR) self.assertEqual( killed_flow.context.status, "Terminated by user test. Reason: Stuck in the worker")
def testCreatorPropagation(self): # Instantiate the flow using one username. session_id = flow.StartAFF4Flow( client_id=self.client_id, flow_name="ParentFlow", sync=False, token=access_control.ACLToken( username="******", reason="testing")) # Run the flow using another user ("test"). flow_test_lib.TestFlowHelper( session_id, ClientMock(), client_id=self.client_id, token=self.token) self.assertEqual(ParentFlow.success, True) subflows = list( obj for obj in aff4.FACTORY.Open(session_id, token=self.token) .OpenChildren() if isinstance(obj, flow.GRRFlow)) self.assertEqual(len(subflows), 1) self.assertEqual(subflows[0].GetRunner().context.creator, "original_user")
def Handle(self, args, token=None): if not args.client_id: raise RuntimeError("Client id has to be specified.") if not args.flow.name: raise RuntimeError("Flow name is not specified.") # Note that runner_args are dropped. From all the arguments We use only # the flow name and the arguments. flow_id = flow.StartAFF4Flow( client_id=args.client_id.ToClientURN(), flow_name=args.flow.name, token=token, args=self.override_flow_args or args.flow.args) with aff4.FACTORY.Open( flow_id, aff4_type=flow.GRRFlow, mode="rw", token=token) as fd: fd.AddLabel(LABEL_NAME_PREFIX + self.robot_id) return api_flow.ApiFlow().InitFromAff4Object( fd, flow_id=flow_id.Basename())
def testNoRequestChildFlowRace(self): manager = queue_manager.QueueManager(token=self.token) self.old_notify = manager._MultiNotifyQueue with utils.Stubber(queue_manager.QueueManager, "_MultiNotifyQueue", self.CollectNotifications): session_id = flow.StartAFF4Flow(client_id=self.client_id, flow_name="NoRequestParentFlow", token=self.token) self.assertIn(session_id, self.notifications) f = aff4.FACTORY.Open(session_id, token=self.token) # Check that the first notification came in after the flow was created. self.assertLess( int(f.Get(f.Schema.TYPE).age), 1e6 * min(self.notifications[session_id]), "There was a notification for a flow before " "the flow was created.")
def testTypeAttributeIsNotAppendedWhenFlowIsClosed(self): session_id = flow.StartAFF4Flow( client_id=self.client_id, flow_name=flow_test_lib.FlowOrderTest.__name__, token=self.token) flow_obj = aff4.FACTORY.Open(session_id, aff4_type=flow_test_lib.FlowOrderTest, age=aff4.ALL_TIMES, mode="rw", token=self.token) flow_obj.Close() flow_obj = aff4.FACTORY.Open(session_id, aff4_type=flow_test_lib.FlowOrderTest, age=aff4.ALL_TIMES, token=self.token) types = list(flow_obj.GetValuesForAttribute(flow_obj.Schema.TYPE)) self.assertEqual(len(types), 1)
def testWaitUntilDoneReturnsWhenFlowCompletes(self): client_urn = self.SetupClient(0) flow_urn = flow.StartAFF4Flow( client_id=client_urn, flow_name=processes.ListProcesses.__name__, token=self.token) result_flow = self.api.Client(client_id=client_urn.Basename()).Flow( flow_urn.Basename()).Get() self.assertEqual(result_flow.data.state, result_flow.data.RUNNING) def ProcessFlow(): time.sleep(1) client_mock = action_mocks.ListProcessesMock([]) flow_test_lib.TestFlowHelper( flow_urn, client_mock, client_id=client_urn, token=self.token) threading.Thread(target=ProcessFlow).start() f = result_flow.WaitUntilDone() self.assertEqual(f.data.state, f.data.TERMINATED)
def ProcessHuntOutputPlugins(self): if data_store.RelationalDBEnabled(): # No processing needed for new style hunts. return if data_store.RelationalDBEnabled(): job = rdf_cronjobs.CronJob( cron_job_id="some/id", lifetime=rdfvalue.DurationSeconds("1h")) run_state = rdf_cronjobs.CronJobRun( cron_job_id="some/id", status="RUNNING", started_at=rdfvalue.RDFDatetime.Now()) process_results.ProcessHuntResultCollectionsCronJob(run_state, job).Run() else: flow_urn = flow.StartAFF4Flow( flow_name=process_results.ProcessHuntResultCollectionsCronFlow .__name__, token=self.token) flow_test_lib.TestFlowHelper(flow_urn, token=self.token) return flow_urn
def ProcessMessage(self, message): """Begins an enrollment flow for this client. Args: message: The Certificate sent by the client. Note that this message is not authenticated. """ cert = rdf_crypto.Certificate(message.payload) queue = self.well_known_session_id.Queue() client_id = message.source # It makes no sense to enrol the same client multiple times, so we # eliminate duplicates. Note, that we can still enroll clients multiple # times due to cache expiration. try: enrolment_cache.Get(client_id) return except KeyError: enrolment_cache.Put(client_id, 1) # Create a new client object for this client. client = aff4.FACTORY.Create(client_id, aff4_grr.VFSGRRClient, mode="rw", token=self.token) # Only enroll this client if it has no certificate yet. if not client.Get(client.Schema.CERT): # Start the enrollment flow for this client. # Note, that the actual CAEnroler class is autogenerated from the # CAEnrolerMixin by the DualDBFlow decorator confusing the linter - hence # the disable directive. flow.StartAFF4Flow( client_id=client_id, flow_name=CAEnroler.__name__, # pylint: disable=undefined-variable csr=cert, queue=queue, token=self.token)
def StartFlowAndWait(client_id, token=None, timeout=DEFAULT_TIMEOUT, **flow_args): """Runs a flow and waits for it to finish. Args: client_id: The client id of the client to run on. token: The datastore access token. timeout: How long to wait for a flow to complete, maximum. **flow_args: Pass through to flow. Returns: The urn of the flow that was run. """ flow_urn = flow.StartAFF4Flow( client_id=client_id, token=token, sync=True, **flow_args) WaitForFlow(flow_urn, token=token, timeout=timeout) return flow_urn
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_urn = flow.StartAFF4Flow( flow_name=flow_test_lib.DummyFlowWithSingleReply.__name__, client_id=client_id, output_plugins=[failing_descriptor], token=self.token) with test_lib.FakeTime(43): flow_test_lib.TestFlowHelper(flow_urn, token=self.token) self.Check("ListFlowOutputPluginErrors", args=flow_plugin.ApiListFlowOutputPluginErrorsArgs( client_id=client_id.Basename(), flow_id=flow_urn.Basename(), plugin_id="FailingDummyHuntOutputPlugin_0"), replace={flow_urn.Basename(): "W:ABCDEF"})
def testChildTermination(self): session_id = flow.StartAFF4Flow(client_id=self.client_id, flow_name="CallClientParentFlow", token=self.token) # The child URN should be contained within the parent session_id URN. flow_obj = aff4.FACTORY.Open(session_id, token=self.token) children = list(obj for obj in flow_obj.OpenChildren() if isinstance(obj, flow.GRRFlow)) self.assertEqual(len(children), 1) reason = "just so" flow.GRRFlow.TerminateAFF4Flow(session_id, reason=reason, token=self.token) flow_obj = aff4.FACTORY.Open(session_id, aff4_type=CallClientParentFlow, token=self.token) runner = flow_obj.GetRunner() self.assertEqual(runner.IsRunning(), False) self.assertEqual(runner.context.state, rdf_flow_runner.FlowContext.State.ERROR) self.assertTrue("user test" in runner.context.status) self.assertTrue(reason in runner.context.status) child = aff4.FACTORY.Open(children[0].urn, aff4_type=CallClientChildFlow, token=self.token) runner = child.GetRunner() self.assertEqual(runner.IsRunning(), False) self.assertEqual(runner.context.state, rdf_flow_runner.FlowContext.State.ERROR) self.assertTrue("user test" in runner.context.status) self.assertTrue("Parent flow terminated." in runner.context.status)
def testShowsNotificationWhenArchiveGenerationIsDone(self): pathspec = rdf_paths.PathSpec( path=os.path.join(self.base_path, "test.plist"), pathtype=rdf_paths.PathSpec.PathType.OS) flow_urn = flow.StartAFF4Flow( flow_name=flows_transfer.GetFile.__name__, client_id=self.client_id, pathspec=pathspec, token=self.token) flow_test_lib.TestFlowHelper( flow_urn, self.action_mock, client_id=self.client_id, token=self.token) self.Open("/#/clients/%s" % self.client_id) self.Click("css=a[grrtarget='client.flows']") self.Click("css=td:contains('GetFile')") self.Click("link=Results") self.Click("css=button.DownloadButton") self.WaitUntil(self.IsTextPresent, "Generation has started") self.WaitUntil(self.IsUserNotificationPresent, "Downloaded archive of flow %s" % flow_urn.Basename())
def testInspect(self): """Test the inspect UI.""" client_id = self.SetupClient(0) self.RequestAndGrantClientApproval(client_id) flow.StartAFF4Flow(client_id=rdf_client.ClientURN(client_id), flow_name=flow_discovery.Interrogate.__name__, token=self.token) mock = flow_test_lib.MockClient(client_id, None, token=self.token) while mock.Next(): pass self.Open("/#/clients/%s/debug-requests" % client_id.Basename()) # Check that the we can see both requests and responses. self.WaitUntil(self.IsTextPresent, "GetPlatformInfo") self.WaitUntil(self.IsTextPresent, "GetConfig") self.WaitUntil(self.IsTextPresent, "EnumerateInterfaces") self.WaitUntil(self.IsTextPresent, "GENERIC_ERROR") self.WaitUntil(self.IsTextPresent, "STATUS") self.WaitUntil(self.IsTextPresent, "Task id")
def testCancelFlowWorksCorrectly(self): """Tests that cancelling flows works.""" flow.StartAFF4Flow( client_id=self.client_id, flow_name=gui_test_lib.RecursiveTestFlow.__name__, token=self.token) # Open client and find the flow self.Open("/") self.Type("client_query", self.client_id) self.Click("client_query_submit") self.WaitUntilEqual(self.client_id, self.GetText, "css=span[type=subject]") self.Click("css=td:contains('0001')") self.Click("css=a[grrtarget='client.flows']") self.Click("css=td:contains('RecursiveTestFlow')") self.Click("css=button[name=cancel_flow]") # The window should be updated now self.WaitUntil(self.IsTextPresent, "Cancelled in GUI")
def testInspect(self): """Test the inspect UI.""" client_urn = self.SetupClient(0) client_id = client_urn.Basename() self.RequestAndGrantClientApproval(client_id) if data_store.RelationalDBEnabled(): flow_id = flow.StartFlow(client_id=client_id, flow_cls=flow_discovery.Interrogate) status = rdf_flow_objects.FlowStatus(client_id=client_id, flow_id=flow_id, request_id=1, response_id=1) data_store.REL_DB.WriteFlowResponses([status]) else: session_id = flow.StartAFF4Flow( client_id=client_urn, flow_name=flow_discovery.Interrogate.__name__, token=self.token) status = rdf_flows.GrrMessage( request_id=1, response_id=1, session_id=session_id, type=rdf_flows.GrrMessage.Type.STATUS, auth_state=rdf_flows.GrrMessage.AuthorizationState. AUTHENTICATED) with queue_manager.QueueManager(token=self.token) as manager: manager.QueueResponse(status) self.Open("/#/clients/%s/debug-requests" % client_id) # Check that the we can see both requests and responses. self.WaitUntil(self.IsTextPresent, "GetPlatformInfo") self.WaitUntil(self.IsTextPresent, "GetConfig") self.WaitUntil(self.IsTextPresent, "EnumerateInterfaces") self.WaitUntil(self.IsTextPresent, "STATUS")
def checkClickingOnDownloadAsStartsDownloadForType(self, mock_method, plugin, plugin_display_name): pathspec = rdf_paths.PathSpec(path=os.path.join( self.base_path, "test.plist"), pathtype=rdf_paths.PathSpec.PathType.OS) flow_urn = flow.StartAFF4Flow( flow_name=flows_transfer.GetFile.__name__, client_id=self.client_id, pathspec=pathspec, token=self.token) flow_test_lib.TestFlowHelper(flow_urn, self.action_mock, client_id=self.client_id, token=self.token) self.Open("/#/clients/%s/flows/%s" % (self.client_id, flow_urn.Basename())) self.Click("link=Results") self.Select("id=plugin-select", plugin_display_name) self.Click("css=grr-download-collection-as button[name='download-as']") def MockMethodIsCalled(): try: # Mock should be called twice: once for HEAD (to check permissions) # and once for GET methods. mock_method.assert_called_with( api_flow.ApiGetExportedFlowResultsArgs( client_id=self.client_id, flow_id=flow_urn.Basename(), plugin_name=plugin), token=mock.ANY) return True except AssertionError: return False self.WaitUntil(MockMethodIsCalled)
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_urn = flow.StartAFF4Flow( flow_name=flow_test_lib.DummyFlowWithSingleReply.__name__, client_id=client_id, output_plugins=[email_descriptor], token=self.token) with test_lib.FakeTime(43): flow_test_lib.TestFlowHelper(flow_urn, token=self.token) self.Check("ListFlowOutputPluginLogs", args=flow_plugin.ApiListFlowOutputPluginLogsArgs( client_id=client_id.Basename(), flow_id=flow_urn.Basename(), plugin_id="EmailOutputPlugin_0"), replace={flow_urn.Basename(): "W:ABCDEF"})
def testExportCommandIsShownForStatEntryResults(self): flow_urn = flow.StartAFF4Flow( flow_name=gui_test_lib.FlowWithOneStatEntryResult.__name__, client_id=self.client_id, token=self.token) flow_test_lib.TestFlowHelper( flow_urn, self.action_mock, client_id=self.client_id, token=self.token) self.Open("/#/clients/%s/flows" % self.client_id) self.Click("css=td:contains('FlowWithOneStatEntryResult')") self.Click("css=li[heading=Results]") self.Click("link=Show export command") self.WaitUntil( self.IsTextPresent, "/usr/bin/grr_api_shell 'http://localhost:8000/' " "--exec_code 'grrapi.Client(\"%s\")." "Flow(\"%s\").GetFilesArchive()." "WriteToFile(\"./flow_results_%s_%s.zip\")'" % ( self.client_id, flow_urn.Basename(), self.client_id.replace(".", "_"), flow_urn.Basename().replace(":", "_"), ))
def testFlowListGetsUpdatedWithChangedFlows(self): f = flow.StartAFF4Flow( client_id=self.client_id, flow_name=gui_test_lib.RecursiveTestFlow.__name__, token=self.token) self.Open("/#/clients/%s" % self.client_id) # Ensure auto-refresh updates happen every second. self.GetJavaScriptValue( "grrUi.flow.flowsListDirective.setAutoRefreshInterval(1000);") # Go to the flows page without refreshing the page, so that # AUTO_REFRESH_INTERVAL_MS setting is not reset. self.Click("css=a[grrtarget='client.flows']") # Check that the flow is running. self.WaitUntil(self.IsElementPresent, "css=tr:contains('%s') div[state=RUNNING]" % f.Basename()) # Cancel the flow and check that the flow state gets updated. flow.GRRFlow.TerminateAFF4Flow(f, "Because I said so", token=self.token) self.WaitUntil(self.IsElementPresent, "css=tr:contains('%s') div[state=ERROR]" % f.Basename())
def RunFlow(self, flow_name, **kwargs): result = {} client_mock = action_mocks.CPULimitClientMock(result) client_mock = flow_test_lib.MockClient( self.client_id, client_mock, token=self.token) worker_mock = ResourcedWorker(check_flow_errors=True, token=self.token) flow.StartAFF4Flow( client_id=self.client_id, flow_name=flow_name, token=self.token, **kwargs) while True: client_processed = client_mock.Next() flows_run = [] for flow_run in worker_mock.Next(): flows_run.append(flow_run) if client_processed == 0 and not flows_run: break return result
def Run(self): client_id = self.SetupClient(0) replace = {} with test_lib.FakeTime(42): flow_urn = flow.StartAFF4Flow( client_id=client_id, flow_name=processes.ListProcesses.__name__, token=self.token) replace[flow_urn.Basename()] = "F:123456" test_process = client_test_lib.MockWindowsProcess( name="test_process") with utils.Stubber(psutil, "Process", lambda: test_process): # Here we emulate a mock client with no actions (None) that # should produce an error. mock = flow_test_lib.MockClient(client_id, None, token=self.token) while mock.Next(): pass manager = queue_manager.QueueManager(token=self.token) requests_responses = manager.FetchRequestsAndResponses(flow_urn) for request, responses in requests_responses: replace[str(request.request.task_id)] = "42" for response in responses: replace[str(response.task_id)] = "43" self.Check("ListClientActionRequests", args=client_plugin.ApiListClientActionRequestsArgs( client_id=client_id.Basename()), replace=replace) self.Check("ListClientActionRequests", args=client_plugin.ApiListClientActionRequestsArgs( client_id=client_id.Basename(), fetch_responses=True), replace=replace)
def testDownloadFilesPanelIsShownWhenNewResultsAreAdded(self): f = flow.StartAFF4Flow( client_id=self.client_id, flow_name=gui_test_lib.RecursiveTestFlow.__name__, token=self.token) with data_store.DB.GetMutationPool() as pool: flow.GRRFlow.ResultCollectionForFID(f).Add( rdfvalue.RDFString("foo-result"), mutation_pool=pool) self.Open("/#/clients/%s" % self.client_id) # Ensure auto-refresh updates happen every second. self.GetJavaScriptValue( "grrUi.core.resultsCollectionDirective.setAutoRefreshInterval(1000);" ) # Go to the flows page without refreshing the page, so that # AUTO_REFRESH_INTERVAL_MS setting is not reset. self.Click("css=a[grrtarget='client.flows']") self.Click("css=tr:contains('%s')" % f.Basename()) self.Click("css=li[heading=Results]:not([disabled]") self.WaitUntil(self.IsElementPresent, "css=grr-results-collection td:contains('foo-result')") self.WaitUntilNot( self.IsElementPresent, "css=grr-results-collection grr-download-collection-files") stat_entry = rdf_client_fs.StatEntry(pathspec=rdf_paths.PathSpec( path="/foo/bar", pathtype=rdf_paths.PathSpec.PathType.OS)) with data_store.DB.GetMutationPool() as pool: flow.GRRFlow.ResultCollectionForFID(f).Add(stat_entry, mutation_pool=pool) self.WaitUntil( self.IsElementPresent, "css=grr-results-collection grr-download-collection-files")
def testWorksCorrectlyWithTestOutputPluginOnFlowWithSingleResult(self): with test_lib.FakeTime(42): flow_urn = flow.StartAFF4Flow( flow_name=flow_test_lib.DummyFlowWithSingleReply.__name__, client_id=self.client_id, token=self.token) flow_test_lib.TestFlowHelper(flow_urn, token=self.token) result = self.handler.Handle(flow_plugin.ApiGetExportedFlowResultsArgs( client_id=self.client_id, flow_id=flow_urn.Basename(), plugin_name=test_plugins.TestInstantOutputPlugin.plugin_name), token=self.token) chunks = list(result.GenerateContent()) self.assertListEqual(chunks, [ "Start: %s" % utils.SmartStr(flow_urn), "Values of type: RDFString", "First pass: oh (source=%s)" % utils.SmartStr(self.client_id), "Second pass: oh (source=%s)" % utils.SmartStr(self.client_id), "Finish: %s" % utils.SmartStr(flow_urn) ]) # pyformat: disable
def testShowsNotificationIfArchiveStreamingFailsInProgress(self): pathspec = rdf_paths.PathSpec(path=os.path.join( self.base_path, "test.plist"), pathtype=rdf_paths.PathSpec.PathType.OS) flow_urn = flow.StartAFF4Flow( flow_name=flows_transfer.GetFile.__name__, client_id=self.client_id, pathspec=pathspec, token=self.token) flow_test_lib.TestFlowHelper(flow_urn, self.action_mock, client_id=self.client_id, token=self.token) def RaisingStub(*unused_args, **unused_kwargs): yield b"foo" yield b"bar" raise RuntimeError("something went wrong") with utils.Stubber(api_call_handler_utils.CollectionArchiveGenerator, "Generate", RaisingStub): self.Open("/#/clients/%s" % self.client_id) self.Click("css=a[grrtarget='client.flows']") self.Click("css=td:contains('GetFile')") self.Click("link=Results") self.Click("css=button.DownloadButton") self.WaitUntil( self.IsUserNotificationPresent, "Archive generation failed for flow %s" % flow_urn.Basename()) # There will be no failure message, as we can't get a status from an # iframe that triggers the download. self.WaitUntilNot(self.IsTextPresent, "Can't generate archive: Unknown error")
def _StartFlow(self, client_id, flow_cls, **kw): if data_store.RelationalDBFlowsEnabled(): flow_id = flow.StartFlow(flow_cls=flow_cls, client_id=client_id, **kw) # Lease the client message. data_store.REL_DB.LeaseClientMessages( client_id, lease_time=rdfvalue.Duration("10000s")) # Write some responses. response = rdf_flow_objects.FlowResponse( client_id=client_id, flow_id=flow_id, request_id=1, response_id=1, payload=rdf_client.Process(name="test_process")) status = rdf_flow_objects.FlowStatus(client_id=client_id, flow_id=flow_id, request_id=1, response_id=2) data_store.REL_DB.WriteFlowResponses([response, status]) return flow_id else: flow_id = flow.StartAFF4Flow( flow_name=compatibility.GetName(flow_cls), client_id=client_id, token=self.token, **kw).Basename() # Have the client write some responses. test_process = rdf_client.Process(name="test_process") mock = flow_test_lib.MockClient(client_id, action_mocks.ListProcessesMock( [test_process]), token=self.token) mock.Next() return flow_id