From 5c49b488ccaa430f49d4093870572c0f987717d4 Mon Sep 17 00:00:00 2001
From: Bernhard Mueller <bernhard.mueller11@gmail.com>
Date: Sun, 18 Nov 2018 13:04:43 +0700
Subject: [PATCH] Revert module loading behavoir

---
 mythril/analysis/security.py | 28 +++++++++++++++++---------
 mythril/interfaces/cli.py    | 39 +++++++++++++++++++-----------------
 2 files changed, 40 insertions(+), 27 deletions(-)

diff --git a/mythril/analysis/security.py b/mythril/analysis/security.py
index 7a98e128..1e89e6e0 100644
--- a/mythril/analysis/security.py
+++ b/mythril/analysis/security.py
@@ -2,6 +2,7 @@ from collections import defaultdict
 from ethereum.opcodes import opcodes
 from mythril.analysis import modules
 import pkgutil
+import importlib.util
 import logging
 
 
@@ -28,16 +29,25 @@ def get_detection_module_hooks():
     return dict(hook_dict)
 
 
-def get_detection_modules(entrypoint, except_modules=()):
-    except_modules = list(except_modules) + ["base"]
+def get_detection_modules(entrypoint, include_modules=()):
+    include_modules = list(include_modules)
+
     _modules = []
 
-    for loader, name, _ in pkgutil.walk_packages(modules.__path__):
-        module = loader.find_module(name).load_module(name)
-        if (
-            module.__name__ not in except_modules
-            and module.detector.entrypoint == entrypoint
-        ):
+    if not len(include_modules):
+
+        for loader, name, _ in pkgutil.walk_packages(modules.__path__):
+            module = loader.find_module(name).load_module(name)
+            if (
+                module.__name__ !=  "base"
+                and module.detector.entrypoint == entrypoint
+            ):
+                _modules.append(module)
+
+    else:
+        for module_name in include_modules:
+            module = importlib.import_module(module_name, modules)
+
             _modules.append(module)
 
     logging.info("Found %s detection modules", len(_modules))
@@ -48,7 +58,7 @@ def fire_lasers(statespace, module_names=()):
     logging.info("Starting analysis")
 
     issues = []
-    for module in get_detection_modules(entrypoint="post", except_modules=module_names):
+    for module in get_detection_modules(entrypoint="post", include_modules=module_names):
         logging.info("Executing " + module.detector.name)
         issues += module.detector.execute(statespace)
 
diff --git a/mythril/interfaces/cli.py b/mythril/interfaces/cli.py
index b529b66d..01635594 100644
--- a/mythril/interfaces/cli.py
+++ b/mythril/interfaces/cli.py
@@ -398,24 +398,27 @@ def main():
                     exit_with_error(args.outform, "Error saving graph: " + str(e))
 
             else:
-                report = mythril.fire_lasers(
-                    strategy=args.strategy,
-                    address=address,
-                    modules=[m.strip() for m in args.modules.strip().split(",")]
-                    if args.modules
-                    else [],
-                    verbose_report=args.verbose_report,
-                    max_depth=args.max_depth,
-                    execution_timeout=args.execution_timeout,
-                    create_timeout=args.create_timeout,
-                    max_transaction_count=args.max_transaction_count,
-                )
-                outputs = {
-                    "json": report.as_json(),
-                    "text": report.as_text(),
-                    "markdown": report.as_markdown(),
-                }
-                print(outputs[args.outform])
+                try:
+                    report = mythril.fire_lasers(
+                        strategy=args.strategy,
+                        address=address,
+                        modules=[m.strip() for m in args.modules.strip().split(",")]
+                        if args.modules
+                        else [],
+                        verbose_report=args.verbose_report,
+                        max_depth=args.max_depth,
+                        execution_timeout=args.execution_timeout,
+                        create_timeout=args.create_timeout,
+                        max_transaction_count=args.max_transaction_count,
+                    )
+                    outputs = {
+                        "json": report.as_json(),
+                        "text": report.as_text(),
+                        "markdown": report.as_markdown(),
+                    }
+                    print(outputs[args.outform])
+                except ModuleNotFoundError as e:
+                    exit_with_error(args.outform, "Error loading analyis modules: " + format(e))
 
         elif args.statespace_json: