From d539cc69ad0963b00f32a17bbaba0620b5e4d2f1 Mon Sep 17 00:00:00 2001 From: Dakota Nelson Date: Tue, 31 Oct 2017 11:15:14 +0000 Subject: [PATCH 1/2] Add basic plugin functionality to Empire --- lib/common/empire.py | 60 +++++++++++++++++++++++++++++++++++++++++++ lib/common/plugins.py | 31 ++++++++++++++++++++++ plugins/__init__.py | 0 plugins/example.py | 35 +++++++++++++++++++++++++ 4 files changed, 126 insertions(+) create mode 100644 lib/common/plugins.py create mode 100644 plugins/__init__.py create mode 100644 plugins/example.py diff --git a/lib/common/empire.py b/lib/common/empire.py index db696bc..494ee85 100644 --- a/lib/common/empire.py +++ b/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,59 @@ 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------") + for name in pluginNames: + activePlugins = self.loadedPlugins.keys() + active = "" + if name in activePlugins: + active = "******" + print("\t" + name + "\t" + active) + + print("") + print(helpers.color("[*] Use \"plugin \" 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))) + # note the 'plugins' package so the loader can find our plugin + fullPluginName = "plugins." + pluginName + plugin = importlib.import_module(fullPluginName) + + self.load_plugin(plugin, pluginName) + else: + raise Exception("[!] Error: the plugin specified does not exist in {}.".format(pluginPath)) + + def load_plugin(self, module, pluginName): + # "self" refers to the mainmenu object + plugin = module.Plugin(self) + self.loadedPlugins[pluginName] = plugin + def postcmd(self, stop, line): if len(self.resourceQueue) > 0: nextcmd = self.resourceQueue.pop(0) diff --git a/lib/common/plugins.py b/lib/common/plugins.py new file mode 100644 index 0000000..cec0726 --- /dev/null +++ b/lib/common/plugins.py @@ -0,0 +1,31 @@ +""" Utilities and helpers and etc. for plugins """ + +import lib.common.helpers as helpers + +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 diff --git a/plugins/__init__.py b/plugins/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/plugins/example.py b/plugins/example.py new file mode 100644 index 0000000..2a6c113 --- /dev/null +++ b/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)) From 3741b0e786ce757cadc522b8afdd204fd7bdd438 Mon Sep 17 00:00:00 2001 From: Dakota Nelson Date: Tue, 31 Oct 2017 13:20:13 +0000 Subject: [PATCH 2/2] Minor refactoring to plugin loading code --- lib/common/empire.py | 14 +++----------- lib/common/plugins.py | 10 ++++++++++ 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/lib/common/empire.py b/lib/common/empire.py index 494ee85..b5e0260 100644 --- a/lib/common/empire.py +++ b/lib/common/empire.py @@ -408,8 +408,8 @@ class MainMenu(cmd.Cmd): if numFound > 0: print("\tName\tActive") print("\t----\t------") + activePlugins = self.loadedPlugins.keys() for name in pluginNames: - activePlugins = self.loadedPlugins.keys() active = "" if name in activePlugins: active = "******" @@ -428,19 +428,11 @@ class MainMenu(cmd.Cmd): pluginNames = [name for _, name, _ in pkgutil.walk_packages([pluginPath])] if pluginName in pluginNames: print(helpers.color("[*] Plugin {} found.".format(pluginName))) - # note the 'plugins' package so the loader can find our plugin - fullPluginName = "plugins." + pluginName - plugin = importlib.import_module(fullPluginName) - - self.load_plugin(plugin, 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 load_plugin(self, module, pluginName): - # "self" refers to the mainmenu object - plugin = module.Plugin(self) - self.loadedPlugins[pluginName] = plugin - def postcmd(self, stop, line): if len(self.resourceQueue) > 0: nextcmd = self.resourceQueue.pop(0) diff --git a/lib/common/plugins.py b/lib/common/plugins.py index cec0726..af54e88 100644 --- a/lib/common/plugins.py +++ b/lib/common/plugins.py @@ -1,7 +1,17 @@ """ 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."