Merge pull request #790 from DakotaNelson/plugins

Add basic plugin functionality to Empire

lib/common/empire.py

@@ -21,6 +21,8 @@ import hashlib
 import time
 import fnmatch
 import shlex
+import pkgutil
+import importlib
 
 # Empire imports
 import helpers
@@ -30,6 +32,7 @@ import listeners
 import modules
 import stagers
 import credentials
+import plugins
 from zlib_wrapper import compress
 from zlib_wrapper import decompress
 
@@ -68,6 +71,10 @@ class MainMenu(cmd.Cmd):
         # globalOptions[optionName] = (value, required, description)
         self.globalOptions = {}
 
+        # currently active plugins:
+        # {'pluginName': classObject}
+        self.loadedPlugins = {}
+
         # empty database object
         self.conn = self.database_connect()
         time.sleep(1)
@@ -381,6 +388,51 @@ class MainMenu(cmd.Cmd):
     # CMD methods
     ###################################################
 
+    def do_plugins(self, args):
+        "List all available and active plugins."
+        pluginPath = os.path.abspath("plugins")
+        print(helpers.color("[*] Searching for plugins at {}".format(pluginPath)))
+        # From walk_packages: "Note that this function must import all packages
+        # (not all modules!) on the given path, in order to access the __path__
+        # attribute to find submodules."
+        pluginNames = [name for _, name, _ in pkgutil.walk_packages([pluginPath])]
+        numFound = len(pluginNames)
+
+        # say how many we found, handling the 1 case
+        if numFound == 1:
+            print(helpers.color("[*] {} plugin found".format(numFound)))
+        else:
+            print(helpers.color("[*] {} plugins found".format(numFound)))
+
+        # if we found any, list them
+        if numFound > 0:
+            print("\tName\tActive")
+            print("\t----\t------")
+            activePlugins = self.loadedPlugins.keys()
+            for name in pluginNames:
+                active = ""
+                if name in activePlugins:
+                    active = "******"
+                print("\t" + name + "\t" + active)
+
+        print("")
+        print(helpers.color("[*] Use \"plugin <plugin name>\" to load a plugin."))
+
+    def do_plugin(self, pluginName):
+        "Load a plugin file to extend Empire."
+        pluginPath = os.path.abspath("plugins")
+        print(helpers.color("[*] Searching for plugins at {}".format(pluginPath)))
+        # From walk_packages: "Note that this function must import all packages
+        # (not all modules!) on the given path, in order to access the __path__
+        # attribute to find submodules."
+        pluginNames = [name for _, name, _ in pkgutil.walk_packages([pluginPath])]
+        if pluginName in pluginNames:
+            print(helpers.color("[*] Plugin {} found.".format(pluginName)))
+            # 'self' is the mainMenu object
+            plugins.load_plugin(self, pluginName)
+        else:
+            raise Exception("[!] Error: the plugin specified does not exist in {}.".format(pluginPath))
+
     def postcmd(self, stop, line):
 	if len(self.resourceQueue) > 0:
 	    nextcmd = self.resourceQueue.pop(0)

lib/common/plugins.py

@@ -0,0 +1,41 @@
+""" Utilities and helpers and etc. for plugins """
+
+import importlib
+
+import lib.common.helpers as helpers
+
+def load_plugin(mainMenu, pluginName):
+    """ Given the name of a plugin and a menu object, load it into the menu """
+    # note the 'plugins' package so the loader can find our plugin
+    fullPluginName = "plugins." + pluginName
+    module = importlib.import_module(fullPluginName)
+    pluginObj = module.Plugin(mainMenu)
+    mainMenu.loadedPlugins[pluginName] = pluginObj
+
+class Plugin(object):
+    # to be overwritten by child
+    description = "This is a description of this plugin."
+
+    def __init__(self, mainMenu):
+        # having these multiple messages should be helpful for debugging
+        # user-reported errors (can narrow down where they happen)
+        print(helpers.color("[*] Initializing plugin..."))
+        # any future init stuff goes here
+
+        print(helpers.color("[*] Doing custom initialization..."))
+        # do custom user stuff
+        self.onLoad()
+
+        # now that everything is loaded, register functions and etc. onto the main menu
+        print(helpers.color("[*] Registering plugin with menu..."))
+        self.register(mainMenu)
+
+    def onLoad(self):
+        """ Things to do during init: meant to be overridden by
+        the inheriting plugin. """
+        pass
+
+    def register(self, mainMenu):
+        """ Any modifications made to the main menu are done here
+        (meant to be overriden by child) """
+        pass

plugins/example.py

@@ -0,0 +1,35 @@
+""" An example of a plugin. """
+
+from lib.common.plugins import Plugin
+import lib.common.helpers as helpers
+
+# anything you simply write out (like a script) will run immediately when the
+# module is imported (before the class is instantiated)
+print("Hello from your new plugin!")
+
+# this class MUST be named Plugin
+class Plugin(Plugin):
+    description = "An example plugin."
+
+    def onLoad(self):
+        """ any custom loading behavior - called by init, so any
+        behavior you'd normally put in __init__ goes here """
+        print("Custom loading behavior happens now.")
+
+        # you can store data here that will persist until the plugin
+        # is unloaded (i.e. Empire closes)
+        self.calledTimes = 0
+
+    def register(self, mainMenu):
+        """ any modifications to the mainMenu go here - e.g.
+        registering functions to be run by user commands """
+        mainMenu.__class__.do_test = self.do_test
+
+    def do_test(self, args):
+        "An example of a plugin function."
+        print("This is executed from a plugin!")
+        print(helpers.color("[*] It can even import Empire functionality!"))
+
+        # you can also store data in the plugin (see onLoad)
+        self.calledTimes += 1
+        print("This function has been called {} times.".format(self.calledTimes))