Make loader more flexible

bin/misp-modules.py

@@ -25,12 +25,12 @@ import tornado.web
 import importlib
 import json
 import logging
-import re
+import fnmatch
 
-runPath = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(os.path.join(runPath, '..'))
 port = 6666
 
+
+def init_logger():
     log = logging.getLogger('misp-modules')
     formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
     handler = logging.StreamHandler(stream=sys.stdout)
@@ -39,22 +39,26 @@ handler.setLevel(logging.INFO)
 
     log.addHandler(handler)
     log.setLevel(logging.INFO)
+    return log
 
-modulesdir = '../modules/expansion'
 
+def load_modules(mod_dir):
+    sys.path.append(mod_dir)
     mhandlers = {}
     modules = []
-for module in os.listdir(modulesdir):
+    for root, dirnames, filenames in os.walk(mod_dir):
-    if ".py" not in module or ".pyc" in module or ".py~" in module:
+        if os.path.basename(root) == '__pycache__':
             continue
-    if re.match("^\.", module):
+        for filename in fnmatch.filter(filenames, '*.py'):
+            if filename == '__init__.py':
                 continue
-    modulename = module.split(".")[0]
+            modulename = filename.split(".")[0]
             moduletype = os.path.split(modulesdir)[1]
             modules.append(modulename)
             log.info('MISP modules {0} imported'.format(modulename))
-    mhandlers[modulename] = importlib.import_module('modules.expansion.' + modulename)
+            mhandlers[modulename] = importlib.import_module(os.path.basename(root) + '.' + modulename)
             mhandlers['type:' + modulename] = moduletype
+    return mhandlers, modules
 
 
 class ListModules(tornado.web.RequestHandler):
@@ -80,6 +84,10 @@ class QueryModule(tornado.web.RequestHandler):
         self.write(json.dumps(ret))
 
 
+if __name__ == '__main__':
+    modulesdir = '../modules'
+    log = init_logger()
+    mhandlers, modules = load_modules(modulesdir)
     service = [(r'/modules', ListModules), (r'/query', QueryModule)]
 
     application = tornado.web.Application(service)