def testNewArtifactLoaded(self):
"""Simulate a new artifact being loaded into the store via the UI."""
cmd_artifact = """name: "TestCmdArtifact"
doc: "Test command artifact for dpkg."
sources:
- type: "COMMAND"
attributes:
cmd: "/usr/bin/dpkg"
args: ["--list"]
labels: [ "Software" ]
supported_os: [ "Linux" ]
"""
no_datastore_artifact = """name: "NotInDatastore"
doc: "Test command artifact for dpkg."
sources:
- type: "COMMAND"
attributes:
cmd: "/usr/bin/dpkg"
args: ["--list"]
labels: [ "Software" ]
supported_os: [ "Linux" ]
"""
test_registry = artifact_registry.ArtifactRegistry()
test_registry.ClearRegistry()
test_registry.AddDatastoreSource(rdfvalue.RDFURN("aff4:/artifact_store"))
test_registry._dirty = False
with utils.Stubber(artifact_registry, "REGISTRY", test_registry):
with self.assertRaises(rdf_artifacts.ArtifactNotRegisteredError):
artifact_registry.REGISTRY.GetArtifact("TestCmdArtifact")
with self.assertRaises(rdf_artifacts.ArtifactNotRegisteredError):
artifact_registry.REGISTRY.GetArtifact("NotInDatastore")
# Add artifact to datastore but not registry
artifact_coll = artifact_registry.ArtifactCollection(
rdfvalue.RDFURN("aff4:/artifact_store"))
with data_store.DB.GetMutationPool() as pool:
for artifact_val in artifact_registry.REGISTRY.ArtifactsFromYaml(
cmd_artifact):
artifact_coll.Add(artifact_val, mutation_pool=pool)
# Add artifact to registry but not datastore
for artifact_val in artifact_registry.REGISTRY.ArtifactsFromYaml(
no_datastore_artifact):
artifact_registry.REGISTRY.RegisterArtifact(
artifact_val, source="datastore", overwrite_if_exists=False)
# We need to reload all artifacts from the data store before trying to get
# the artifact.
artifact_registry.REGISTRY.ReloadDatastoreArtifacts()
self.assertTrue(artifact_registry.REGISTRY.GetArtifact("TestCmdArtifact"))
# We registered this artifact with datastore source but didn't
# write it into aff4. This simulates an artifact that was
# uploaded in the UI then later deleted. We expect it to get
# cleared when the artifacts are reloaded from the datastore.
with self.assertRaises(rdf_artifacts.ArtifactNotRegisteredError):
artifact_registry.REGISTRY.GetArtifact("NotInDatastore")
def testSystemArtifactOverwrite(self):
content = """
name: WMIActiveScriptEventConsumer
doc: here's the doc
sources:
- type: COMMAND
attributes:
args: ["-L", "-v", "-n"]
cmd: /sbin/iptables
supported_os: [Linux]
"""
artifact_registry.REGISTRY.ClearRegistry()
artifact_registry.REGISTRY.ClearSources()
artifact_registry.REGISTRY._CheckDirty()
# System artifacts come from the test file.
self.LoadTestArtifacts()
# Uploaded files go to this collection.
artifact_store_urn = aff4.ROOT_URN.Add("artifact_store")
artifact_registry.REGISTRY.AddDatastoreSources([artifact_store_urn])
# WMIActiveScriptEventConsumer is a system artifact, we can't overwrite it.
with self.assertRaises(rdf_artifacts.ArtifactDefinitionError):
artifact.UploadArtifactYamlFile(content)
# Override the check and upload anyways. This simulates the case
# where an artifact ends up shadowing a system artifact somehow -
# for example when the system artifact was created after the
# artifact was uploaded to the data store for testing.
artifact.UploadArtifactYamlFile(content,
overwrite_system_artifacts=True)
# The shadowing artifact is at this point stored in the
# collection. On the next full reload of the registry, there will
# be an error that we can't overwrite the system artifact. The
# artifact should automatically get deleted from the collection to
# mitigate the problem.
with self.assertRaises(rdf_artifacts.ArtifactDefinitionError):
artifact_registry.REGISTRY._ReloadArtifacts()
# As stated above, now this should work.
artifact_registry.REGISTRY._ReloadArtifacts()
# Make sure the artifact is now loaded and it's the version from the file.
self.assertIn("WMIActiveScriptEventConsumer",
artifact_registry.REGISTRY._artifacts)
artifact_obj = artifact_registry.REGISTRY.GetArtifact(
"WMIActiveScriptEventConsumer")
self.assertStartsWith(artifact_obj.loaded_from, "file:")
# The artifact is gone from the collection.
coll = artifact_registry.ArtifactCollection(artifact_store_urn)
self.assertNotIn("WMIActiveScriptEventConsumer", coll)
def UploadArtifactYamlFile(file_content,
overwrite=True,
overwrite_system_artifacts=False):
"""Upload a yaml or json file as an artifact to the datastore."""
loaded_artifacts = []
registry_obj = artifact_registry.REGISTRY
# Make sure all artifacts are loaded so we don't accidentally overwrite one.
registry_obj.GetArtifacts(reload_datastore_artifacts=True)
new_artifacts = registry_obj.ArtifactsFromYaml(file_content)
new_artifact_names = set()
# A quick syntax check before we upload anything.
for artifact_value in new_artifacts:
artifact_registry.ValidateSyntax(artifact_value)
new_artifact_names.add(artifact_value.name)
# Iterate through each artifact adding it to the collection.
artifact_coll = artifact_registry.ArtifactCollection(
ARTIFACT_STORE_ROOT_URN)
current_artifacts = list(artifact_coll)
# We need to remove artifacts we are overwriting.
filtered_artifacts = [
art for art in current_artifacts if art.name not in new_artifact_names
]
artifact_coll.Delete()
with data_store.DB.GetMutationPool() as pool:
for artifact_value in filtered_artifacts:
artifact_coll.Add(artifact_value, mutation_pool=pool)
for artifact_value in new_artifacts:
registry_obj.RegisterArtifact(
artifact_value,
source="datastore:%s" % ARTIFACT_STORE_ROOT_URN,
overwrite_if_exists=overwrite,
overwrite_system_artifacts=overwrite_system_artifacts)
artifact_coll.Add(artifact_value, mutation_pool=pool)
if data_store.RelationalDBWriteEnabled():
data_store.REL_DB.WriteArtifact(artifact_value)
loaded_artifacts.append(artifact_value)
name = artifact_value.name
logging.info("Uploaded artifact %s to %s", name,
ARTIFACT_STORE_ROOT_URN)
# Once all artifacts are loaded we can validate dependencies. Note that we do
# not have to perform a syntax validation because it is already done after
# YAML is parsed.
for artifact_value in loaded_artifacts:
artifact_registry.ValidateDependencies(artifact_value)
def testArtifactUpload(self):
file_path = os.path.join(self.base_path, "hello.exe")
artifact_source = """
name: ArtifactFromSource
doc: My first artifact.
labels:
- Logs
- Authentication
supported_os:
- Linux
sources:
- type: FILE
attributes:
paths:
- %s
""" % file_path
artifact_obj = artifact_registry.REGISTRY.ArtifactsFromYaml(
artifact_source)[0]
artifact_registry.REGISTRY.AddDatastoreSources(
[artifact.ARTIFACT_STORE_ROOT_URN])
artifact_registry.REGISTRY._CheckDirty()
if data_store.RelationalDBEnabled():
data_store.REL_DB.WriteArtifact(artifact_obj)
else:
with data_store.DB.GetMutationPool() as pool:
artifact_coll = artifact_registry.ArtifactCollection(
artifact.ARTIFACT_STORE_ROOT_URN)
artifact_coll.Delete()
artifact_coll.Add(artifact_obj, mutation_pool=pool)
# Make sure that the artifact is not yet registered and the flow will have
# to read it from the data store.
with self.assertRaises(rdf_artifacts.ArtifactNotRegisteredError):
artifact_registry.REGISTRY.GetArtifact("ArtifactFromSource")
self._GetArtifact("ArtifactFromSource")
