new highest-priority-first scheduler
This commit is contained in:
@@ -3,14 +3,15 @@ Data provider modules for DNSRecon.
|
||||
Contains implementations for various reconnaissance data sources.
|
||||
"""
|
||||
|
||||
from .base_provider import BaseProvider, RateLimiter
|
||||
from .base_provider import BaseProvider
|
||||
from .crtsh_provider import CrtShProvider
|
||||
from .dns_provider import DNSProvider
|
||||
from .shodan_provider import ShodanProvider
|
||||
from core.rate_limiter import GlobalRateLimiter
|
||||
|
||||
__all__ = [
|
||||
'BaseProvider',
|
||||
'RateLimiter',
|
||||
'GlobalRateLimiter',
|
||||
'CrtShProvider',
|
||||
'DNSProvider',
|
||||
'ShodanProvider'
|
||||
|
||||
@@ -3,44 +3,12 @@
|
||||
import time
|
||||
import requests
|
||||
import threading
|
||||
import redis
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
|
||||
from core.logger import get_forensic_logger
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
"""Simple rate limiter for API calls."""
|
||||
|
||||
def __init__(self, requests_per_minute: int):
|
||||
"""
|
||||
Initialize rate limiter.
|
||||
|
||||
Args:
|
||||
requests_per_minute: Maximum requests allowed per minute
|
||||
"""
|
||||
self.requests_per_minute = requests_per_minute
|
||||
self.min_interval = 60.0 / requests_per_minute
|
||||
self.last_request_time = 0
|
||||
|
||||
def __getstate__(self):
|
||||
"""RateLimiter is fully picklable, return full state."""
|
||||
return self.__dict__.copy()
|
||||
|
||||
def __setstate__(self, state):
|
||||
"""Restore RateLimiter state."""
|
||||
self.__dict__.update(state)
|
||||
|
||||
def wait_if_needed(self) -> None:
|
||||
"""Wait if necessary to respect rate limits."""
|
||||
current_time = time.time()
|
||||
time_since_last = current_time - self.last_request_time
|
||||
|
||||
if time_since_last < self.min_interval:
|
||||
sleep_time = self.min_interval - time_since_last
|
||||
time.sleep(sleep_time)
|
||||
|
||||
self.last_request_time = time.time()
|
||||
from core.rate_limiter import GlobalRateLimiter
|
||||
|
||||
|
||||
class BaseProvider(ABC):
|
||||
@@ -72,7 +40,6 @@ class BaseProvider(ABC):
|
||||
actual_timeout = timeout
|
||||
|
||||
self.name = name
|
||||
self.rate_limiter = RateLimiter(actual_rate_limit)
|
||||
self.timeout = actual_timeout
|
||||
self._local = threading.local()
|
||||
self.logger = get_forensic_logger()
|
||||
@@ -172,8 +139,6 @@ class BaseProvider(ABC):
|
||||
print(f"Request cancelled before start: {url}")
|
||||
return None
|
||||
|
||||
self.rate_limiter.wait_if_needed()
|
||||
|
||||
start_time = time.time()
|
||||
response = None
|
||||
error = None
|
||||
@@ -297,5 +262,5 @@ class BaseProvider(ABC):
|
||||
'failed_requests': self.failed_requests,
|
||||
'success_rate': (self.successful_requests / self.total_requests * 100) if self.total_requests > 0 else 0,
|
||||
'relationships_found': self.total_relationships_found,
|
||||
'rate_limit': self.rate_limiter.requests_per_minute
|
||||
'rate_limit': self.config.get_rate_limit(self.name)
|
||||
}
|
||||
Reference in New Issue
Block a user