refactor(scanner): reorganize files and merge logger into multithreading

Author: ayanrajpoot10

bugscanx/modules/scanners/port_scanner.py

@@ -1,14 +1,12 @@
 import socket
 from concurrent.futures import ThreadPoolExecutor, as_completed
-
 from rich.console import Console
 from rich.progress import (
     BarColumn,
     Progress,
     TextColumn,
     TimeRemainingColumn,
 )
-
 from bugscanx.utils.prompts import get_input
 
 console = Console()

bugscanx/modules/scanners/scanners/base.py

@@ -2,7 +2,7 @@
 import threading
 from datetime import datetime
 
-from .concurrency.multithread import MultiThread
+from .multithread import MultiThread
 
 
 class BaseScanner(MultiThread):

bugscanx/modules/scanners/scanners/concurrency/logger.py

@@ -11,7 +11,6 @@
 class DirectScannerBase(BaseScanner):
     requests = requests
     DEFAULT_TIMEOUT = 3
-    DEFAULT_RETRY = 1
 
     def __init__(
         self,
@@ -30,25 +29,16 @@ def __init__(
     def request(self, method, url, **kwargs):
         method = method.upper()
         kwargs['timeout'] = self.timeout
-        max_attempts = self.DEFAULT_RETRY
-
-        for attempt in range(max_attempts):
-            self.progress(method, url)
-            try:
-                return self.requests.request(method, url, **kwargs)
-            except (
-                requests.exceptions.ConnectionError,
-                requests.exceptions.ReadTimeout,
-                requests.exceptions.Timeout,
-                requests.exceptions.RequestException
-            ) as e:
-                wait_time = (1 if isinstance(e, requests.exceptions.ConnectionError)
-                           else 3)
-                for _ in self.sleep(wait_time):
-                    self.progress(method, url)
-                if attempt == max_attempts - 1:
-                    return None
-        return None
+        self.progress(method, url)
+        try:
+            return self.requests.request(method, url, **kwargs)
+        except (
+            requests.exceptions.ConnectionError,
+            requests.exceptions.ReadTimeout,
+            requests.exceptions.Timeout,
+            requests.exceptions.RequestException
+        ):
+            return None
 
     def task(self, payload):
         method = payload['method']

bugscanx/modules/scanners/scanners/direct.py

@@ -1,9 +1,52 @@
-import time
+import os
+import sys
 from abc import ABC, abstractmethod
 from queue import Queue, Empty
 from threading import Thread, RLock
 
-from .logger import Logger, CursorManager
+
+class Logger:
+    COLORS = {
+        'ORANGE': '\033[33m',
+        'MAGENTA': '\033[35m',
+        'CYAN': '\033[36m',
+        'LGRAY': '\033[37m',
+        'GRAY': '\033[90m',
+        'RED': '\033[91m',
+        'GREEN': '\033[92m',
+        'YELLOW': '\033[93m',
+        'BLUE': '\033[94m',
+    }
+    RESET = '\033[0m'
+    CLEAR_LINE = '\033[2K'
+
+    def __init__(self):
+        self._lock = RLock()
+
+    @classmethod
+    def colorize(cls, text, color):
+        return f"{cls.COLORS.get(color, '')}{text}{cls.RESET}"
+
+    def replace(self, message):
+        cols = os.get_terminal_size().columns
+        msg = f"{message[:cols - 3]}..." if len(message) > cols else message
+        with self._lock:
+            sys.stdout.write(f'{self.CLEAR_LINE}{msg}{self.RESET}\r')
+            sys.stdout.flush()
+
+    def log(self, message):
+        with self._lock:
+            sys.stderr.write(f'\r{self.CLEAR_LINE}{message}{self.RESET}\n')
+            sys.stderr.flush()
+
+
+class CursorManager:
+    def __enter__(self):
+        print('\033[?25l', end='', flush=True)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        print('\033[?25h', end='', flush=True)
 
 
 class MultiThread(ABC):
@@ -76,12 +119,6 @@ def progress(self, *extra):
         ] + [str(x) for x in extra if x]
         self.logger.replace(" - ".join(parts))
 
-    def sleep(self, seconds):
-        while seconds > 0:
-            yield seconds
-            time.sleep(1)
-            seconds -= 1
-
     @abstractmethod
     def generate_tasks(self): pass