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(artifact_registry.ArtifactDefinitionError):
artifact.UploadArtifactYamlFile(content, token=self.token)
# 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,
token=self.token)
# 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(artifact_registry.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.assertTrue(artifact_obj.loaded_from.startswith("file:"))
# The artifact is gone from the collection.
coll = artifact_registry.ArtifactCollection(artifact_store_urn,
token=self.token)
self.assertNotIn("WMIActiveScriptEventConsumer", coll)
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._dirty = False
with utils.Stubber(artifact_registry, "REGISTRY", test_registry):
collect_flow = collectors.ArtifactCollectorFlow(None, token=self.token)
with self.assertRaises(artifact_registry.ArtifactNotRegisteredError):
artifact_registry.REGISTRY.GetArtifact("TestCmdArtifact")
with self.assertRaises(artifact_registry.ArtifactNotRegisteredError):
artifact_registry.REGISTRY.GetArtifact("NotInDatastore")
# Add artifact to datastore but not registry
artifact_coll = artifact_registry.ArtifactCollection(
rdfvalue.RDFURN("aff4:/artifact_store"), token=self.token)
for artifact_val in artifact_registry.REGISTRY.ArtifactsFromYaml(
cmd_artifact):
artifact_coll.Add(artifact_val)
# 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)
# This should succeeded because the artifacts will be reloaded from the
# datastore.
self.assertTrue(collect_flow._GetArtifactFromName("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(artifact_registry.ArtifactNotRegisteredError):
artifact_registry.REGISTRY.GetArtifact("NotInDatastore")
def UploadArtifactYamlFile(file_content,
base_urn=None,
token=None,
overwrite=True,
overwrite_system_artifacts=False):
"""Upload a yaml or json file as an artifact to the datastore."""
loaded_artifacts = []
if not base_urn:
base_urn = aff4.ROOT_URN.Add("artifact_store")
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_value.ValidateSyntax()
new_artifact_names.add(artifact_value.name)
# Iterate through each artifact adding it to the collection.
artifact_coll = artifact_registry.ArtifactCollection(base_urn, token=token)
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()
for artifact_value in filtered_artifacts:
artifact_coll.Add(artifact_value)
for artifact_value in new_artifacts:
registry_obj.RegisterArtifact(
artifact_value,
source="datastore:%s" % base_urn,
overwrite_if_exists=overwrite,
overwrite_system_artifacts=overwrite_system_artifacts)
artifact_coll.Add(artifact_value)
loaded_artifacts.append(artifact_value)
logging.info("Uploaded artifact %s to %s", artifact_value.name,
base_urn)
# Once all artifacts are loaded we can validate, as validation of dependencies
# requires the group are all loaded before doing the validation.
for artifact_value in loaded_artifacts:
artifact_value.Validate()
return base_urn