Add ARTIFACT_GROUP support for image_export.py

plaso/engine/artifact_filters.py

@@ -43,6 +43,9 @@ def __init__(self, artifacts_registry, artifact_filters, knowledge_base):
     self._artifacts = artifact_filters
     self._artifacts_registry = artifacts_registry
     self._knowledge_base = knowledge_base
+    self.find_specs_per_source_type = {
+        artifact_types.TYPE_INDICATOR_FILE: [],
+        artifact_types.TYPE_INDICATOR_WINDOWS_REGISTRY_KEY: []}
 
   @staticmethod
   def CheckKeyCompatibility(key_path):
@@ -73,10 +76,6 @@ def BuildFindSpecs(self, environment_variables=None):
       environment_variables (Optional[list[EnvironmentVariableArtifact]]):
           environment variables.
     """
-    find_specs_per_source_type = {
-        artifact_types.TYPE_INDICATOR_FILE: [],
-        artifact_types.TYPE_INDICATOR_WINDOWS_REGISTRY_KEY: []}
-
     for name in self._artifacts:
       definition = self._artifacts_registry.GetDefinitionByName(name)
       if not definition:
@@ -90,7 +89,7 @@ def BuildFindSpecs(self, environment_variables=None):
             find_specs = self.BuildFindSpecsFromFileArtifact(
                 path_entry, source.separator, environment_variables,
                 self._knowledge_base.user_accounts)
-            find_specs_per_source_type[
+            self.find_specs_per_source_type[
                 artifact_types.TYPE_INDICATOR_FILE].extend(find_specs)
 
         elif (source.type_indicator ==
@@ -100,7 +99,7 @@ def BuildFindSpecs(self, environment_variables=None):
           for key_path in set(source.keys):
             if self.CheckKeyCompatibility(key_path):
               find_specs = self.BuildFindSpecsFromRegistryArtifact(key_path)
-              find_specs_per_source_type[
+              self.find_specs_per_source_type[
                   artifact_types.TYPE_INDICATOR_WINDOWS_REGISTRY_KEY].extend(
                       find_specs)
 
@@ -118,17 +117,24 @@ def BuildFindSpecs(self, environment_variables=None):
               key_value['key'] for key_value in source.key_value_pairs]):
             if self.CheckKeyCompatibility(key_path):
               find_specs = self.BuildFindSpecsFromRegistryArtifact(key_path)
-              find_specs_per_source_type[
+              self.find_specs_per_source_type[
                   artifact_types.TYPE_INDICATOR_WINDOWS_REGISTRY_KEY].extend(
                       find_specs)
 
+        elif (source.type_indicator ==
+              artifact_types.TYPE_INDICATOR_ARTIFACT_GROUP):
+          self._artifacts.remove(name)
+          for name_entry in set(source.names):
+            self._artifacts.append(name_entry)
+          self.BuildFindSpecs(environment_variables=environment_variables)
+
         else:
           logger.warning(
               'Unsupported artifact definition source type: "{0:s}"'.format(
                   source.type_indicator))
 
     self._knowledge_base.SetValue(
-        self.KNOWLEDGE_BASE_VALUE, find_specs_per_source_type)
+        self.KNOWLEDGE_BASE_VALUE, self.find_specs_per_source_type)
 
   def BuildFindSpecsFromFileArtifact(
       self, source_path, path_separator, environment_variables, user_accounts):

test_data/artifacts/artifacts_filters.yaml

@@ -1,5 +1,16 @@
 # Artifact definitions.
 
+name: TestGroupExport
+doc: Test Group
+sources:
+- type: ARTIFACT_GROUP
+  attributes:
+    names:
+     - 'TestFiles3'
+     - 'TestFiles4'
+labels: [System]
+supported_os: [Windows]
+---
 name: TestFiles
 doc: Test Doc
 sources:
@@ -26,6 +37,30 @@ sources:
 labels: [System]
 supported_os: [Windows]
 ---
+name: TestFiles3
+doc: Test Doc3
+sources:
+- type: FILE
+  attributes:
+    paths:
+      - '\a_directory\'
+      - '\a_directory\a_file'
+    separator: '\'
+labels: [System]
+supported_os: [Windows]
+---
+name: TestFiles4
+doc: Test Doc4
+sources:
+- type: FILE
+  attributes:
+    paths:
+      - '\a_directory\another_file'
+      - '\passwords.txt'
+    separator: '\'
+labels: [System]
+supported_os: [Windows]
+---
 name: TestRegistry
 doc: Test Registry Doc
 sources:

tests/cli/image_export_tool.py

@@ -411,7 +411,7 @@ def testProcessSourcesExtractWithArtifactsFilter(self):
     options.image = self._GetTestFilePath(['image.qcow2'])
     options.quiet = True
     options.artifact_filters = 'TestFilesImageExport'
-
+    
     with shared_test_lib.TempDirectory() as temp_directory:
       options.path = temp_directory
 
@@ -427,6 +427,36 @@ def testProcessSourcesExtractWithArtifactsFilter(self):
       extracted_files = self._RecursiveList(temp_directory)
 
     self.assertEqual(sorted(extracted_files), expected_extracted_files)
+
+  @shared_test_lib.skipUnlessHasTestFile(['artifacts'])
+  @shared_test_lib.skipUnlessHasTestFile(['image.qcow2'])
+  def testProcessSourcesExtractWithArtifactsGroupFilter(self):
+    """Tests the ProcessSources function with a filter file."""
+    output_writer = test_lib.TestOutputWriter(encoding='utf-8')
+    test_tool = image_export_tool.ImageExportTool(output_writer=output_writer)
+
+    options = test_lib.TestOptions()
+    options.artifact_definitions_path = self._GetTestFilePath(['artifacts'])
+    options.image = self._GetTestFilePath(['image.qcow2'])
+    options.quiet = True
+    options.artifact_filters = 'TestGroupExport'
+
+    with shared_test_lib.TempDirectory() as temp_directory:
+      options.path = temp_directory
+
+      test_tool.ParseOptions(options)
+
+      test_tool.ProcessSources()
+
+      expected_extracted_files = sorted([
+          os.path.join(temp_directory, 'a_directory'),
+          os.path.join(temp_directory, 'a_directory', 'another_file'),
+          os.path.join(temp_directory, 'a_directory', 'a_file'),
+          os.path.join(temp_directory, 'passwords.txt')])
+
+      extracted_files = self._RecursiveList(temp_directory)
+
+    self.assertEqual(sorted(extracted_files), expected_extracted_files)
 
   @shared_test_lib.skipUnlessHasTestFile(['syslog_image.dd'])
   def testProcessSourcesExtractWithSignaturesFilter(self):