def _GetActionInstance(self, action_cls, grr_worker=None):
"""Run an action and generate responses.
This basically emulates GRRClientWorker.HandleMessage().
Args:
action_cls: The action class to run.
grr_worker: The GRRClientWorker instance to use. If not provided we make
a new one.
Returns:
A list of response protobufs.
"""
# A mock SendReply() method to collect replies.
def mock_send_reply(mock_self, reply=None, **kwargs):
if reply is None:
reply = mock_self.out_rdfvalues[0](**kwargs)
self.results.append(reply)
if grr_worker is None:
grr_worker = worker_mocks.FakeClientWorker()
action = action_cls(grr_worker=grr_worker)
action.SendReply = types.MethodType(mock_send_reply, action)
return action
def testSuspendableListDirectory(self):
request = rdf_client.ListDirRequest()
request.pathspec.path = self.base_path
request.pathspec.pathtype = "OS"
request.iterator.number = 2
results = []
grr_worker = worker_mocks.FakeClientWorker()
while request.iterator.state != request.iterator.State.FINISHED:
responses = self.RunAction("SuspendableListDirectory",
request,
grr_worker=grr_worker)
results.extend(responses)
for response in responses:
if isinstance(response, rdf_client.Iterator):
request.iterator = response
filenames = [
os.path.basename(r.pathspec.path) for r in results
if isinstance(r, rdf_client.StatEntry)
]
self.assertItemsEqual(filenames, os.listdir(self.base_path))
iterators = [r for r in results if isinstance(r, rdf_client.Iterator)]
# One for two files plus one extra with the FINISHED status.
nr_files = len(os.listdir(self.base_path))
expected_iterators = (nr_files / 2) + 1
if nr_files % 2:
expected_iterators += 1
self.assertEqual(len(iterators), expected_iterators)
# Make sure the thread has been deleted.
self.assertEqual(grr_worker.suspended_actions, {})
def __init__(self, *action_classes, **kwargs):
self.client_id = kwargs.get("client_id")
self.action_classes = {cls.__name__: cls for cls in action_classes}
self.action_counts = dict((cls_name, 0) for cls_name in self.action_classes)
self.recorded_args = {}
self.client_worker = (kwargs.get("client_worker", None) or
worker_mocks.FakeClientWorker())
def __init__(self, *action_classes, **kwargs):
self.client_id = kwargs.get("client_id")
self.action_classes = {cls.__name__: cls for cls in action_classes}
self.action_counts = dict(
(cls_name, 0) for cls_name in self.action_classes)
self.recorded_args = {}
# Create a single long lived client worker mock.
self.client_worker = worker_mocks.FakeClientWorker()
def __init__(self, *action_names, **kwargs):
self.client_id = kwargs.get("client_id")
self.action_names = action_names
self.action_classes = dict(
[(k, v) for (k, v) in actions.ActionPlugin.classes.items()
if k in action_names])
self.action_counts = dict((x, 0) for x in action_names)
# Create a single long lived client worker mock.
self.client_worker = worker_mocks.FakeClientWorker()
def testSuspendableActionException(self):
class testActionWorker(actions.ClientActionWorker):
def run(self):
try:
return super(testActionWorker, self).run()
except Exception as e: # pylint: disable=broad-except
logging.info("Expected exception: %s", e)
class RaisingListDirectory(standard.SuspendableListDirectory):
iterations = 3
def Suspend(self):
RaisingListDirectory.iterations -= 1
if not RaisingListDirectory.iterations:
raise IOError("Ran out of iterations.")
return super(RaisingListDirectory, self).Suspend()
p = rdf_paths.PathSpec(path=self.base_path,
pathtype=rdf_paths.PathSpec.PathType.OS)
request = rdf_client.ListDirRequest(pathspec=p)
request.iterator.number = 2
results = []
grr_worker = worker_mocks.FakeClientWorker()
while request.iterator.state != request.iterator.State.FINISHED:
responses = self.ExecuteAction("RaisingListDirectory",
request,
grr_worker=grr_worker,
action_worker_cls=testActionWorker)
results.extend(responses)
for response in responses:
if isinstance(response, rdf_client.Iterator):
request.iterator = response
status = responses[-1]
self.assertTrue(isinstance(status, rdf_flows.GrrStatus))
if status.status != rdf_flows.GrrStatus.ReturnedStatus.OK:
break
if len(results) > 100:
self.fail("Endless loop detected.")
self.assertIn("Ran out of iterations", status.error_message)
self.assertEqual(grr_worker.suspended_actions, {})
def testSuspendableActionException(self):
class RaisingListDirectory(standard.SuspendableListDirectory):
iterations = 3
def Suspend(self):
RaisingListDirectory.iterations -= 1
if not RaisingListDirectory.iterations:
raise IOError("Ran out of iterations.")
return super(RaisingListDirectory, self).Suspend()
logging.info("The following test is expected to raise a "
"'Ran out of iterations' backtrace.")
p = rdfvalue.PathSpec(path=self.base_path,
pathtype=rdfvalue.PathSpec.PathType.OS)
request = rdfvalue.ListDirRequest(pathspec=p)
request.iterator.number = 2
results = []
grr_worker = worker_mocks.FakeClientWorker()
while request.iterator.state != request.iterator.State.FINISHED:
responses = self.ExecuteAction("RaisingListDirectory",
request,
grr_worker=grr_worker)
results.extend(responses)
for response in responses:
if isinstance(response, rdfvalue.Iterator):
request.iterator = response
status = responses[-1]
self.assertTrue(isinstance(status, rdfvalue.GrrStatus))
if status.status != rdfvalue.GrrStatus.ReturnedStatus.OK:
break
if len(results) > 100:
self.fail("Endless loop detected.")
self.assertIn("Ran out of iterations", status.error_message)
self.assertEqual(grr_worker.suspended_actions, {})
