def testPendingFlowTermination(self):
client_mock = ClientMock()
flow_id = flow.StartFlow(flow_cls=ParentFlow, client_id=self.client_id)
flow_obj = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(flow_obj.flow_state, "RUNNING")
pending_termination = rdf_flow_objects.PendingFlowTermination(
reason="testing")
data_store.REL_DB.UpdateFlow(self.client_id,
flow_id,
pending_termination=pending_termination)
with flow_test_lib.TestWorker() as worker:
with test_lib.SuppressLogs():
flow_test_lib.RunFlow(self.client_id,
flow_id,
client_mock=client_mock,
worker=worker,
check_flow_errors=False)
flow_obj = data_store.REL_DB.ReadFlowObject(
self.client_id, flow_id)
self.assertEqual(flow_obj.flow_state, "ERROR")
self.assertEqual(flow_obj.error_message, "testing")
def testGetFileIndicatesCollectedSizeAfterCollection(self):
# Find the file with FileFinder stat action, so that we can reference it
# and trigger "Collect" operation on it.
client_ref = self.api.Client(client_id=self.client_id)
args = rdf_file_finder.FileFinderArgs(
paths=[os.path.join(self.base_path, "numbers.txt")],
action=rdf_file_finder.FileFinderAction.Stat()).AsPrimitiveProto()
client_ref.CreateFlow(name=file_finder.FileFinder.__name__, args=args)
client_mock = action_mocks.FileFinderClientMock()
flow_test_lib.FinishAllFlowsOnClient(
self.client_id, client_mock=client_mock)
f = client_ref.File("fs/os" + os.path.join(self.base_path, "numbers.txt"))
with flow_test_lib.TestWorker():
operation = f.Collect()
self.assertEqual(operation.GetState(), operation.STATE_RUNNING)
flow_test_lib.FinishAllFlowsOnClient(
self.client_id, client_mock=client_mock)
self.assertEqual(operation.GetState(), operation.STATE_FINISHED)
f = f.Get()
self.assertNotEqual(f.data.hash.sha256, b"")
self.assertGreater(f.data.hash.num_bytes, 0)
self.assertGreater(f.data.last_collected, 0)
self.assertGreater(f.data.last_collected_size, 0)
def testCreatorPropagation(self):
username = u"original user"
data_store.REL_DB.WriteGRRUser(username)
client_mock = ClientMock()
flow_id = flow.StartFlow(flow_cls=ParentFlow,
client_id=self.client_id,
creator=username)
worker = flow_test_lib.TestWorker(token=True)
data_store.REL_DB.RegisterFlowProcessingHandler(worker.ProcessFlow)
flow_test_lib.RunFlow(self.client_id,
flow_id,
client_mock=client_mock,
worker=worker)
flow_obj = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(flow_obj.creator, username)
child_flows = data_store.REL_DB.ReadChildFlowObjects(
self.client_id, flow_id)
self.assertEqual(len(child_flows), 1)
child_flow = child_flows[0]
self.assertEqual(child_flow.creator, username)
def TestHuntHelperWithMultipleMocks(client_mocks,
check_flow_errors=False,
token=None,
iteration_limit=None):
"""Runs a hunt with a given set of clients mocks.
Args:
client_mocks: Dictionary of (client_id->client_mock) pairs. Client mock
objects are used to handle client actions. Methods names of a client mock
object correspond to client actions names. For an example of a client mock
object, see SampleHuntMock.
check_flow_errors: If True, raises when one of hunt-initiated flows fails.
token: An instance of access_control.ACLToken security token.
iteration_limit: If None, hunt will run until it's finished. Otherwise,
worker_mock.Next() will be called iteration_limit number of times. Every
iteration processes worker's message queue. If new messages are sent to
the queue during the iteration processing, they will be processed on next
iteration,
"""
total_flows = set()
# Worker always runs with absolute privileges, therefore making the token
# SetUID().
token = token.SetUID()
client_mocks = [
flow_test_lib.MockClient(client_id, client_mock, token=token)
for client_id, client_mock in iteritems(client_mocks)
]
with flow_test_lib.TestWorker(token=True) as rel_db_worker:
worker_mock = worker_test_lib.MockWorker(
check_flow_errors=check_flow_errors, token=token)
# Run the clients and worker until nothing changes any more.
while iteration_limit is None or iteration_limit > 0:
client_processed = 0
for client_mock in client_mocks:
client_processed += client_mock.Next()
flows_run = []
for flow_run in worker_mock.Next():
total_flows.add(flow_run)
flows_run.append(flow_run)
worker_processed = rel_db_worker.ResetProcessedFlows()
flows_run.extend(worker_processed)
if client_processed == 0 and not flows_run:
break
if iteration_limit:
iteration_limit -= 1
if check_flow_errors:
flow_test_lib.CheckFlowErrors(total_flows, token=token)
def testRaisesIfFlowProcessingRequestDoesNotTriggerAnyProcessing(self):
with flow_test_lib.TestWorker() as worker:
flow_id = flow.StartFlow(flow_cls=CallClientParentFlow,
client_id=self.client_id)
fpr = rdf_flows.FlowProcessingRequest(client_id=self.client_id,
flow_id=flow_id)
with self.assertRaises(worker_lib.FlowHasNothingToProcessError):
worker.ProcessFlow(fpr)
def RunHuntWithClientCrashes(self, client_ids):
with flow_test_lib.TestWorker() as test_worker:
client_mocks = dict([(client_id,
flow_test_lib.CrashClientMock(client_id))
for client_id in client_ids])
self.AssignTasksToClients(client_ids=client_ids,
worker=test_worker)
return TestHuntHelperWithMultipleMocks(client_mocks)
def RunHuntWithClientCrashes(self, client_ids):
with flow_test_lib.TestWorker(threadpool_size=0, token=True) as test_worker:
client_mocks = dict([
(client_id, flow_test_lib.CrashClientMock(client_id, self.token))
for client_id in client_ids
])
self.AssignTasksToClients(client_ids=client_ids, worker=test_worker)
return TestHuntHelperWithMultipleMocks(client_mocks, False, self.token)
def TestHuntHelperWithMultipleMocks(client_mocks,
iteration_limit=None,
worker=None):
"""Runs a hunt with a given set of clients mocks.
Args:
client_mocks: Dictionary of (client_id->client_mock) pairs. Client mock
objects are used to handle client actions. Methods names of a client mock
object correspond to client actions names. For an example of a client mock
object, see SampleHuntMock.
iteration_limit: If None, hunt will run until it's finished. Otherwise,
worker_mock.Next() will be called iteration_limit number of times. Every
iteration processes worker's message queue. If new messages are sent to
the queue during the iteration processing, they will be processed on next
iteration.
worker: flow_test_lib.TestWorker object to use.
Returns:
A number of iterations complete.
"""
client_mocks = [
flow_test_lib.MockClient(client_id, client_mock)
for client_id, client_mock in client_mocks.items()
]
if worker is None:
rel_db_worker = flow_test_lib.TestWorker()
data_store.REL_DB.RegisterFlowProcessingHandler(
rel_db_worker.ProcessFlow)
else:
rel_db_worker = worker
num_iterations = 0
try:
# Run the clients and worker until nothing changes any more.
while iteration_limit is None or num_iterations < iteration_limit:
data_store.REL_DB.delegate.WaitUntilNoFlowsToProcess(timeout=10)
worker_processed = rel_db_worker.ResetProcessedFlows()
client_processed = 0
for client_mock in client_mocks:
client_processed += int(client_mock.Next())
num_iterations += 1
if client_processed == 0 and not worker_processed:
break
finally:
if worker is None:
data_store.REL_DB.UnregisterFlowProcessingHandler(timeout=60)
rel_db_worker.Shutdown()
return num_iterations
def RunHunt(self,
client_ids=None,
iteration_limit=None,
client_mock=None,
**mock_kwargs):
with flow_test_lib.TestWorker() as test_worker:
self.AssignTasksToClients(client_ids=client_ids,
worker=test_worker)
if client_mock is None:
client_mock = SampleHuntMock(**mock_kwargs)
return TestHuntHelper(client_mock,
client_ids or self.client_ids,
iteration_limit=iteration_limit,
worker=test_worker)
def AssignTasksToClients(self, client_ids=None, worker=None):
# Pretend to be the foreman now and dish out hunting jobs to all the
# clients..
client_ids = client_ids or self.client_ids
foreman_obj = foreman.Foreman()
def Assign():
for client_id in client_ids:
foreman_obj.AssignTasksToClient(client_id)
if worker is None:
with flow_test_lib.TestWorker():
Assign()
else:
Assign()
def testRefreshRecursivelyWaitUntilDone(self):
f = self.api.Client(client_id=self.client_id).File("fs/os/c/Downloads")
with flow_test_lib.TestWorker():
operation = f.RefreshRecursively(max_depth=5)
self.assertEqual(operation.GetState(), operation.STATE_RUNNING)
def ProcessOperation():
time.sleep(1)
flow_test_lib.FinishAllFlowsOnClient(self.client_id)
threading.Thread(target=ProcessOperation).start()
result_f = operation.WaitUntilDone().target_file
self.assertEqual(f.path, result_f.path)
self.assertEqual(operation.GetState(), operation.STATE_FINISHED)
def AssignTasksToClients(self, client_ids=None, worker=None):
# Pretend to be the foreman now and dish out hunting jobs to all the
# clients..
client_ids = client_ids or self.client_ids
foreman_obj = foreman.GetForeman(token=self.token)
def Assign():
for client_id in client_ids:
foreman_obj.AssignTasksToClient(
rdf_client.ClientURN(client_id).Basename())
if worker is None:
with flow_test_lib.TestWorker(threadpool_size=0, token=True):
Assign()
else:
Assign()
def testCollectWaitUntilDone(self):
f = self.api.Client(
client_id=self.client_urn.Basename()).File("fs/os/c/Downloads/a.txt")
with flow_test_lib.TestWorker(token=self.token):
operation = f.Collect()
self.assertEqual(operation.GetState(), operation.STATE_RUNNING)
def ProcessOperation():
time.sleep(1)
flow_test_lib.FinishAllFlowsOnClient(self.client_urn.Basename())
threading.Thread(target=ProcessOperation).start()
result_f = operation.WaitUntilDone().target_file
self.assertEqual(f.path, result_f.path)
self.assertEqual(operation.GetState(), operation.STATE_FINISHED)
def RunHunt(self,
client_ids=None,
iteration_limit=None,
client_mock=None,
**mock_kwargs):
with flow_test_lib.TestWorker(threadpool_size=0, token=True) as test_worker:
self.AssignTasksToClients(client_ids=client_ids, worker=test_worker)
if client_mock is None:
client_mock = SampleHuntMock(**mock_kwargs)
return TestHuntHelper(
client_mock,
client_ids or self.client_ids,
check_flow_errors=False,
iteration_limit=iteration_limit,
worker=test_worker,
token=self.token)
def TestHuntHelperWithMultipleMocks(client_mocks,
check_flow_errors=False,
token=None,
iteration_limit=None,
worker=None):
"""Runs a hunt with a given set of clients mocks.
Args:
client_mocks: Dictionary of (client_id->client_mock) pairs. Client mock
objects are used to handle client actions. Methods names of a client mock
object correspond to client actions names. For an example of a client mock
object, see SampleHuntMock.
check_flow_errors: If True, raises when one of hunt-initiated flows fails.
token: An instance of access_control.ACLToken security token.
iteration_limit: If None, hunt will run until it's finished. Otherwise,
worker_mock.Next() will be called iteration_limit number of times. Every
iteration processes worker's message queue. If new messages are sent to
the queue during the iteration processing, they will be processed on next
iteration.
worker: flow_test_lib.TestWorker object to use.
Returns:
A number of iterations complete.
"""
if token is None:
token = access_control.ACLToken(username="******")
total_flows = set()
# Worker always runs with absolute privileges, therefore making the token
# SetUID().
token = token.SetUID()
client_mocks = [
flow_test_lib.MockClient(client_id, client_mock, token=token)
for client_id, client_mock in iteritems(client_mocks)
]
if worker is None:
rel_db_worker = flow_test_lib.TestWorker(threadpool_size=0, token=True)
data_store.REL_DB.RegisterFlowProcessingHandler(rel_db_worker.ProcessFlow)
else:
rel_db_worker = worker
num_iterations = 0
try:
worker_mock = worker_test_lib.MockWorker(
check_flow_errors=check_flow_errors, token=token)
# Run the clients and worker until nothing changes any more.
while iteration_limit is None or num_iterations < iteration_limit:
worker_processed = []
if data_store.RelationalDBEnabled():
data_store.REL_DB.delegate.WaitUntilNoFlowsToProcess(timeout=10)
worker_processed = rel_db_worker.ResetProcessedFlows()
client_processed = 0
for client_mock in client_mocks:
client_processed += client_mock.Next()
flows_run = []
for flow_run in worker_mock.Next():
total_flows.add(flow_run)
flows_run.append(flow_run)
flows_run.extend(worker_processed)
num_iterations += 1
if client_processed == 0 and not flows_run and not worker_processed:
break
if check_flow_errors:
flow_test_lib.CheckFlowErrors(total_flows, token=token)
finally:
if worker is None:
data_store.REL_DB.UnregisterFlowProcessingHandler(timeout=60)
rel_db_worker.Shutdown()
return num_iterations
