95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
import imp
|
|
import os
|
|
import sys
|
|
import cme
|
|
from logging import getLogger
|
|
from cme.context import Context
|
|
from cme.logger import CMEAdapter
|
|
from cme.cmeserver import CMEServer
|
|
|
|
class ModuleLoader:
|
|
|
|
def __init__(self, args, db, logger):
|
|
self.args = args
|
|
self.db = db
|
|
self.logger = logger
|
|
self.server_port = args.server_port
|
|
self.cme_path = os.path.expanduser('~/.cme')
|
|
|
|
def module_is_sane(self, module, module_path):
|
|
module_error = False
|
|
|
|
if not hasattr(module, 'name'):
|
|
self.logger.error('{} missing the name variable'.format(module_path))
|
|
module_error = True
|
|
|
|
elif not hasattr(module, 'description'):
|
|
self.logger.error('{} missing the description variable'.format(module_path))
|
|
module_error = True
|
|
|
|
elif not hasattr(module, 'options'):
|
|
self.logger.error('{} missing the options function'.format(module_path))
|
|
module_error = True
|
|
|
|
elif not hasattr(module, 'on_login') and not (module, 'on_admin_login'):
|
|
self.logger.error('{} missing the on_login/on_admin_login function(s)'.format(module_path))
|
|
module_error = True
|
|
|
|
if module_error: return False
|
|
|
|
return True
|
|
|
|
def load_module(self, module_path):
|
|
module = imp.load_source('payload_module', module_path).CMEModule()
|
|
if self.module_is_sane(module, module_path):
|
|
return module
|
|
|
|
return None
|
|
|
|
def get_modules(self):
|
|
modules = {}
|
|
|
|
modules_paths = [os.path.join(os.path.dirname(cme.__file__), 'modules'), os.path.join(self.cme_path, 'modules')]
|
|
|
|
for path in modules_paths:
|
|
for module in os.listdir(path):
|
|
if module[-3:] == '.py' and module != 'example_module.py':
|
|
module_path = os.path.join(path, module)
|
|
m = self.load_module(os.path.join(path, module))
|
|
if m:
|
|
modules[m.name] = {'path': os.path.join(path, module), 'description': m.description, 'options': m.options.__doc__}
|
|
|
|
return modules
|
|
|
|
def init_module(self, module_path):
|
|
|
|
module = None
|
|
server = None
|
|
context = None
|
|
|
|
module = self.load_module(module_path)
|
|
|
|
if module:
|
|
module_logger = CMEAdapter(getLogger('CME'), {'module': module.name.upper()})
|
|
context = Context(self.db, module_logger, self.args)
|
|
|
|
module_options = {}
|
|
|
|
for option in self.args.module_options:
|
|
key, value = option.split('=', 1)
|
|
module_options[str(key).upper()] = value
|
|
|
|
module.options(context, module_options)
|
|
|
|
if hasattr(module, 'on_request') or hasattr(module, 'has_response'):
|
|
|
|
if hasattr(module, 'required_server'):
|
|
args.server = getattr(module, 'required_server')
|
|
|
|
if not self.server_port:
|
|
self. server_port = self.args.server_port
|
|
|
|
server = CMEServer(module, context, self.logger, self.args.server_host, self.server_port, self.args.server)
|
|
server.start()
|
|
|
|
return module, context, server |