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 testResponsesForUnknownRequest(self):
client_id, flow_id = self._SetupClientAndFlow()
request = rdf_flow_objects.FlowRequest(client_id=client_id,
flow_id=flow_id,
request_id=1)
self.db.WriteFlowRequests([request])
# Write two responses at a time, one request exists, the other doesn't.
with test_lib.SuppressLogs():
self.db.WriteFlowResponses([
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=1),
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=2,
response_id=1)
])
# We should have one response in the db.
read = self.db.ReadAllFlowRequestsAndResponses(client_id, flow_id)
self.assertEqual(len(read), 1)
request, responses = read[0]
self.assertEqual(len(responses), 1)
def testInterrogateCloudMetadataLinux(self):
"""Check google cloud metadata on linux."""
client_id, client_urn = self._SetupMinimalClient()
with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.OS,
vfs_test_lib.FakeTestDataVFSHandler):
with test_lib.ConfigOverrider({
"Artifacts.knowledge_base": [
"LinuxWtmp", "NetgroupConfiguration", "LinuxRelease"
],
"Artifacts.netgroup_filter_regexes": [r"^login$"]
}):
client_mock = action_mocks.InterrogatedClient()
client_mock.InitializeClient()
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
discovery.Interrogate.__name__,
client_mock,
token=self.token,
client_id=client_id)
if data_store.RelationalDBReadEnabled():
client = self._OpenClient(client_id)
self._CheckCloudMetadataRelational(client)
else:
client = aff4.FACTORY.Open(client_urn, token=self.token)
self._CheckCloudMetadataAFF4(client)
def testArtifactOutput(self):
"""Check we can run command based artifacts."""
client_id = test_lib.TEST_CLIENT_ID
with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.OS,
vfs_test_lib.FakeTestDataVFSHandler):
# Will raise if something goes wrong.
self.RunCollectorAndGetResults(["TestFilesArtifact"],
client_mock=self.client_mock,
client_id=client_id)
# Will raise if something goes wrong.
self.RunCollectorAndGetResults(["TestFilesArtifact"],
client_mock=self.client_mock,
client_id=client_id,
split_output_by_artifact=True)
# Test the error_on_no_results option.
with self.assertRaises(RuntimeError) as context:
with test_lib.SuppressLogs():
self.RunCollectorAndGetResults(
["NullArtifact"],
client_mock=self.client_mock,
client_id=client_id,
split_output_by_artifact=True,
error_on_no_results=True)
if "collector returned 0 responses" not in str(context.exception):
raise RuntimeError("0 responses should have been returned")
def testKnowledgeBaseRetrievalLinuxPasswd(self):
"""Check we can retrieve a Linux kb."""
with vfs_test_lib.FakeTestDataVFSOverrider():
with test_lib.ConfigOverrider({
"Artifacts.knowledge_base":
["LinuxWtmp", "LinuxPasswdHomedirs", "LinuxReleaseInfo"],
"Artifacts.knowledge_base_additions": [],
"Artifacts.knowledge_base_skip": []
}):
with test_lib.SuppressLogs():
kb = self._RunKBI()
self.assertEqual(kb.os_major_version, 14)
self.assertEqual(kb.os_minor_version, 4)
# user 1,2,3 from wtmp.
self.assertCountEqual([x.username for x in kb.users],
["user1", "user2", "user3"])
user = kb.GetUser(username="******")
self.assertEqual(user.last_logon.AsSecondsSinceEpoch(), 1296552099)
self.assertEqual(user.homedir, "/home/user1")
user = kb.GetUser(username="******")
self.assertEqual(user.last_logon.AsSecondsSinceEpoch(), 1296552102)
self.assertEqual(user.homedir, "/home/user2")
self.assertFalse(kb.GetUser(username="******"))
def testSystemDriveArtifact(self):
client_id = self.SetupClient(0, system="Windows", os_version="6.2")
class BrokenClientMock(action_mocks.ActionMock):
def StatFile(self, _):
raise IOError
def ListDirectory(self, _):
raise IOError
# No registry, broken filesystem, this should just raise.
with self.assertRaises(RuntimeError):
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
collectors.ArtifactCollectorFlow.__name__,
BrokenClientMock(),
artifact_list=["WindowsEnvironmentVariableSystemDrive"],
token=self.token,
client_id=client_id)
# No registry, so this should use the fallback flow
with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.OS,
vfs_test_lib.ClientVFSHandlerFixture):
self._CheckDriveAndRoot()
# Registry is present, so this should use the regular artifact collection
with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.REGISTRY,
vfs_test_lib.FakeRegistryVFSHandler):
self._CheckDriveAndRoot()
def testNoNotificationRescheduling(self):
"""Test that no notifications are rescheduled when a flow raises."""
with test_lib.FakeTime(10000):
flow_obj = self.FlowSetup("RaisingTestFlow")
session_id = flow_obj.session_id
flow_obj.Close()
# Send the flow some messages.
self.SendResponse(session_id, "Hello1", request_id=1)
self.SendResponse(session_id, "Hello2", request_id=2)
self.SendResponse(session_id, "Hello3", request_id=3)
worker_obj = self._TestWorker()
# Process all messages.
with test_lib.SuppressLogs():
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
delay = flow_runner.FlowRunner.notification_retry_interval
with test_lib.FakeTime(10000 + 100 + delay):
manager = queue_manager.QueueManager(token=self.token)
self.assertFalse(
manager.GetNotificationsForAllShards(session_id.Queue()))
def testBrokenFlow(self):
"""Check that flows which call to incorrect states raise."""
client_mock = action_mocks.ActionMock(standard.ReadBuffer)
with self.assertRaises(RuntimeError):
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(flow_test_lib.BrokenFlow.__name__,
client_mock,
client_id=self.client_id,
check_flow_errors=True,
token=self.token)
def testFailedDumpImage(self):
"""Fail to dump flash."""
client_mock = FailDumpMock()
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(hardware.DumpFlashImage.__name__,
client_mock,
client_id=self.client_id,
token=self.token,
check_flow_errors=False)
def testClientFileFinderUploadLimit(self):
paths = [os.path.join(self.base_path, "{**,.}/*.plist")]
action = rdf_file_finder.FileFinderAction.Download()
# TODO(hanuszczak): Instead of catching arbitrary runtime errors, we should
# catch specific instance that was thrown. Unfortunately, all errors are
# intercepted in the `MockWorker` class and converted to runtime errors.
with self.assertRaisesRegexp(RuntimeError, "exceeded network send limit"):
with test_lib.SuppressLogs():
self._RunClientFileFinder(paths, action, network_bytes_limit=1500)
def _Process(self, client_mocks, worker_obj):
while True:
client_msgs_processed = 0
for client_mock in client_mocks:
client_msgs_processed += client_mock.Next()
with test_lib.SuppressLogs():
worker_msgs_processed = worker_obj.RunOnce()
worker_obj.thread_pool.Join()
if not client_msgs_processed and not worker_msgs_processed:
break
def testBrokenChainedFlow(self):
BrokenParentFlow.success = False
# Run the flow in the simulated way
with test_lib.SuppressLogs():
flow_test_lib.StartAndRunFlow(BrokenParentFlow,
client_mock=ClientMock(),
client_id=self.client_id,
check_flow_errors=False)
self.assertEqual(BrokenParentFlow.success, True)
def testCollectNoHashesError(self):
"""Tests exception when no results is returned."""
client_mock = CollectEfiNoHashesMock()
with self.assertRaises(RuntimeError) as err:
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
apple_firmware.CollectEfiHashes.__name__,
client_mock,
client_id=self.client_id,
token=self.token)
self.assertIn("No hash collected.", str(err.exception))
def testCollectHashesError(self):
"""Tests fail collection."""
client_mock = CollectEfiHashesFailMock()
with self.assertRaises(RuntimeError) as err:
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
apple_firmware.CollectEfiHashes.__name__,
client_mock,
client_id=self.client_id,
token=self.token)
self.assertIn("Unable to collect the hashes.", str(err.exception))
def testDumpImageFail(self):
"""Tests EFI Failed dump."""
client_mock = DumpEfiImageFailMock()
with self.assertRaises(RuntimeError) as err:
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
apple_firmware.DumpEfiImage.__name__,
client_mock,
client_id=self.client_id,
token=self.token)
self.assertIn("Unable to dump the flash image", str(err.exception))
def testBrokenChainedFlow(self):
"""Test that exceptions are properly handled in chain flows."""
BrokenParentFlow.success = False
# Run the flow in the simulated way
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper("BrokenParentFlow",
ClientMock(),
client_id=self.client_id,
check_flow_errors=False,
token=self.token)
self.assertEqual(BrokenParentFlow.success, True)
def testResponsesForUnknownFlow(self):
client_id = u"C.1234567890123456"
flow_id = u"1234ABCD"
# This will not raise but also not write anything.
with test_lib.SuppressLogs():
self.db.WriteFlowResponses([
rdf_flow_objects.FlowResponse(client_id=client_id,
flow_id=flow_id,
request_id=1,
response_id=1)
])
read = self.db.ReadAllFlowRequestsAndResponses(client_id, flow_id)
self.assertEqual(read, [])
def testSourceFlowIdIsSet(self):
client_id = self._SetupMinimalClient()
client_mock = action_mocks.InterrogatedClient()
client_mock.InitializeClient()
with test_lib.SuppressLogs():
flow_id = flow_test_lib.TestFlowHelper(
discovery.Interrogate.__name__,
client_mock,
creator=self.test_username,
client_id=client_id)
client = self._OpenClient(client_id)
self.assertNotEmpty(client.metadata.source_flow_id)
self.assertEqual(client.metadata.source_flow_id, flow_id)
def testListDirectoryOnFile(self):
"""OS ListDirectory on a file will raise."""
client_mock = action_mocks.ListDirectoryClientMock()
pb = rdf_paths.PathSpec(path=os.path.join(self.base_path,
"test_img.dd"),
pathtype=rdf_paths.PathSpec.PathType.OS)
# Make sure the flow raises.
with self.assertRaises(RuntimeError):
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(filesystem.ListDirectory.__name__,
client_mock,
client_id=self.client_id,
pathspec=pb,
token=self.token)
def testKnowledgeBaseRetrievalLinuxNoUsers(self):
"""Cause a users.username dependency failure."""
with test_lib.ConfigOverrider({
"Artifacts.knowledge_base": [
"NetgroupConfiguration", "NssCacheLinuxPasswdHomedirs",
"LinuxReleaseInfo"
],
"Artifacts.netgroup_filter_regexes": ["^doesntexist$"]
}):
with vfs_test_lib.FakeTestDataVFSOverrider():
with test_lib.SuppressLogs():
kb = self._RunKBI(require_complete=False)
self.assertEqual(kb.os_major_version, 14)
self.assertEqual(kb.os_minor_version, 4)
self.assertCountEqual([x.username for x in kb.users], [])
def testListDirectoryOnNonexistentDir(self):
"""Test that the ListDirectory flow works."""
client_mock = action_mocks.ListDirectoryClientMock()
pb = rdf_paths.PathSpec(
path=os.path.join(self.base_path, "test_img.dd"),
pathtype=rdf_paths.PathSpec.PathType.OS)
pb.Append(path="doesnotexist", pathtype=rdf_paths.PathSpec.PathType.TSK)
with self.assertRaises(RuntimeError):
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
filesystem.ListDirectory.__name__,
client_mock,
client_id=self.client_id,
pathspec=pb,
token=self.token)
def testNetworkLimitExceeded(self):
"""This tests that the network limit for flows is working."""
client_mock = action_mocks.CPULimitClientMock(user_cpu_usage=[10],
system_cpu_usage=[10],
network_usage=[1000])
with test_lib.SuppressLogs():
flow_id = flow_test_lib.StartAndRunFlow(flow_test_lib.CPULimitFlow,
client_mock=client_mock,
client_id=self.client_id,
cpu_limit=1000,
network_bytes_limit=1500,
check_flow_errors=False)
rdf_flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(rdf_flow.flow_state, "ERROR")
self.assertIn("bytes limit exceeded", rdf_flow.error_message)
def testInterrogateLinuxWithWtmp(self):
"""Test the Interrogate flow."""
self.client_id = self.SetupClient(0,
system="Linux",
os_version="12.04")
data_store.REL_DB.WriteClientMetadata(self.client_id.Basename(),
fleetspeak_enabled=False)
with vfs_test_lib.FakeTestDataVFSOverrider():
with test_lib.ConfigOverrider({
"Artifacts.knowledge_base":
["LinuxWtmp", "NetgroupConfiguration", "LinuxRelease"],
"Artifacts.netgroup_filter_regexes": [r"^login$"]
}):
client_mock = action_mocks.InterrogatedClient()
client_mock.InitializeClient()
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
discovery.Interrogate.__name__,
client_mock,
token=self.token,
client_id=self.client_id)
self.fd = aff4.FACTORY.Open(self.client_id, token=self.token)
self._CheckBasicInfo("test_node.test", "Linux", 100 * 1000000)
self._CheckClientInfo()
self._CheckGRRConfig()
self._CheckNotificationsCreated()
self._CheckClientSummary("Linux",
"14.4",
release="Ubuntu",
kernel="3.13.0-39-generic")
self._CheckRelease("Ubuntu", "14.4")
# users 1,2,3 from wtmp
# users yagharek, isaac from netgroup
self._CheckUsers(["yagharek", "isaac", "user1", "user2", "user3"])
self._CheckNetworkInfo()
self._CheckVFS()
self._CheckLabels()
self._CheckLabelIndex()
self._CheckClientKwIndex(["Linux"], 1)
self._CheckClientKwIndex(["Label2"], 1)
self._CheckClientLibraries()
self._CheckMemory()
def testRuntimeLimitExceeded(self):
client_mock = action_mocks.CPULimitClientMock(
user_cpu_usage=[1],
system_cpu_usage=[1],
network_usage=[1],
runtime_us=[rdfvalue.Duration.From(4, rdfvalue.SECONDS)])
with test_lib.SuppressLogs():
flow_id = flow_test_lib.StartAndRunFlow(
flow_test_lib.CPULimitFlow,
client_mock=client_mock,
client_id=self.client_id,
runtime_limit=rdfvalue.Duration.From(9, rdfvalue.SECONDS),
check_flow_errors=False)
rdf_flow = data_store.REL_DB.ReadFlowObject(self.client_id, flow_id)
self.assertEqual(rdf_flow.flow_state, "ERROR")
self.assertIn("Runtime limit exceeded", rdf_flow.error_message)
def testInterrogateLinuxWithWtmp(self):
"""Test the Interrogate flow."""
client_id = self._SetupMinimalClient()
with vfs_test_lib.FakeTestDataVFSOverrider():
with test_lib.ConfigOverrider({
"Artifacts.knowledge_base": [
"LinuxWtmp", "NetgroupConfiguration", "LinuxReleaseInfo"
],
"Artifacts.netgroup_filter_regexes": [r"^login$"]
}):
client_mock = action_mocks.InterrogatedClient()
client_mock.InitializeClient(version="14.4", release="Ubuntu")
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
discovery.Interrogate.__name__,
client_mock,
token=self.token,
client_id=client_id)
client = self._OpenClient(client_id)
self._CheckBasicInfo(client, "test_node.test", "Linux", 100 * 1000000)
self._CheckClientInfo(client)
self._CheckGRRConfig(client)
self._CheckNotificationsCreated(self.token.username, client_id)
self._CheckClientSummary(
client_id,
client.GetSummary(),
"Linux",
"14.4",
release="Ubuntu",
kernel="3.13.0-39-generic")
self._CheckRelease(client, "Ubuntu", "14.4")
# users 1,2,3 from wtmp, users yagharek, isaac from netgroup
self._CheckUsers(client, ["yagharek", "isaac", "user1", "user2", "user3"])
self._CheckNetworkInfo(client)
# No VFS test when running on the relational db.
self._CheckLabels(client_id)
self._CheckLabelIndex(client_id)
self._CheckClientKwIndex(["Linux"], 1)
self._CheckClientKwIndex(["Label2"], 1)
self._CheckClientLibraries(client)
self._CheckMemory(client)
def _RunTSKFileFinder(self, paths):
image_path = os.path.join(self.base_path, "ntfs_img.dd")
with utils.Stubber(
vfs, "_VFS_VIRTUALROOTS", {
rdf_paths.PathSpec.PathType.TSK:
rdf_paths.PathSpec(
path=image_path, pathtype="OS", offset=63 * 512)
}):
action = rdf_file_finder.FileFinderAction.Action.DOWNLOAD
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
file_finder.FileFinder.__name__,
self.client_mock,
client_id=self.client_id,
paths=paths,
pathtype=rdf_paths.PathSpec.PathType.TSK,
action=rdf_file_finder.FileFinderAction(action_type=action),
creator=self.test_username)
def testStuckFlowGetsTerminated(self):
worker_obj = self._TestWorker()
initial_time = rdfvalue.RDFDatetime.FromSecondsSinceEpoch(100)
with test_lib.SuppressLogs():
try:
with test_lib.FakeTime(initial_time.AsSecondsSinceEpoch()):
session_id = flow.StartAFF4Flow(
flow_name=WorkerStuckableTestFlow.__name__,
client_id=self.client_id,
token=self.token,
sync=False)
# Process all messages
while worker_obj.RunOnce():
pass
# Wait until worker thread starts processing the flow.
WorkerStuckableTestFlow.WaitUntilWorkerStartsProcessing()
# Set the time to max worker flow duration + 1 minute. The flow is
# currently blocked because of the way semaphores are set up.
# Worker should consider the flow to be stuck and terminate it.
stuck_flows_timeout = flow_runner.FlowRunner.stuck_flows_timeout
future_time = (initial_time + rdfvalue.Duration("1m") +
stuck_flows_timeout)
with test_lib.FakeTime(future_time.AsSecondsSinceEpoch()):
worker_obj.RunOnce()
finally:
# Release the semaphore so that worker thread unblocks and finishes
# processing the flow.
WorkerStuckableTestFlow.StopFlow()
WorkerStuckableTestFlow.LetWorkerFinishProcessing()
worker_obj.thread_pool.Join()
killed_flow = aff4.FACTORY.Open(session_id, token=self.token)
self.assertEqual(killed_flow.context.state,
rdf_flow_runner.FlowContext.State.ERROR)
self.assertEqual(
killed_flow.context.status,
"Terminated by user test. Reason: Stuck in the worker")
def testInterrogateCloudMetadataLinux(self):
"""Check google cloud metadata on linux."""
client_id = self._SetupMinimalClient()
with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.OS,
vfs_test_lib.FakeTestDataVFSHandler):
with test_lib.ConfigOverrider({
"Artifacts.knowledge_base":
["LinuxWtmp", "NetgroupConfiguration", "LinuxReleaseInfo"],
"Artifacts.netgroup_filter_regexes": [r"^login$"]
}):
client_mock = action_mocks.InterrogatedClient()
client_mock.InitializeClient()
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
discovery.Interrogate.__name__,
client_mock,
creator=self.test_username,
client_id=client_id)
client = self._OpenClient(client_id)
self._CheckCloudMetadata(client)
def testKnowledgeBaseRetrievalLinux(self):
"""Check we can retrieve a Linux kb."""
with test_lib.ConfigOverrider({
"Artifacts.knowledge_base": [
"LinuxWtmp", "NetgroupConfiguration",
"LinuxPasswdHomedirs", "LinuxReleaseInfo"
],
"Artifacts.netgroup_filter_regexes": ["^login$"],
"Artifacts.netgroup_user_blacklist": ["isaac"]
}):
with vfs_test_lib.FakeTestDataVFSOverrider():
with test_lib.SuppressLogs():
kb = self._RunKBI()
self.assertEqual(kb.os_major_version, 14)
self.assertEqual(kb.os_minor_version, 4)
# user 1,2,3 from wtmp. yagharek from netgroup.
self.assertCountEqual([x.username for x in kb.users],
["user1", "user2", "user3", "yagharek"])
user = kb.GetUser(username="******")
self.assertEqual(user.last_logon.AsSecondsSinceEpoch(), 1296552099)
self.assertEqual(user.homedir, "/home/user1")
def testInterrogateCloudMetadataLinux(self):
"""Check google cloud metadata on linux."""
self.client_id = self.SetupClient(0,
system="Linux",
os_version="12.04")
with vfs_test_lib.VFSOverrider(rdf_paths.PathSpec.PathType.OS,
vfs_test_lib.FakeTestDataVFSHandler):
with test_lib.ConfigOverrider({
"Artifacts.knowledge_base":
["LinuxWtmp", "NetgroupConfiguration", "LinuxRelease"],
"Artifacts.netgroup_filter_regexes": [r"^login$"]
}):
client_mock = action_mocks.InterrogatedClient()
client_mock.InitializeClient()
with test_lib.SuppressLogs():
flow_test_lib.TestFlowHelper(
discovery.Interrogate.__name__,
client_mock,
token=self.token,
client_id=self.client_id)
self.fd = aff4.FACTORY.Open(self.client_id, token=self.token)
self._CheckCloudMetadata()
