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):
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"))
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)
# 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 CreateDefaultArtifactRegistry():
r = artifact_registry.ArtifactRegistry()
r.AddDefaultSources()
return r
def CreateTestArtifactRegistry():
r = artifact_registry.ArtifactRegistry()
r.AddDirSource(os.path.join(config.CONFIG["Test.data_dir"], "artifacts"))
r.AddDatastoreSources([aff4.ROOT_URN.Add("artifact_store")])
return r
def CreateDatastoreOnlyArtifactRegistry():
r = artifact_registry.ArtifactRegistry()
r.AddDatastoreSources([aff4.ROOT_URN.Add("artifact_store")])
return r