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 Execute(self, action_cls, args): responses = list() def SendReply(value, session_id=None, message_type=rdf_flows.GrrMessage.Type.MESSAGE): if message_type != rdf_flows.GrrMessage.Type.MESSAGE: return if str(session_id) in self.wkfs: message = rdf_flows.GrrMessage(name=action_cls.__name__, payload=value, auth_state="AUTHENTICATED", session_id=session_id) self.wkfs[str(session_id)].ProcessMessage(message) else: responses.append(value) message = rdf_flows.GrrMessage(name=action_cls.__name__, payload=args, auth_state="AUTHENTICATED", session_id=rdfvalue.SessionID()) action = action_cls(grr_worker=worker_mocks.FakeClientWorker()) action.SendReply = SendReply action.Execute(message) return responses
def _RunClientFileFinder(self, paths, action, network_bytes_limit=None, client_id=None): client_id = client_id or self.SetupClient(0) with test_lib.ConfigOverrider({"Client.server_urls": [self.base_url]}): client = comms.GRRHTTPClient( ca_cert=config.CONFIG["CA.certificate"], private_key=config.CONFIG.Get("Client.private_key", default=None), worker_cls=worker_mocks.DisabledNannyClientWorker) client.client_worker = worker_mocks.FakeClientWorker(client=client) client.server_certificate = config.CONFIG["Frontend.certificate"] for s in flow_test_lib.TestFlowHelper( file_finder.ClientFileFinder.__name__, action_mocks.ClientFileFinderClientMock( client_worker=client.client_worker), client_id=client_id, paths=paths, pathtype=rdf_paths.PathSpec.PathType.OS, action=action, process_non_regular_files=True, network_bytes_limit=network_bytes_limit, token=self.token): session_id = s return session_id
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 = collections.defaultdict(lambda: 0) self.recorded_args = {} self._recorded_messages = [] self.client_worker = (kwargs.get("client_worker", None) or worker_mocks.FakeClientWorker())
def _RunClientFileFinder(self, paths, action, network_bytes_limit=None, client_id=None): client_id = client_id or self.SetupClient(0) with test_lib.ConfigOverrider({"Client.server_urls": [self.base_url]}): session_id = flow_test_lib.TestFlowHelper( file_finder.ClientFileFinder.__name__, action_mocks.ClientFileFinderClientMock( client_worker=worker_mocks.FakeClientWorker()), client_id=client_id, paths=paths, pathtype=rdf_paths.PathSpec.PathType.OS, action=action, process_non_regular_files=True, network_bytes_limit=network_bytes_limit, token=self.token) return session_id
def _RunAction(self): fake_worker = worker_mocks.FakeClientWorker() self.RunAction(admin.SendStartupInfo, grr_worker=fake_worker) return [m.payload for m in fake_worker.responses]