WIP

Author: ostefano

misp_modules/__init__.py

@@ -17,8 +17,6 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
-
-import abc
 import pathlib
 import signal
 import sys
@@ -189,38 +187,16 @@ def main():
     signal.signal(signal.SIGINT, handle_signal)
     signal.signal(signal.SIGTERM, handle_signal)
 
-    arg_parser = argparse.ArgumentParser(
-        description="misp-modules server", formatter_class=argparse.RawTextHelpFormatter
-    )
+    arg_parser = argparse.ArgumentParser(description="misp-modules", formatter_class=argparse.RawTextHelpFormatter)
     arg_parser.add_argument("-t", "--test", default=False, action="store_true", help="Test mode")
     arg_parser.add_argument("-d", "--debug", default=False, action="store_true", help="Enable debugging")
-    arg_parser.add_argument("-p", "--port", default=6666, help="misp-modules TCP port (default 6666)")
-    arg_parser.add_argument(
-        "-l",
-        "--listen",
-        default="localhost",
-        help="misp-modules listen address (default localhost)",
-    )
-    arg_parser.add_argument(
-        "-m",
-        default=[],
-        action="append",
-        help="Register a custom module",
-    )
-    arg_parser.add_argument(
-        "--devel",
-        default=False,
-        action="store_true",
-        help="""Start in development mode, enable debug, start only the module(s) listed in -m.\nExample: -m misp_modules.modules.expansion.bgpranking""",
-    )
+    arg_parser.add_argument("-p", "--port", default=6666, help="port (default 6666)")
+    arg_parser.add_argument("-l", "--listen", default="localhost", help="address (default localhost)")
+    arg_parser.add_argument("-c", "--custom", default=None, help="custom root")
     args = arg_parser.parse_args()
 
     # Initialize
-    if args.devel:
-        init_logger(debug=True)
-        log.info("Launch MISP modules server in development mode. Enable debug, load a list of modules is -m is used.")
-    else:
-        init_logger(debug=args.debug)
+    init_logger(debug=args.debug)
 
     # Load helpers
     for helper in iterate_helpers(importlib.resources.files(__package__).joinpath(HELPERS_DIR)):
@@ -255,18 +231,20 @@ def main():
             log.info("MISP module %s (type=%s) imported", module_name, module_type.name)
 
     # Load custom modules
-    for module_type, module in iterate_modules(pathlib.Path("/Users/ortolanis/Downloads/test/")):
-        # remove extension from name
-        module_name = module.name.split(".")[0]
-        try:
-            imported_module = import_from_path(module_name, str(module_type.joinpath(module.name)))
-        except ImportError as e:
-            log.warning("CUSTOM MISP module %s (type=%s) failed: %s", module_name, module_type.name, e)
-            continue
-        else:
-            MODULES_HANDLERS[module_name] = imported_module
-            MODULES_HANDLERS[f"type:{module_name}"] = module_type.name
-            log.info("CUSTOM MISP module %s (type=%s) imported", module_name, module_type.name)
+    if args.custom:
+        log.info("Parsing custom modules from root directory: %s", args.custom)
+        for module_type, module in iterate_modules(pathlib.Path(args.custom)):
+            # remove extension from name
+            module_name = module.name.split(".")[0]
+            try:
+                imported_module = import_from_path(module_name, str(module_type.joinpath(module.name)))
+            except ImportError as e:
+                log.warning("CUSTOM MISP module %s (type=%s) failed: %s", module_name, module_type.name, e)
+                continue
+            else:
+                MODULES_HANDLERS[module_name] = imported_module
+                MODULES_HANDLERS[f"type:{module_name}"] = module_type.name
+                log.info("CUSTOM MISP module %s (type=%s) imported", module_name, module_type.name)
 
     service = [
         (r"/modules", ListModules),