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