def testDelayedCallState(self):
"""Tests the ability to delay a CallState invocation."""
with test_lib.FakeTime(10000):
client_mock = ClientMock()
client_mock = test_lib.MockClient(self.client_id, client_mock,
token=self.token)
worker_mock = test_lib.MockWorker(check_flow_errors=True,
token=self.token)
flow.GRRFlow.StartFlow(
client_id=self.client_id, flow_name="DelayedCallStateFlow",
token=self.token)
self.Work(client_mock, worker_mock)
# We should have done the first CallState so far.
self.assertEqual(DelayedCallStateFlow.flow_ran, 1)
with test_lib.FakeTime(10050):
# 50 seconds more is not enough.
self.Work(client_mock, worker_mock)
self.assertEqual(DelayedCallStateFlow.flow_ran, 1)
with test_lib.FakeTime(10100):
# But 100 is.
self.Work(client_mock, worker_mock)
self.assertEqual(DelayedCallStateFlow.flow_ran, 2)
def Run(self):
with test_lib.FakeTime(42):
flow_urn = flow.GRRFlow.StartFlow(
flow_name=processes.ListProcesses.__name__,
client_id=self.client_id,
token=self.token)
mock = test_lib.MockClient(self.client_id, None, token=self.token)
while mock.Next():
pass
replace = {flow_urn.Basename(): "W:ABCDEF"}
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)] = "42"
self.Check("ListFlowRequests",
args=flow_plugin.ApiListFlowRequestsArgs(
client_id=self.client_id.Basename(),
flow_id=flow_urn.Basename()),
replace=replace)
def Run(self):
client_ids = self.SetupClients(1)
client_id = client_ids[0]
replace = {}
with test_lib.FakeTime(42):
flow_urn = flow.GRRFlow.StartFlow(
client_id=client_id,
flow_name=processes.ListProcesses.__name__,
token=self.token)
replace[flow_urn.Basename()] = "F:123456"
# Here we emulate a mock client with no actions (None) that should produce
# an error.
mock = 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 testCPULimitForFlows(self):
"""This tests that the client actions are limited properly."""
result = {}
client_mock = CPULimitClientMock(result)
client_mock = test_lib.MockClient(self.client_id,
client_mock,
token=self.token)
client_mock.EnableResourceUsage(user_cpu_usage=[10],
system_cpu_usage=[10],
network_usage=[1000])
worker_obj = worker.GRRWorker(token=self.token)
flow.GRRFlow.StartFlow(client_id=self.client_id,
flow_name="CPULimitFlow",
cpu_limit=1000,
network_bytes_limit=10000,
token=self.token)
self._Process([client_mock], worker_obj)
self.assertEqual(result["cpulimit"], [1000, 980, 960])
self.assertEqual(result["networklimit"], [10000, 9000, 8000])
return result
def setUp(self):
super(GRRFuseTest, self).setUp()
self.client_id = self.SetupClients(1)[0]
self.client_name = str(self.client_id)[len("aff4:/"):]
with aff4.FACTORY.Open(self.client_id, token=self.token,
mode="rw") as fd:
fd.Set(fd.Schema.SYSTEM("Linux"))
kb = fd.Schema.KNOWLEDGE_BASE()
fd.Set(kb)
with aff4.FACTORY.Create(self.client_id.Add("fs/os"),
aff4_standard.VFSDirectory,
mode="rw",
token=self.token) as fd:
fd.Set(fd.Schema.PATHSPEC(path="/", pathtype="OS"))
# Ignore cache so our tests always get client side updates.
self.grr_fuse = fuse_mount.GRRFuse(root="/",
token=self.token,
ignore_cache=True)
self.action_mock = action_mocks.ActionMock(
admin.GetClientInfo,
admin.GetConfiguration,
admin.GetPlatformInfo,
file_fingerprint.FingerprintFile,
linux.EnumerateFilesystems,
linux.EnumerateInterfaces,
linux.EnumerateUsers,
linux.GetInstallDate,
searching.Find,
standard.HashBuffer,
standard.ListDirectory,
standard.StatFile,
standard.TransferBuffer,
)
self.client_mock = test_lib.MockClient(self.client_id,
self.action_mock,
token=self.token)
self.update_stubber = utils.Stubber(self.grr_fuse,
"_RunAndWaitForVFSFileUpdate",
self._RunAndWaitForVFSFileUpdate)
self.update_stubber.Start()
self.start_flow_stubber = utils.Stubber(flow_utils, "StartFlowAndWait",
self.StartFlowAndWait)
self.start_flow_stubber.Start()
def testWorkerPrioritization(self):
"""Test that flow priorities work on the worker side."""
result = []
client_mock = PriorityClientMock(result)
client_mock = test_lib.MockClient(self.client_id,
client_mock,
token=self.token)
worker_mock = test_lib.MockWorker(check_flow_errors=True,
token=self.token)
# Start some flows with different priorities.
# pyformat: disable
args = [
(rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority"),
(rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority"),
(rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority2"),
(rdf_flows.GrrMessage.Priority.HIGH_PRIORITY, "high priority"),
(rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority2")
]
# pyformat: enable
server_result = []
PriorityFlow.storage = server_result
for (priority, msg) in args:
flow.GRRFlow.StartFlow(client_id=self.client_id,
flow_name="PriorityFlow",
msg=msg,
priority=priority,
token=self.token)
while True:
# Run all the clients first so workers have messages to choose from.
client_processed = 1
while client_processed:
client_processed = client_mock.Next()
# Now process the results, this should happen in the correct order.
flows_run = []
for flow_run in worker_mock.Next():
flows_run.append(flow_run)
if not flows_run:
break
# The flows should be run in order of priority.
self.assertEqual(server_result[0:1], [u"high priority"])
self.assertEqual(sorted(server_result[1:3]),
[u"medium priority", u"medium priority2"])
self.assertEqual(sorted(server_result[3:5]),
[u"low priority", u"low priority2"])
def testClientPrioritization(self):
"""Test that flow priorities work on the client side."""
result = []
client_mock = PriorityClientMock(result)
client_mock = test_lib.MockClient(self.client_id,
client_mock,
token=self.token)
worker_mock = test_lib.MockWorker(check_flow_errors=True,
token=self.token)
# Start some flows with different priorities.
# pyformat: disable
args = [
(rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority"),
(rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority"),
(rdf_flows.GrrMessage.Priority.LOW_PRIORITY, "low priority2"),
(rdf_flows.GrrMessage.Priority.HIGH_PRIORITY, "high priority"),
(rdf_flows.GrrMessage.Priority.MEDIUM_PRIORITY, "medium priority2")
]
# pyformat: enable
for (priority, msg) in args:
flow.GRRFlow.StartFlow(client_id=self.client_id,
flow_name="PriorityFlow",
msg=msg,
priority=priority,
token=self.token)
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
# The flows should be run in order of priority.
self.assertEqual(result[0:1], [u"high priority"])
self.assertEqual(sorted(result[1:3]),
[u"medium priority", u"medium priority2"])
self.assertEqual(sorted(result[3:5]),
[u"low priority", u"low priority2"])
def __init__(self, method_name=None):
super(GRRFuseTest, self).__init__(method_name)
# Set up just once for the whole test suite, since we don't have any
# per-test setup to do.
super(GRRFuseTest, self).setUp()
self.client_name = str(self.client_id)[len("aff4:/"):]
with aff4.FACTORY.Open(self.client_id, token=self.token, mode="rw") as fd:
fd.Set(fd.Schema.SYSTEM("Linux"))
kb = fd.Schema.KNOWLEDGE_BASE()
fd.Set(kb)
# Ignore cache so our tests always get client side updates.
self.grr_fuse = fuse_mount.GRRFuse(root="/", token=self.token,
ignore_cache=True)
self.action_mock = test_lib.ActionMock("TransferBuffer", "StatFile", "Find",
"HashFile", "HashBuffer",
"UpdateVFSFile",
"EnumerateInterfaces",
"EnumerateFilesystems",
"GetConfiguration",
"GetConfig", "GetClientInfo",
"GetInstallDate", "GetPlatformInfo",
"EnumerateUsers", "ListDirectory")
client_mock = test_lib.MockClient(self.client_id, self.action_mock,
token=self.token)
worker_mock = test_lib.MockWorker(check_flow_errors=True, token=self.token)
# All the flows we've run so far. We'll check them for errors at the end of
# each test.
self.total_flows = set()
# We add the thread as a class variable since we'll be referring to it in
# the tearDownClass method, and we want all tests to share it.
self.__class__.fake_server_thread = threading.Thread(
target=self.RunFakeWorkerAndClient,
args=(client_mock,
worker_mock))
self.fake_server_thread.start()
def RunFlow(self, flow_name, **kwargs):
result = {}
client_mock = CPULimitClientMock(result)
client_mock = test_lib.MockClient(self.client_id, client_mock,
token=self.token)
worker_mock = ResourcedWorker(check_flow_errors=True,
token=self.token)
flow.GRRFlow.StartFlow(
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 testInspect(self):
"""Test the inspect UI."""
client_id = self.SetupClients(1)[0]
self.RequestAndGrantClientApproval(client_id)
flow.GRRFlow.StartFlow(client_id=client_id,
flow_name=flow_discovery.Interrogate.__name__,
token=self.token)
mock = 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 testInspect(self):
"""Test the inspect UI."""
with self.ACLChecksDisabled():
self.GrantClientApproval("C.0000000000000001")
self.Open("/")
self.Type("client_query", "0001")
self.Click("client_query_submit")
self.WaitUntilEqual(u"C.0000000000000001", self.GetText,
"css=span[type=subject]")
# Choose client 1
self.Click("css=td:contains('0001')")
self.Click("css=a[grrtarget=LaunchFlows]")
self.Click("css=#_Administrative ins")
self.Click("css=a:contains(Interrogate)")
self.Click("css=button.Launch")
# Open the "Advanced" dropdown.
self.Click("css=a[href='#HostAdvanced']")
# Click on the "Debug client requests".
self.Click("css=a[grrtarget=DebugClientRequestsView]")
self.WaitUntil(self.IsElementPresent,
"css=td:contains(GetPlatformInfo)")
# Check that the we can see the requests in the table.
for request in "GetPlatformInfo GetConfig EnumerateInterfaces".split():
self.assertTrue(
self.IsElementPresent("css=td:contains(%s)" % request))
self.Click("css=td:contains(GetPlatformInfo)")
# Check that the proto is rendered inside the tab.
self.WaitUntil(
self.IsElementPresent,
"css=.tab-content td.proto_value:contains(GetPlatformInfo)")
# Check that the request tab is currently selected.
self.assertTrue(
self.IsElementPresent("css=li.active:contains(Request)"))
# Here we emulate a mock client with no actions (None) this should produce
# an error.
with self.ACLChecksDisabled():
mock = test_lib.MockClient(
rdfvalue.ClientURN("C.0000000000000001"),
None,
token=self.token)
while mock.Next():
pass
# Now select the Responses tab:
self.Click("css=li a:contains(Responses)")
self.WaitUntil(self.IsElementPresent,
"css=td:contains('flow:response:')")
self.assertTrue(
self.IsElementPresent(
"css=.tab-content td.proto_value:contains(GENERIC_ERROR)"))
self.assertTrue(
self.IsElementPresent(
"css=.tab-content td.proto_value:contains(STATUS)"))
def testCPULimitForHunts(self):
worker_obj = worker.GRRWorker(token=self.token)
client_ids = ["C.%016X" % i for i in xrange(10, 20)]
result = {}
client_mocks = []
for client_id in client_ids:
client_mock = CPULimitClientMock(result)
client_mock = test_lib.MockClient(rdf_client.ClientURN(client_id),
client_mock,
token=self.token)
client_mock.EnableResourceUsage(user_cpu_usage=[10],
system_cpu_usage=[10],
network_usage=[1000])
client_mocks.append(client_mock)
flow_runner_args = rdf_flows.FlowRunnerArgs(flow_name="CPULimitFlow")
with hunts.GRRHunt.StartHunt(hunt_name="GenericHunt",
flow_runner_args=flow_runner_args,
cpu_limit=5000,
network_bytes_limit=1000000,
client_rate=0,
token=self.token) as hunt:
hunt.GetRunner().Start()
hunts.GRRHunt.StartClients(hunt.session_id, client_ids[:1])
self._Process(client_mocks, worker_obj)
hunts.GRRHunt.StartClients(hunt.session_id, client_ids[1:2])
self._Process(client_mocks, worker_obj)
hunts.GRRHunt.StartClients(hunt.session_id, client_ids[2:3])
self._Process(client_mocks, worker_obj)
# The limiting factor here is the overall hunt limit of 5000 cpu
# seconds. Clients that finish should decrease the remaining quota
# and the following clients should get the reduced quota.
self.assertEqual(result["cpulimit"], [
5000.0, 4980.0, 4960.0, 4940.0, 4920.0, 4900.0, 4880.0, 4860.0,
4840.0
])
self.assertEqual(result["networklimit"], [
1000000L, 999000L, 998000L, 997000L, 996000L, 995000L, 994000L,
993000L, 992000L
])
result.clear()
with hunts.GRRHunt.StartHunt(hunt_name="GenericHunt",
flow_runner_args=flow_runner_args,
per_client_cpu_limit=3000,
per_client_network_limit_bytes=3000000,
client_rate=0,
token=self.token) as hunt:
hunt.GetRunner().Start()
hunts.GRRHunt.StartClients(hunt.session_id, client_ids[:1])
self._Process(client_mocks, worker_obj)
hunts.GRRHunt.StartClients(hunt.session_id, client_ids[1:2])
self._Process(client_mocks, worker_obj)
hunts.GRRHunt.StartClients(hunt.session_id, client_ids[2:3])
self._Process(client_mocks, worker_obj)
# This time, the per client limit is 3000s / 3000000 bytes. Every
# client should get the same limit.
self.assertEqual(result["cpulimit"], [
3000.0, 2980.0, 2960.0, 3000.0, 2980.0, 2960.0, 3000.0, 2980.0,
2960.0
])
self.assertEqual(result["networklimit"], [
3000000, 2999000, 2998000, 3000000, 2999000, 2998000, 3000000,
2999000, 2998000
])
result.clear()
for client_mock in client_mocks:
client_mock.EnableResourceUsage(user_cpu_usage=[500],
system_cpu_usage=[500],
network_usage=[1000000])
with hunts.GRRHunt.StartHunt(hunt_name="GenericHunt",
flow_runner_args=flow_runner_args,
per_client_cpu_limit=3000,
cpu_limit=5000,
per_client_network_limit_bytes=3000000,
network_bytes_limit=5000000,
client_rate=0,
token=self.token) as hunt:
hunt.GetRunner().Start()
hunts.GRRHunt.StartClients(hunt.session_id, client_ids[:1])
self._Process(client_mocks, worker_obj)
hunts.GRRHunt.StartClients(hunt.session_id, client_ids[1:2])
self._Process(client_mocks, worker_obj)
hunts.GRRHunt.StartClients(hunt.session_id, client_ids[2:3])
self._Process(client_mocks, worker_obj)
# The first client gets the full per client limit of 3000s, and
# uses all of it. The hunt has a limit of just 5000 total so the
# second client gets started with a limit of 2000. It can only run
# two of three states, the last client will not be started at all
# due to out of quota.
self.assertEqual(result["cpulimit"],
[3000.0, 2000.0, 1000.0, 2000.0, 1000.0])
self.assertEqual(result["networklimit"],
[3000000, 2000000, 1000000, 2000000, 1000000])
errors = list(hunt.GetClientsErrors())
self.assertEqual(len(errors), 2)
# Client side out of cpu.
self.assertIn("CPU limit exceeded", errors[0].log_message)
# Server side out of cpu.
self.assertIn("Out of CPU quota", errors[1].backtrace)
