Improvements - skip not applicable scripts, fix workers system

.gitignore

@@ -1,6 +1,7 @@
 # Custom
 .idea/
 .DS_Store
+results/*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

example_scenario.yaml

@@ -17,13 +17,38 @@
   description: Recon example for 'facebook.com' website
   kwargs:
     url: https://facebook.com
+- case: recon
+  name: vk-website
+  description: Recon example for 'vk.com' website
+  kwargs:
+    url: https://vk.com
+- case: recon
+  name: mail-website
+  description: Recon example for 'mail.ru' website
+  kwargs:
+    url: https://mail.ru
 - case: recon
   name: 8-8-8-8-host
   description: Recon example for '8.8.8.8' host
   kwargs:
     ip: 8.8.8.8
+- case: recon
+  name: 92-63-64-162-host
+  description: Recon example for '92.63.64.162' host
+  kwargs:
+    ip: 92.63.64.162
 - case: recon
   name: 13-91-95-74-host
   description: Recon example for '13.91.95.74' host
   kwargs:
     ip: 13.91.95.74
+- case: recon
+  name: 87-240-190-78-host
+  description: Recon example for '87.240.190.78' host
+  kwargs:
+    ip: 87.240.190.78
+- case: osint
+  name: phone-check
+  description: check information about the phone number
+  kwargs:
+    phone: 89138111111

main.py

@@ -4,13 +4,13 @@
 Main runner.
 """
 
-from logging import basicConfig, INFO
-from pprint import pprint
-from yaml import safe_load
+from datetime import datetime
 from json import dump
-from sys import argv
+from logging import basicConfig, INFO
 from pathlib import Path
-from datetime import datetime
+from sys import argv
+
+from yaml import safe_load
 
 from src.core.runner.manager import CaseManager
 from src.core.utils.log import Logger
@@ -66,7 +66,7 @@ def save_results(results: dict or list, name: str or None = "scenario") -> None:
         file=str(DefaultValues.RESULTS_DIR.joinpath(f"{name}_{current_time}.json")),
         mode="w",
     ) as results_file:
-        dump(results, results_file, indent=2, default=str)
+        dump(results, results_file, indent=2, default=str, ensure_ascii=False)
 
 
 if __name__ == "__main__":
@@ -84,7 +84,7 @@ def save_results(results: dict or list, name: str or None = "scenario") -> None:
     logger.info(f"Start framework for {len(scenario_cases)} cases")
 
     # Define CaseManager class
-    manager = CaseManager(cases=scenario_cases, max_workers=5)
+    manager = CaseManager(cases=scenario_cases)
 
     # Run all the cases in parallel way
     multiple_results = list(manager.multi_case_runner())

src/core/base/base.py

@@ -11,6 +11,7 @@ class BaseRunner:
     """
     Defines base runner class to create childs from
     """
+    required = []
 
     def __init__(self, logger: str = __name__):
         """

src/core/base/osint.py

@@ -20,6 +20,8 @@ class PossibleKeys:
 
 
 class OsintRunner(BaseRunner):
+    required = []
+
     def __init__(self, logger: str = __name__):
         super(OsintRunner, self).__init__(logger)
 

src/core/base/recon.py

@@ -14,10 +14,14 @@ class PossibleKeys:
     Defines default values for the function arguments (kwargs, named args)
     """
 
+    # fmt: off
     KEYS = ["host", "hostname", "ip", "url", "torrent_api_key"]
+    # fmt: on
 
 
 class ReconRunner(BaseRunner):
+    required = []
+
     def __init__(self, logger: str = __name__):
         super(ReconRunner, self).__init__(logger)
 

src/core/runner/manager.py

@@ -4,13 +4,13 @@
 Defines multiprocessing manager
 """
 
-from concurrent.futures import ProcessPoolExecutor
-from functools import partial
-from src.core.utils.log import Logger
+from concurrent.futures import ProcessPoolExecutor, as_completed
 
 from src.core.case.base import BaseCase
 from src.core.case.osint import OsintCase
 from src.core.case.recon import ReconCase
+from src.core.utils.log import Logger
+from src.core.values.defaults import CoreDefaults
 
 logger = Logger.get_logger(name=__name__)
 
@@ -27,7 +27,7 @@ class CaseManager:
         "default": BaseCase,
     }
 
-    def __init__(self, cases: list or None = None, max_workers: int = 3):
+    def __init__(self, cases: list or None = None, max_workers: int = CoreDefaults.MAX_PROCESSES):
         """
         Init manager
         :param cases: cases to run
@@ -62,6 +62,10 @@ def single_case_runner(
             "case_class": self.MAPPING.get(case_class, "default").__name__,
             "case_name": case_name,
             "case_description": case_description,
+            "case_data": {
+                "args": args,
+                "kwargs": kwargs
+            },
             "case_results": case.get_results() or {},
         }
 
@@ -74,22 +78,18 @@ def multi_case_runner(
         :param max_workers: maximum processes
         :return: results
         """
-        futures = []
         with ProcessPoolExecutor(
             max_workers=max_workers or self.max_workers
         ) as executor:
-            for case in cases or self.cases:
-                futures.append(
-                    executor.submit(
-                        partial(
-                            self.single_case_runner,
-                            case_class=case.get("case"),
-                            case_name=case.get("name"),
-                            case_description=case.get("description"),
-                            *case.get("args", []),
-                            **case.get("kwargs", {}),
-                        )
-                    )
-                )
-        for future in futures:
+            futures = [
+                executor.submit(
+                    self.single_case_runner,
+                    case_class=case.get("case"),
+                    case_name=case.get("name"),
+                    case_description=case.get("description"),
+                    *case.get("args", []),
+                    **case.get("kwargs", {}),
+                ) for case in cases or self.cases
+            ]
+        for future in as_completed(futures):
             yield future.result()

src/core/runner/runner.py

@@ -4,15 +4,15 @@
 Defines basic scripts runner
 """
 
-from concurrent.futures import ThreadPoolExecutor
-from functools import partial
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from importlib.machinery import SourceFileLoader
 from pathlib import Path
 from types import ModuleType
 
 import urllib3
 
 from src.core.utils.response import ScriptResponse
+from src.core.values.defaults import CoreDefaults
 
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
@@ -77,7 +77,12 @@ def exec_script(
 
         # fmt: off
         try:
-            class_instance = getattr(module, script_class)(logger=path.parent.stem)
+            module_class = getattr(module, script_class)
+            applicable = set(module_class.required).intersection(kwargs.keys())
+            # If the current script is not applicable for the current set of arguments - skip it
+            if not applicable:
+                return
+            class_instance = module_class(logger=path.parent.stem)
             result.update(getattr(class_instance, function)(*args, **kwargs))
         except Exception as unexp_err:
             result.update(ScriptResponse.error(message=f"Unexpected execution error: {str(unexp_err)}"))
@@ -102,28 +107,37 @@ def get_scripts(self) -> dict:
                 self.scripts[directory.stem].append(file)
         return self.scripts
 
-    def run_category(self, category: str, max_workers: int = 10, *args, **kwargs) -> None:
+    def run_category(
+        self,
+        category: str,
+        max_workers: int = CoreDefaults.MAX_THREADS,
+        timeout: int = CoreDefaults.CASE_TIMEOUT,
+        *args,
+        **kwargs
+    ) -> None:
         """
         Run a category with scripts
         :param category: category to run
         :param max_workers: max quantity of workers
+        :param timeout: timeout to wait in seconds
         :return: nothing
         """
         if not self.scripts:
             self.get_scripts()
-        futures = []
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
-            for script in self.scripts.get(category, []):
-                futures.append(
-                    executor.submit(
-                        fn=partial(
-                            self.exec_script, path=script, args=args, kwargs=kwargs
-                        )
-                    )
-                )
-        for future in futures:
-            result = future.result()
-            self.results.update({result.get("script"): result})
+            futures = [
+                executor.submit(self.exec_script, path=script, args=args, kwargs=kwargs)
+                for script in self.scripts.get(category, [])
+            ]
+        try:
+            for future in as_completed(futures, timeout=timeout):
+                result = future.result()
+                if not result:
+                    continue
+                self.results.update({result.get("script"): result})
+        except TimeoutError:
+            # Tasks took too much time - kill the execution process and return empty results
+            self.results = {}
 
     def get_results(self) -> dict:
         """

src/core/values/defaults.py

@@ -7,6 +7,12 @@
 from random import randrange
 
 
+class CoreDefaults:
+    MAX_PROCESSES = 10
+    MAX_THREADS = 10
+    CASE_TIMEOUT = 5 * 60
+
+
 class TestDefaults:
     HOST = "127.0.0.1"
     DOWN_TIMEOUT = 0.1

src/scripts/convert/email_generator/module.py

@@ -48,6 +48,11 @@ def generate(self, username: str) -> list:
 
 
 class Runner(BaseRunner):
+    """
+    Class that generates different email addresses
+    """
+    required = ["username"]
+
     def __init__(self, logger: str = __name__):
         super(Runner, self).__init__(logger)
 

src/scripts/convert/phone_num_generator/module.py

@@ -7,6 +7,11 @@
 
 
 class Runner(OsintRunner):
+    """
+    Class that generates different phone numbers
+    """
+    required = ["phone", "region"]
+
     def __init__(self, logger: str = __name__):
         super(Runner, self).__init__(logger)
 

src/scripts/osint/check_nickname/module.py

@@ -64,6 +64,11 @@ async def check_nickname_async(nickname: str, social) -> list:
 
 
 class Runner(OsintRunner):
+    """
+    Class that performs nickname check
+    """
+    required = ["username"]
+
     def __init__(self, logger: str = __name__):
         super().__init__(logger=logger)
 

src/scripts/osint/email_verifier/module.py

@@ -8,6 +8,11 @@
 
 
 class Runner(OsintRunner):
+    """
+    Class that performs email verification
+    """
+    required = ["email"]
+
     def __init__(self, logger: str = __name__):
         super(Runner, self).__init__(logger)
 

src/scripts/osint/region_check/module.py

@@ -8,6 +8,11 @@
 
 
 class Runner(OsintRunner):
+    """
+    Class that performs region check by phone number
+    """
+    required = ["phone"]
+
     def __init__(self, logger: str = __name__):
         super(Runner, self).__init__(logger)
 

src/scripts/other/network_usage/module.py

@@ -7,6 +7,11 @@
 
 
 class Runner(ReconRunner):
+    """
+    Simple example module to returh headers
+    """
+    required = ["hostname"]
+
     def __init__(self, logger: str = __name__):
         super(Runner, self).__init__(logger)
 

src/scripts/other/simple_example/module.py

@@ -9,6 +9,7 @@ class Runner(BaseRunner):
     """
     Basic example
     """
+    required = []
 
     def __init__(self, logger: str = __name__):
         """

src/scripts/other/user_greeting/module.py

@@ -6,6 +6,11 @@
 
 
 class Runner(OsintRunner):
+    """
+    Class that greets users
+    """
+    required = ["username"]
+
     def __init__(self, logger: str = __name__):
         super(Runner, self).__init__(logger)
 

src/scripts/recon/allowed_methods/module.py

@@ -53,6 +53,11 @@ def get_random_method():
 
 
 class Runner(ReconRunner):
+    """
+    Class that performs allowed methods check
+    """
+    required = ["url"]
+
     def __init__(self, logger: str = __name__):
         super(Runner, self).__init__(logger)
 

src/scripts/recon/cookie_checker/module.py

@@ -11,6 +11,7 @@ class Runner(ReconRunner):
     """
     Class that performs cookie flags checking.
     """
+    required = ["url"]
 
     def __init__(self, logger: str = __name__):
         """