From d2e4c6ee492681031c7be7320059267de86f6606 Mon Sep 17 00:00:00 2001
From: overcuriousity <overcuriousity@posteo.org>
Date: Fri, 12 Sep 2025 14:47:12 +0200
Subject: [PATCH] fix attempt

---
 core/graph_manager.py      |  8 ++++++++
 core/logger.py             |  4 +++-
 core/scanner.py            | 13 +++++++++++++
 providers/base_provider.py | 13 +++++++++----
 4 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/core/graph_manager.py b/core/graph_manager.py
index b28bc13..279af86 100644
--- a/core/graph_manager.py
+++ b/core/graph_manager.py
@@ -63,6 +63,14 @@ class GraphManager:
         self.creation_time = datetime.now(timezone.utc).isoformat()
         self.last_modified = self.creation_time
 
+    def __getstate__(self):
+        """GraphManager is fully picklable, return full state."""
+        return self.__dict__.copy()
+
+    def __setstate__(self, state):
+        """Restore GraphManager state."""
+        self.__dict__.update(state)
+
     def add_node(self, node_id: str, node_type: NodeType,
                  metadata: Optional[Dict[str, Any]] = None) -> bool:
         """
diff --git a/core/logger.py b/core/logger.py
index 19c850f..845c595 100644
--- a/core/logger.py
+++ b/core/logger.py
@@ -81,13 +81,15 @@ class ForensicLogger:
             self.logger.addHandler(console_handler)
 
     def __getstate__(self):
+        """Prepare ForensicLogger for pickling by excluding unpicklable objects."""
         state = self.__dict__.copy()
-        # Exclude the unpickleable 'logger' attribute
+        # Remove the unpickleable 'logger' attribute
         if 'logger' in state:
             del state['logger']
         return state
 
     def __setstate__(self, state):
+        """Restore ForensicLogger after unpickling by reconstructing logger."""
         self.__dict__.update(state)
         # Re-initialize the 'logger' attribute
         self.logger = logging.getLogger(f'dnsrecon.{self.session_id}')
diff --git a/core/scanner.py b/core/scanner.py
index 6571a6e..c80ee89 100644
--- a/core/scanner.py
+++ b/core/scanner.py
@@ -97,6 +97,13 @@ class Scanner:
             if attr in state:
                 del state[attr]
         
+        # Handle providers separately to ensure they're picklable
+        if 'providers' in state:
+            # The providers should be picklable now, but let's ensure clean state
+            for provider in state['providers']:
+                if hasattr(provider, '_stop_event'):
+                    provider._stop_event = None
+        
         return state
 
     def __setstate__(self, state):
@@ -107,6 +114,12 @@ class Scanner:
         self.stop_event = threading.Event()
         self.scan_thread = None
         self.executor = None
+        
+        # Re-set stop events for providers
+        if hasattr(self, 'providers'):
+            for provider in self.providers:
+                if hasattr(provider, 'set_stop_event'):
+                    provider.set_stop_event(self.stop_event)
 
     def _initialize_providers(self) -> None:
         """Initialize all available providers based on session configuration."""
diff --git a/providers/base_provider.py b/providers/base_provider.py
index 8c63106..6051339 100644
--- a/providers/base_provider.py
+++ b/providers/base_provider.py
@@ -96,16 +96,21 @@ class BaseProvider(ABC):
         print(f"Initialized {name} provider with session-specific config (rate: {actual_rate_limit}/min)")
 
     def __getstate__(self):
+        """Prepare BaseProvider for pickling by excluding unpicklable objects."""
         state = self.__dict__.copy()
-        # Exclude the unpickleable '_local' attribute
-        if '_local' in state:
-            del state['_local']
+        # Exclude the unpickleable '_local' attribute and stop event
+        unpicklable_attrs = ['_local', '_stop_event']
+        for attr in unpicklable_attrs:
+            if attr in state:
+                del state[attr]
         return state
 
     def __setstate__(self, state):
+        """Restore BaseProvider after unpickling by reconstructing threading objects."""
         self.__dict__.update(state)
-        # Re-initialize the '_local' attribute
+        # Re-initialize the '_local' attribute and stop event
         self._local = threading.local()
+        self._stop_event = None
 
     @property
     def session(self):