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, {})