refactor(iplookup): restructure module using OOP principles and cleaner architecture

Improved code structure by removing ResultManager class and integrating its functionality directly

Enhanced readability and maintainability through modular design and clean code practices

Simplified logic and eliminated unnecessary codes

Author: ayanrajpoot10

bugscanx/modules/scrapers/iplookup/iplookup.py

@@ -1,64 +1,93 @@
-import concurrent.futures
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
 from bugscanx.utils.common import get_input, is_cidr
 from .sources import get_scrapers
 from .utils import process_input, process_file
-from .result_manager import ResultManager
-from .logger import IPLookupConsole, console
-
-def extract_domains(ip, scrapers, ip_console):
-    ip_console.start_ip_scan(ip)
-    domains = []
-    for scraper in scrapers:
-        domain_list = scraper.fetch_domains(ip)
-        if domain_list:
-            domains.extend(domain_list)
-            
-    domains = sorted(set(domains))
-    return (ip, domains)
-
-def process_ips(ips, output_file):
-    if not ips:
-        console.print("[bold red] No valid IPs/CIDRs to process.[/bold red]")
-        return 0
-        
-    scrapers = get_scrapers()
-    ip_console = IPLookupConsole()
-    result_manager = ResultManager(output_file, ip_console)
-    
-    def process_ip(ip):
-        ip, domains = extract_domains(ip, scrapers, ip_console)
+from .logger import IPLookupConsole
+
+
+class IPLookup:
+    def __init__(self):
+        self.console = IPLookupConsole()
+        self.completed = 0
+
+    def _fetch_from_source(self, source, ip):
+        try:
+            return source.fetch(ip)
+        except Exception:
+            return set()
+
+    def _save_domains(self, domains, output_file):
         if domains:
-            result_manager.save_result(ip, domains)
-        return ip, domains
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
-        futures = {executor.submit(process_ip, ip): ip for ip in ips}
-        for future in concurrent.futures.as_completed(futures):
-            future.result()
-    
-    for scraper in scrapers:
-        scraper.close()
-        
-    ip_console.print_final_summary(output_file)
-    return ip_console.total_domains
-
-def get_input_interactively():
+            with open(output_file, "a", encoding="utf-8") as f:
+                f.write("\n".join(sorted(domains)) + "\n")
+
+    def process_ip(self, ip, output_file, scrapers, total):
+        self.console.print_ip_start(ip)
+        self.console.print_progress(self.completed, total)
+
+        with ThreadPoolExecutor(max_workers=3) as executor:
+            futures = [
+                executor.submit(self._fetch_from_source, source, ip)
+                for source in scrapers
+            ]
+            results = [f.result() for f in as_completed(futures)]
+
+        domains = set().union(*results) if results else set()
+
+        self.console.update_ip_stats(ip, len(domains))
+        self.console.print_ip_complete(ip, len(domains))
+        self._save_domains(domains, output_file)
+
+        self.completed += 1
+        self.console.print_progress(self.completed, total)
+        return domains
+
+    def run(self, ips, output_file, scrapers=None):
+        if not ips:
+            self.console.print_error("No valid IPs provided")
+            return
+
+        os.makedirs(os.path.dirname(output_file) or '.', exist_ok=True)
+        self.completed = 0
+        all_domains = set()
+        total = len(ips)
+        scrapers = scrapers or get_scrapers()
+
+        with ThreadPoolExecutor(max_workers=5) as executor:
+            futures = [
+                executor.submit(self.process_ip, ip, output_file, scrapers, total)
+                for ip in ips
+            ]
+            for future in as_completed(futures):
+                try:
+                    result = future.result()
+                    all_domains.update(result)
+                except Exception as e:
+                    self.console.print_error(f"Error processing IP: {str(e)}")
+
+        for scraper in scrapers:
+            scraper.close()
+
+        self.console.print_final_summary(output_file)
+        return all_domains
+
+
+def main():
     ips = []
-    
-    input_choice = get_input("Select input mode", "choice", 
-                           choices=["Manual", "File"])
-    
-    if input_choice == "Manual":
-        cidr = get_input("Enter IP or CIDR", validators=[is_cidr])
-        ips.extend(process_input(cidr))
+    input_type = get_input("Select input mode", "choice",
+                          choices=["Manual", "File"])
+
+    if input_type == "Manual":
+        ip_input = get_input("Enter IP or CIDR", validators=[is_cidr])
+        ips.extend(process_input(ip_input))
+        default_output = f"{ip_input}_domains.txt".replace("/", "-")
     else:
         file_path = get_input("Enter filename", "file")
         ips.extend(process_file(file_path))
-        
-    output_file = get_input("Enter output filename")
-    return ips, output_file
-
-def main(ips=None, output_file=None):
-    if ips is None or output_file is None:
-        ips, output_file = get_input_interactively()
-    process_ips(ips, output_file)
+        default_output = f"{file_path.rsplit('.', 1)[0]}_domains.txt"
+
+    output_file = get_input("Enter output filename", default=default_output)
+    iplookup = IPLookup()
+    iplookup.run(ips, output_file)

bugscanx/modules/scrapers/iplookup/logger.py

@@ -1,41 +1,36 @@
-from threading import RLock
 from rich.console import Console
 
-console = Console()
 
-class Logger:
+class IPLookupConsole(Console):
     def __init__(self):
