Initial Commit

util/generic/log.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+from logging.handlers import RotatingFileHandler
+from enum import IntEnum
+from typing import Optional
+import logging
+
+@dataclass
+class Log:
+    class Level(IntEnum):
+        CRITICAL = 50
+        FATAL = CRITICAL
+        ERROR = 40
+        WARNING = 30
+        WARN = WARNING
+        INFO = 20
+        DEBUG = 10
+        NOTSET = 0
+
+    name: str
+    level: Level
+    file: Optional[str] = None
+    format: str = '[%(name)s] %(asctime)s - %(levelname)s - %(message)s'
+    date_format: str = '%Y-%m-%d %H:%M:%S'
+
+    def create_logger(self):
+        logger = logging.getLogger(self.name)
+        logger.setLevel(self.level)
+        formatter = logging.Formatter(self.format, self.date_format)
+        self._add_handler(logger, logging.StreamHandler(), formatter)
+        if self.file:
+            self._add_handler(logger, RotatingFileHandler(self.file, maxBytes=1000000, backupCount=3), formatter)
+        return logger
+    
+    def _add_handler(self, logger: logging.Logger, handler: logging.Handler, formatter: logging.Formatter):
+        handler.setLevel(self.level)
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)

util/generic/multithreader.py

@@ -0,0 +1,129 @@
+from util.generic.log import Log
+
+from dataclasses import dataclass
+from enum import Enum
+import sys
+import time
+from datetime import datetime
+from multiprocessing import Pool, Manager
+import threading
+from logging import Logger
+from typing import Optional
+import re
+
+@dataclass
+class Spinner:
+    char_list: list[str]
+    delay: float
+
+class Spinners(Enum):
+    BASIC = Spinner(["|", "/", "-", "\\"], 0.1)
+    SPIN_TRI_BLOCK = Spinner(["▙", "▛", "▜", "▟"], 0.1)
+    SPIN_RIGHT_ANGLE = Spinner(["🮤", "🮧", "🮥", "🮦"], 0.1)
+    SPIN_OPEN_CUBE = Spinner(["🮪", "🮫", "🮭", "🮬"], 0.1)
+    SPIN_DOTS = Spinner(["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"], 0.1)
+    GUY_DANCING = Spinner(["🯇", "🯈"], 0.3)
+    GUY_JUMPING = Spinner(["🯅", "🯆"], 0.3)
+    SHIFTING_PATTERN = Spinner(["🮕", "🮖"], 0.1)
+    SHIFTING_GRADIENT = Spinner(["░▒▓█▓▒░", "▒▓█▓▒░░", "▓█▓▒░▒▓", "█▓▒░▒▓█", "▓▒░▒▓█▓", "░░▒▓█▓▒"], 0.1)
+
+@dataclass
+class WorkerProcess:
+    function: callable
+    instructions: tuple
+
+def dispatch_worker(worker_process: WorkerProcess):
+    try:
+        result = worker_process.function(worker_process.instructions)
+        return result
+    except Exception as e:
+        raise
+    
+@dataclass
+class WorkerReturn:
+    status: bool
+    command: str
+    stdout: str = ""
+    stderr: str = ""
+    error: Optional[str] = None
+
+    def output_result(self, logger: Logger, success_regex: str = None, success_message: str = None):
+        if self.status:
+            result_level = Log.Level.INFO
+            if success_regex and success_message:
+                success_match = re.search(success_regex, self.command)
+                success_string = success_match.group(1) if success_match else "unknown_target"
+                logger.log(result_level, success_message.format(output = success_string))
+            else: 
+                logger.log(result_level, f"SUCCESS: {self.command}")
+        else:
+            result_level = Log.Level.ERROR
+            logger.log(result_level, f"FAILURE: {self.command}")
+            if self.error:
+                logger.log(result_level, f"Error: {self.error}")
+        
+        if self.stdout:
+            logger.log(result_level, f"STDOUT:\n{self.stdout}")
+        
+        if self.stderr:
+            logger.log(result_level, f"STDERR:\n{self.stderr}")
+
+class MultiprocessWorker:
+    def __init__(self, logger: Logger = None, max_processes=4, spinner_set:Spinners = Spinners.BASIC):
+        if logger is None: logger = Log(self.__class__.__name__, Log.Level.DEBUG).create_logger()
+        self.logger = logger
+        
+        self.tasks: list[WorkerProcess] = []
+        self.max_processes = max_processes
+        self.spinner:Spinner = spinner_set.value
+        self.manager = Manager()
+        self.active_processes = self.manager.Value('i', 0)
+        self.stop_display = self.manager.Event()
+        self._process_lock = threading.Lock()
+
+    def add_tasks(self, instructions: list[WorkerProcess]):
+        self.logger.log(Log.Level.DEBUG, f"Adding {len(instructions)} tasks.")
+        self.tasks.extend(instructions)
+    
+    def add_task(self, instruction: WorkerProcess):
+        self.logger.log(Log.Level.DEBUG, f"Adding task with callable {instruction.function.__name__}")
+        self.tasks.append(instruction)
+
+    def _spinner_display(self, total_tasks, results):
+        idx = 0
+        list_length = len(self.spinner.char_list)
+
+        while len(results) < total_tasks:
+            sys.stdout.write(f"\r{self.spinner.char_list[idx]} | Completed: {len(results)}/{total_tasks} | Time: {datetime.now().strftime('%H:%M:%S')} ")
+            sys.stdout.flush()
+            idx = (idx + 1) % list_length
+            time.sleep(self.spinner.delay)
+
+        sys.stdout.write("\r")
+        sys.stdout.flush()
+        
+    def run(self):
+        self.logger.log(Log.Level.DEBUG, f"Starting with {len(self.tasks)} tasks and max {self.max_processes} processes.")
+        results = []
+        spinner_thread = threading.Thread(target=self._spinner_display, args=(len(self.tasks), results))
+        spinner_thread.daemon = True
+        spinner_thread.start()
+
+        try:
+            with Pool(self.max_processes) as pool:
+                for wp in self.tasks:
+                    pool.apply_async(
+                        dispatch_worker,
+                        args=(wp,),
+                        callback=lambda res: results.append(res)
+                    )
+                pool.close()
+                pool.join()
+        except Exception as e:
+            self.logger.exception("Unexpected error during multiprocessing run")
+        finally:
+            self.stop_display.set()
+            spinner_thread.join()
+
+        self.logger.log(Log.Level.DEBUG, "All multiprocessing tasks completed.")
+        return results