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: