From ae9e33292f8944ac419ca4b4bd106970677b4512 Mon Sep 17 00:00:00 2001
From: Mike Auty <mike.auty@gmail.com>
Date: Sun, 8 Oct 2017 20:22:10 +0100
Subject: [PATCH] Add in restrictions on automagics should be used for each
 type of plugin.

---
 volatility/cli/__init__.py                 | 20 ++++++++++++++++++++
 volatility/framework/automagic/__init__.py | 12 ++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/volatility/cli/__init__.py b/volatility/cli/__init__.py
index 588e22e9e..0cbc9458b 100644
--- a/volatility/cli/__init__.py
+++ b/volatility/cli/__init__.py
@@ -150,6 +150,9 @@ class CommandLine(object):
                 address, value = extension[:extension.find('=')], json.loads(extension[extension.find('=') + 1:])
                 ctx.config[address] = value
 
+        # It should be up to the UI to determine which automagics to run, so this is before BACK TO THE FRAMEWORK
+        automagics = self.choose_automagic(automagics, plugin)
+
         ###
         # BACK TO THE FRAMEWORK
         ###
@@ -176,6 +179,23 @@ class CommandLine(object):
         # Construct and run the plugin
         text.QuickTextRenderer().render(constructed.run())
 
+    def choose_automagic(self, automagics, plugin):
+        """Chooses which automagics to run, maintaining the order they were handed in"""
+        plugin_category = plugin.__module__.split('.')[2]
+        vollog.info("Detected a {} category plugin".format(plugin_category))
+        output = []
+        for amagic in automagics:
+            if plugin_category == 'windows':
+                if amagic.__class__.__name__ in automagic.windows_automagic:
+                    output += [amagic]
+            elif plugin_category == 'linux':
+                if amagic.__class__.__name__ in automagic.linux_automagic:
+                    output += [amagic]
+            else:
+                return automagics
+        vollog.info("Restricting automagics to: {}".format([x.__class__.__name__ for x in output]))
+        return output
+
     def populate_requirements_argparse(self, parser, configurable):
         """Adds the plugin's simple requirements to the provided parser
 
diff --git a/volatility/framework/automagic/__init__.py b/volatility/framework/automagic/__init__.py
index 175013742..d96dd05be 100644
--- a/volatility/framework/automagic/__init__.py
+++ b/volatility/framework/automagic/__init__.py
@@ -16,6 +16,18 @@ from volatility.framework.configuration import requirements
 
 vollog = logging.getLogger(__name__)
 
+windows_automagic = ['ConstructionMagic',
+                     'LayerStacker',
+                     'NlpDtbfinder',
+                     'WintelHelper',
+                     'KernelPDBScanner']
+
+linux_automagic = ['ConstructionMagic',
+                   'LayerStacker',
+                   'LinuxSymbolCache',
+                   'NlpDtbfinder',
+                   'LinuxSymbolFinder']
+
 
 def available(context):
     """Returns an ordered list of all subclasses of :class:`~volatility.framework.interfaces.automagic.AutomagicInterface`.