diff --git a/core/scanner.py b/core/scanner.py index 0a54d22..6571a6e 100644 --- a/core/scanner.py +++ b/core/scanner.py @@ -82,6 +82,32 @@ class Scanner: traceback.print_exc() raise + def __getstate__(self): + """Prepare object for pickling by excluding unpicklable attributes.""" + state = self.__dict__.copy() + + # Remove unpicklable threading objects + unpicklable_attrs = [ + 'stop_event', + 'scan_thread', + 'executor' + ] + + for attr in unpicklable_attrs: + if attr in state: + del state[attr] + + return state + + def __setstate__(self, state): + """Restore object after unpickling by reconstructing threading objects.""" + self.__dict__.update(state) + + # Reconstruct threading objects + self.stop_event = threading.Event() + self.scan_thread = None + self.executor = None + def _initialize_providers(self) -> None: """Initialize all available providers based on session configuration.""" self.providers = [] diff --git a/core/session_manager.py b/core/session_manager.py index d8b74a3..6e53f2a 100644 --- a/core/session_manager.py +++ b/core/session_manager.py @@ -34,17 +34,21 @@ class SessionManager: print(f"SessionManager initialized with Redis backend and {session_timeout_minutes}min timeout") def __getstate__(self): + """Prepare SessionManager for pickling.""" state = self.__dict__.copy() - # Exclude the unpickleable 'lock' and 'cleanup_thread' attributes - if 'lock' in state: - del state['lock'] - if 'cleanup_thread' in state: - del state['cleanup_thread'] + # Exclude unpickleable attributes - Redis client and threading objects + unpicklable_attrs = ['lock', 'cleanup_thread', 'redis_client'] + for attr in unpicklable_attrs: + if attr in state: + del state[attr] return state def __setstate__(self, state): + """Restore SessionManager after unpickling.""" self.__dict__.update(state) - # Re-initialize the 'lock' and 'cleanup_thread' attributes + # Re-initialize unpickleable attributes + import redis + self.redis_client = redis.StrictRedis(db=0, decode_responses=False) self.lock = threading.Lock() self.cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True) self.cleanup_thread.start() diff --git a/providers/base_provider.py b/providers/base_provider.py index 9679f45..8c63106 100644 --- a/providers/base_provider.py +++ b/providers/base_provider.py @@ -26,6 +26,14 @@ class RateLimiter: 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()