def setUp(self):
super(ArtifactFilesDownloaderFlowTest, self).setUp()
with aff4.FACTORY.Open(self.client_id, token=self.token,
mode="rw") as fd:
fd.Set(fd.Schema.SYSTEM("Windows"))
kb = fd.Schema.KNOWLEDGE_BASE()
artifact.SetCoreGRRKnowledgeBaseValues(kb, fd)
fd.Set(kb)
self.output_path = "analysis/artifact_files_downloader"
self.stubbers = []
self.collector_replies = []
def ArtifactCollectorStub(this):
for r in self.collector_replies:
this.SendReply(r)
stubber = utils.Stubber(collectors.ArtifactCollectorFlow, "Start",
ArtifactCollectorStub)
stubber.Start()
self.stubbers.append(stubber)
self.start_file_fetch_args = []
self.received_files = []
self.failed_files = []
def StartFileFetch(this, pathspec, request_data=None):
self.start_file_fetch_args.append(pathspec)
for r in self.received_files:
this.ReceiveFetchedFile(r, None, request_data=request_data)
for r in self.failed_files:
this.FileFetchFailed(pathspec,
"StatFile",
request_data=request_data)
stubber = utils.Stubber(transfer.MultiGetFileMixin, "StartFileFetch",
StartFileFetch)
stubber.Start()
self.stubbers.append(stubber)
def setUp(self):
"""Make sure things are initialized."""
super(TestArtifactCollectors, self).setUp()
test_artifacts_file = os.path.join(config_lib.CONFIG["Test.data_dir"],
"artifacts", "test_artifacts.json")
artifact_registry.REGISTRY.AddFileSource(test_artifacts_file)
self.fakeartifact = artifact_registry.REGISTRY.GetArtifact(
"FakeArtifact")
self.fakeartifact2 = artifact_registry.REGISTRY.GetArtifact(
"FakeArtifact2")
self.output_count = 0
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()
artifact.SetCoreGRRKnowledgeBaseValues(kb, fd)
fd.Set(kb)
def UpdateCoreKBAttributes(self):
fd = aff4.FACTORY.Open(self.client_id, token=self.token, mode="rw")
kb = fd.Get(fd.Schema.KNOWLEDGE_BASE)
artifact.SetCoreGRRKnowledgeBaseValues(kb, fd)
fd.Set(fd.Schema.KNOWLEDGE_BASE, kb)
fd.Flush()
def CreateClientObject(self, vfs_fixture):
"""Make a new client object."""
# First remove the old fixture just in case its still there.
aff4.FACTORY.Delete(self.client_id, token=self.token)
# Create the fixture at a fixed time.
with test_lib.FakeTime(self.age):
for path, (aff4_type, attributes) in vfs_fixture:
path %= self.args
aff4_object = aff4.FACTORY.Create(self.client_id.Add(path),
aff4_type,
mode="rw",
token=self.token)
for attribute_name, value in attributes.items():
attribute = aff4.Attribute.PREDICATES[attribute_name]
if isinstance(value, (str, unicode)):
# Interpolate the value
value %= self.args
# Is this supposed to be an RDFValue array?
if aff4.issubclass(attribute.attribute_type,
rdf_protodict.RDFValueArray):
rdfvalue_object = attribute()
for item in value:
new_object = rdfvalue_object.rdf_type.FromTextFormat(
utils.SmartStr(item))
rdfvalue_object.Append(new_object)
# It is a text serialized protobuf.
elif aff4.issubclass(attribute.attribute_type,
rdf_structs.RDFProtoStruct):
# Use the alternate constructor - we always write protobufs in
# textual form:
rdfvalue_object = attribute.attribute_type.FromTextFormat(
utils.SmartStr(value))
elif aff4.issubclass(attribute.attribute_type,
rdfvalue.RDFInteger):
rdfvalue_object = attribute(int(value))
else:
rdfvalue_object = attribute(value)
# If we don't already have a pathspec, try and get one from the stat.
if aff4_object.Get(aff4_object.Schema.PATHSPEC) is None:
# If the attribute was a stat, it has a pathspec nested in it.
# We should add that pathspec as an attribute.
if attribute.attribute_type == rdf_client.StatEntry:
stat_object = attribute.attribute_type.FromTextFormat(
utils.SmartStr(value))
if stat_object.pathspec:
pathspec_attribute = aff4.Attribute(
"aff4:pathspec", rdf_paths.PathSpec,
"The pathspec used to retrieve "
"this object from the client.", "pathspec")
aff4_object.AddAttribute(
pathspec_attribute, stat_object.pathspec)
if attribute in ["aff4:content", "aff4:content"]:
# For AFF4MemoryStreams we need to call Write() instead of
# directly setting the contents..
aff4_object.Write(rdfvalue_object)
else:
aff4_object.AddAttribute(attribute, rdfvalue_object)
# Populate the KB from the client attributes.
if aff4_type == aff4_grr.VFSGRRClient:
kb = rdf_client.KnowledgeBase()
artifact.SetCoreGRRKnowledgeBaseValues(kb, aff4_object)
aff4_object.Set(aff4_object.Schema.KNOWLEDGE_BASE, kb)
# Make sure we do not actually close the object here - we only want to
# sync back its attributes, not run any finalization code.
aff4_object.Flush()
if aff4_type == aff4_grr.VFSGRRClient:
index = client_index.CreateClientIndex(token=self.token)
index.AddClient(aff4_object)
