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
