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]