-        self._lock = RLock()
+        super().__init__()
+        self.total_domains = 0
+        self.ip_stats = {}
+        self.disable_cursor()
 
-    def clear_line(self):
-        with self._lock:
-            print("\033[2K\r", end='', flush=True)
+    def disable_cursor(self):
+        print('\033[?25l', end='', flush=True)
 
-    def replace(self, message):
-        with self._lock:
-            print(f"{message}", end='', flush=True)
+    def enable_cursor(self):
+        print('\033[?25h', end='', flush=True)
 
-class IPLookupConsole:
-    def __init__(self):
-        self.total_domains = 0
-        self.ip_stats = {}
-        self.logger = Logger()
-    
-    def start_ip_scan(self, ip):
-        self.logger.clear_line()
-        console.print(f"[cyan] Processing: {ip}[/cyan]")
+    def print_ip_start(self, ip):
+        self.print(f"[cyan]Processing: {ip}[/cyan]")
 
     def update_ip_stats(self, ip, count):
         self.ip_stats[ip] = count
         self.total_domains += count
 
-    def print_ip_complete(self, ip, domains_count):
-        self.logger.clear_line()
-        console.print(f"[green] {ip}: {domains_count} domains found[/green]")
+    def print_ip_complete(self, ip, count):
+        self.print(f"[green]{ip}: {count} domains found[/green]")
 
     def print_final_summary(self, output_file):
-        console.print(f"\n[green] Total [bold]{self.total_domains}[/bold] domains found[/green]")
-        console.print(f"[green] Results saved to {output_file}[/green]")
-        
+        self.print(f"\n[green]Total: [bold]{self.total_domains}[/bold] domains found")
+        self.print(f"[green]Results saved to {output_file}[/green]")
+        self.enable_cursor()
+
+    def print_progress(self, current, total):
+        self.print(f"Progress: {current} / {total}", end="\r")
+    
     def print_error(self, message):
-        console.print(f"[bold red]✗ ERROR: {message}[/bold red]")
+        self.print(f"[red]{message}[/red]")

bugscanx/modules/scrapers/iplookup/result_manager.py

@@ -1,57 +1,53 @@
-import random
-import requests
+from abc import ABC, abstractmethod
+
 from bs4 import BeautifulSoup
-from bugscanx.utils.http import USER_AGENTS, EXTRA_HEADERS
-from .utils import RateLimiter
-
-class DomainScraper:
-    def __init__(self, rate_limit=1.0):
-        self.headers = {
-            "User-Agent": random.choice(USER_AGENTS),
-            **EXTRA_HEADERS
-        }
-        self.rate_limiter = RateLimiter(rate_limit)
-        self.session = requests.Session()
-        self.session.timeout = 10.0
-
-    def _make_request(self, url, method='get', data=None):
-        self.rate_limiter.acquire()
-        try:
-            if method == 'get':
-                response = self.session.get(url, headers=self.headers)
-            else:
-                response = self.session.post(url, headers=self.headers, data=data)
-            response.raise_for_status()
-            return response
-        except requests.RequestException:
-            return None
-
-    def fetch_domains(self, ip):
-        raise NotImplementedError
-
-    def close(self):
-        self.session.close()
-
-class RapidDNSScraper(DomainScraper):
-    def fetch_domains(self, ip):
-        response = self._make_request(f"https://rapiddns.io/sameip/{ip}")
-        if not response:
-            return []
-        soup = BeautifulSoup(response.content, 'html.parser')
-        return [row.find_all('td')[0].text.strip() 
-                for row in soup.find_all('tr') if row.find_all('td')]
-
-class YouGetSignalScraper(DomainScraper):
-    def fetch_domains(self, ip):
+
+from .utils import RequestHandler
+
+
+class DomainSource(RequestHandler, ABC):
+    def __init__(self, name):
+        super().__init__()
+        self.name = name
+        self.domains = set()
+
+    @abstractmethod
+    def fetch(self, ip):
+        pass
+
+
+class RapidDNSSource(DomainSource):
+    def __init__(self):
+        super().__init__("RapidDNS")
+
+    def fetch(self, ip):
+        response = self.get(f"https://rapiddns.io/sameip/{ip}")
+        if response:
+            soup = BeautifulSoup(response.content, 'html.parser')
+            self.domains.update(
+                row.find_all('td')[0].text.strip()
+                for row in soup.find_all('tr')
+                if row.find_all('td')
+            )
+        return self.domains
+
+
+class YouGetSignalSource(DomainSource):
+    def __init__(self):
+        super().__init__("YouGetSignal")
+
+    def fetch(self, ip):
         data = {'remoteAddress': ip, 'key': '', '_': ''}
-        response = self._make_request("https://domains.yougetsignal.com/domains.php",
-                                method='post', data=data)
-        if not response:
-            return []
-        return [domain[0] for domain in response.json().get("domainArray", [])]
+        response = self.post("https://domains.yougetsignal.com/domains.php", data=data)
+        if response:
+            self.domains.update(
+                domain[0] for domain in response.json().get("domainArray", [])
+            )
+        return self.domains
+
 
 def get_scrapers():
     return [
-        RapidDNSScraper(),
-        YouGetSignalScraper()
+        RapidDNSSource(),
+        YouGetSignalSource()
     ]