def testIteratedListDirectory(self):
"""Tests iterated listing of directories."""
p = rdfvalue.PathSpec(path=self.base_path,
pathtype=rdfvalue.PathSpec.PathType.OS)
non_iterated_results = self.RunAction(
"ListDirectory", rdfvalue.ListDirRequest(pathspec=p))
# Make sure we get some results.
l = len(non_iterated_results)
self.assertTrue(l > 0)
iterated_results = []
request = rdfvalue.ListDirRequest(pathspec=p)
request.iterator.number = 2
while True:
responses = self.RunAction("IteratedListDirectory", request)
results = responses[:-1]
if not results: break
for result in results:
iterated_results.append(result)
for x, y in zip(non_iterated_results, iterated_results):
# Reset the st_atime in the results to avoid potential flakiness.
x.st_atime = y.st_atime = 0
self.assertRDFValueEqual(x, y)
def testSuspendableListDirectory(self):
request = rdfvalue.ListDirRequest()
request.pathspec.path = self.base_path
request.pathspec.pathtype = "OS"
request.iterator.number = 2
results = []
grr_worker = worker_mocks.FakeClientWorker()
while request.iterator.state != request.iterator.State.FINISHED:
responses = self.RunAction("SuspendableListDirectory",
request,
grr_worker=grr_worker)
results.extend(responses)
for response in responses:
if isinstance(response, rdfvalue.Iterator):
request.iterator = response
filenames = [
os.path.basename(r.pathspec.path) for r in results
if isinstance(r, rdfvalue.StatEntry)
]
self.assertItemsEqual(filenames, os.listdir(self.base_path))
iterators = [r for r in results if isinstance(r, rdfvalue.Iterator)]
# One for two files plus one extra with the FINISHED status.
nr_files = len(os.listdir(self.base_path))
expected_iterators = (nr_files / 2) + 1
if nr_files % 2:
expected_iterators += 1
self.assertEqual(len(iterators), expected_iterators)
# Make sure the thread has been deleted.
self.assertEqual(grr_worker.suspended_actions, {})
def Start(self):
"""Get information about the file from the client."""
self.state.Register("max_chunk_number",
max(2, self.args.read_length/self.CHUNK_SIZE))
self.state.Register("current_chunk_number", 0)
self.state.Register("file_size", 0)
self.state.Register("fd", None)
self.state.Register("stat", None)
self.CallClient("StatFile", rdfvalue.ListDirRequest(
pathspec=self.args.pathspec), next_state="Stat")
def testHashFile(self):
"""Can we hash a file?"""
path = os.path.join(self.base_path, "morenumbers.txt")
p = rdfvalue.PathSpec(path=path,
pathtype=rdfvalue.PathSpec.PathType.OS)
# The action returns a DataBlob object.
result = self.RunAction("HashFile",
rdfvalue.ListDirRequest(
pathspec=p))[0]
self.assertEqual(result.data,
hashlib.sha256(open(path).read()).digest())
def Start(self):
"""Issue a request to list the directory."""
self.state.Register("responses", [])
self.state.Register("urn", None)
# We use data to pass the path to the callback:
self.state.Register("request", rdfvalue.ListDirRequest(
pathspec=self.state.args.pathspec))
# For this example we will use a really small number to force many round
# trips with the client. This is a performance killer.
self.state.request.iterator.number = 50
self.CallClient("IteratedListDirectory", self.state.request,
next_state="List")
def testListDirectory(self):
"""Tests listing directories."""
p = rdfvalue.PathSpec(path=self.base_path, pathtype=0)
results = self.RunAction("ListDirectory",
rdfvalue.ListDirRequest(pathspec=p))
# Find the number.txt file
result = None
for result in results:
if os.path.basename(result.pathspec.path) == "morenumbers.txt":
break
self.assert_(result)
self.assertEqual(result.__class__, rdfvalue.StatEntry)
self.assertEqual(result.pathspec.Basename(), "morenumbers.txt")
self.assertEqual(result.st_size, 3893)
self.assert_(stat.S_ISREG(result.st_mode))
def testClientAction(self):
client_mock = test_lib.ActionMock("ListDirectory")
pathspec = rdfvalue.PathSpec(path=os.path.join(self.base_path,
"test_img.dd"),
pathtype=rdfvalue.PathSpec.PathType.OS)
request = rdfvalue.ListDirRequest(pathspec=pathspec)
for _ in test_lib.TestFlowHelper("ClientAction",
client_mock,
client_id=self.client_id,
action="ListDirectory",
break_pdb=False,
action_args=request,
token=self.token):
pass
def testSuspendableActionException(self):
class RaisingListDirectory(standard.SuspendableListDirectory):
iterations = 3
def Suspend(self):
RaisingListDirectory.iterations -= 1
if not RaisingListDirectory.iterations:
raise IOError("Ran out of iterations.")
return super(RaisingListDirectory, self).Suspend()
logging.info("The following test is expected to raise a "
"'Ran out of iterations' backtrace.")
p = rdfvalue.PathSpec(path=self.base_path,
pathtype=rdfvalue.PathSpec.PathType.OS)
request = rdfvalue.ListDirRequest(pathspec=p)
request.iterator.number = 2
results = []
grr_worker = worker_mocks.FakeClientWorker()
while request.iterator.state != request.iterator.State.FINISHED:
responses = self.ExecuteAction("RaisingListDirectory",
request,
grr_worker=grr_worker)
results.extend(responses)
for response in responses:
if isinstance(response, rdfvalue.Iterator):
request.iterator = response
status = responses[-1]
self.assertTrue(isinstance(status, rdfvalue.GrrStatus))
if status.status != rdfvalue.GrrStatus.ReturnedStatus.OK:
break
if len(results) > 100:
self.fail("Endless loop detected.")
self.assertIn("Ran out of iterations", status.error_message)
self.assertEqual(grr_worker.suspended_actions, {})
def _ProcessResponse(self, response, component_paths, base_wildcard=False):
for component_path in component_paths:
regexes_to_get = []
recursions_to_get = {}
node = self.FindNode(component_path)
if not node:
# Node is empty representing a leaf node - we found a hit - report it.
self.GlobReportMatch(response)
return
# There are further components in the tree - iterate over them.
for component_str, next_node in node.items():
component = rdfvalue.PathSpec(component_str)
next_component = component_path + [component_str]
# If we reach this point, we are instructed to go deeper into the
# directory structure. We only want to actually do this if
# - the last response was a proper directory,
# - or it was a file (an image) that was explicitly given meaning
# no wildcards or groupings,
# - or no_file_type_check was set.
#
# This reduces the number of TSK opens on the client that may
# sometimes lead to instabilities due to bugs in the library.
if response and (
not (stat.S_ISDIR(response.st_mode) or
not base_wildcard or
self.state.no_file_type_check)):
continue
if component.path_options == component.Options.RECURSIVE:
recursions_to_get.setdefault(
component.recursion_depth, []).append(component)
elif component.path_options == component.Options.REGEX:
regexes_to_get.append(component)
elif component.path_options == component.Options.CASE_INSENSITIVE:
# Here we need to create the next pathspec by appending the current
# component to what we already have. If we don't have anything yet, we
# fall back to the root path. If there is no root path either, the
# current component becomes the new base.
base_pathspec = self._GetBasePathspec(response)
if base_pathspec:
pathspec = base_pathspec.Append(component)
else:
pathspec = component
if not next_node:
# Check for the existence of the last node.
request = rdfvalue.ListDirRequest(pathspec=pathspec)
if (response is None or
(response and
(response.st_mode == 0 or
not stat.S_ISREG(response.st_mode)))):
# If next node is empty, this node is a leaf node, we therefore
# must stat it to check that it is there. There is a special case
# here where this pathspec points to a file/directory in the root
# directory. In this case, response will be None but we still need
# to stat it.
self.CallClient(
"StatFile", request, next_state="ProcessEntry",
request_data=dict(component_path=next_component))
else:
# There is no need to go back to the client for intermediate
# paths in the prefix tree, just emulate this by recursively
# calling this state inline.
self.CallStateInline(
[rdfvalue.StatEntry(pathspec=pathspec)],
next_state="ProcessEntry",
request_data=dict(component_path=next_component))
if recursions_to_get or regexes_to_get:
# Recursions or regexes need a base pathspec to operate on. If we
# have neither a response or a root path, we send a default pathspec
# that opens the root with pathtype "OS".
base_pathspec = self._GetBasePathspec(response)
if not base_pathspec:
base_pathspec = rdfvalue.PathSpec(path="/", pathtype="OS")
for depth, recursions in recursions_to_get.iteritems():
path_regex = "(?i)^" + "$|^".join(
set([c.path for c in recursions])) + "$"
findspec = rdfvalue.FindSpec(pathspec=base_pathspec,
cross_devs=True,
max_depth=depth,
path_regex=path_regex)
findspec.iterator.number = self.FILE_MAX_PER_DIR
self.CallClient("Find", findspec,
next_state="ProcessEntry",
request_data=dict(base_path=component_path))
if regexes_to_get:
path_regex = "(?i)^" + "$|^".join(
set([c.path for c in regexes_to_get])) + "$"
findspec = rdfvalue.FindSpec(pathspec=base_pathspec,
max_depth=1,
path_regex=path_regex)
findspec.iterator.number = self.FILE_MAX_PER_DIR
self.CallClient("Find", findspec,
next_state="ProcessEntry",
request_data=dict(base_path=component_path))
def ProcessEntry(self, responses):
"""Process the responses from the client."""
if not responses.success:
return
component_path = responses.request_data["component_path"]
node = self.FindNode(component_path)
# If we get a response with an unfinished iterator then we missed some
# files. Call Find on the client until we're done.
if (responses.iterator and
responses.iterator.state != responses.iterator.State.FINISHED):
findspec = rdfvalue.FindSpec(responses.request.request.args)
findspec.iterator = responses.iterator
self.CallClient("Find",
findspec,
next_state="ProcessEntry",
request_data=responses.request_data)
regexes_to_get = []
recursions_to_get = {}
for response in responses:
# The Find client action does not return a StatEntry but a
# FindSpec. Normalize to a StatEntry.
if isinstance(response, rdfvalue.FindSpec):
response = response.hit
if node:
# There are further components in the tree - iterate over them.
for component_str, next_node in node.items():
component = rdfvalue.PathSpec(component_str)
next_component = component_path + [component_str]
# Use the pathtype from the flow args.
component.pathtype = self.state.args.pathtype
if component.path_options == component.Options.RECURSIVE:
recursions_to_get.setdefault(component.recursion_depth,
[]).append(component)
elif component.path_options == component.Options.REGEX:
regexes_to_get.append(component)
elif component.path_options == component.Options.CASE_INSENSITIVE:
# Check for the existence of the last node.
if not next_node:
pathspec = response.pathspec.Copy().AppendPath(
component.path)
request = rdfvalue.ListDirRequest(
pathspec=pathspec)
if response.st_mode == 0 or not stat.S_ISREG(
response.st_mode):
# If next node is empty, this node is a leaf node, we therefore
# must stat it to check that it is there.
self.CallClient(
"StatFile",
request,
next_state="ProcessEntry",
request_data=dict(
component_path=next_component))
else:
pathspec = response.pathspec.Copy().AppendPath(
component.path)
# There is no need to go back to the client for intermediate paths
# in the prefix tree, just emulate this by recursively calling
# this state inline.
self.CallStateInline(
[rdfvalue.StatEntry(pathspec=pathspec)],
next_state="ProcessEntry",
request_data=dict(
component_path=next_component))
if recursions_to_get:
for depth, recursions in recursions_to_get.iteritems():
path_regex = "(?i)^" + "$|^".join(
set([c.path for c in recursions])) + "$"
findspec = rdfvalue.FindSpec(
pathspec=response.pathspec,
cross_devs=True,
max_depth=depth,
path_regex=path_regex)
findspec.iterator.number = self.FILE_MAX_PER_DIR
self.CallClient(
"Find",
findspec,
next_state="ProcessEntry",
request_data=dict(component_path=next_component))
if regexes_to_get:
path_regex = "(?i)^" + "$|^".join(
set([c.path for c in regexes_to_get])) + "$"
findspec = rdfvalue.FindSpec(pathspec=response.pathspec,
max_depth=1,
path_regex=path_regex)
findspec.iterator.number = self.FILE_MAX_PER_DIR
self.CallClient(
"Find",
findspec,
next_state="ProcessEntry",
request_data=dict(component_path=next_component))
else:
# Node is empty representing a leaf node - we found a hit - report it.
self.ReportMatch(response)
