fix attempt
This commit is contained in:
parent
9e66fd0785
commit
d2e4c6ee49
@ -63,6 +63,14 @@ class GraphManager:
|
|||||||
self.creation_time = datetime.now(timezone.utc).isoformat()
|
self.creation_time = datetime.now(timezone.utc).isoformat()
|
||||||
self.last_modified = self.creation_time
|
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,
|
def add_node(self, node_id: str, node_type: NodeType,
|
||||||
metadata: Optional[Dict[str, Any]] = None) -> bool:
|
metadata: Optional[Dict[str, Any]] = None) -> bool:
|
||||||
"""
|
"""
|
||||||
|
@ -81,13 +81,15 @@ class ForensicLogger:
|
|||||||
self.logger.addHandler(console_handler)
|
self.logger.addHandler(console_handler)
|
||||||
|
|
||||||
def __getstate__(self):
|
def __getstate__(self):
|
||||||
|
"""Prepare ForensicLogger for pickling by excluding unpicklable objects."""
|
||||||
state = self.__dict__.copy()
|
state = self.__dict__.copy()
|
||||||
# Exclude the unpickleable 'logger' attribute
|
# Remove the unpickleable 'logger' attribute
|
||||||
if 'logger' in state:
|
if 'logger' in state:
|
||||||
del state['logger']
|
del state['logger']
|
||||||
return state
|
return state
|
||||||
|
|
||||||
def __setstate__(self, state):
|
def __setstate__(self, state):
|
||||||
|
"""Restore ForensicLogger after unpickling by reconstructing logger."""
|
||||||
self.__dict__.update(state)
|
self.__dict__.update(state)
|
||||||
# Re-initialize the 'logger' attribute
|
# Re-initialize the 'logger' attribute
|
||||||
self.logger = logging.getLogger(f'dnsrecon.{self.session_id}')
|
self.logger = logging.getLogger(f'dnsrecon.{self.session_id}')
|
||||||
|
@ -97,6 +97,13 @@ class Scanner:
|
|||||||
if attr in state:
|
if attr in state:
|
||||||
del state[attr]
|
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
|
return state
|
||||||
|
|
||||||
def __setstate__(self, state):
|
def __setstate__(self, state):
|
||||||
@ -108,6 +115,12 @@ class Scanner:
|
|||||||
self.scan_thread = None
|
self.scan_thread = None
|
||||||
self.executor = 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:
|
def _initialize_providers(self) -> None:
|
||||||
"""Initialize all available providers based on session configuration."""
|
"""Initialize all available providers based on session configuration."""
|
||||||
self.providers = []
|
self.providers = []
|
||||||
|
@ -96,16 +96,21 @@ class BaseProvider(ABC):
|
|||||||
print(f"Initialized {name} provider with session-specific config (rate: {actual_rate_limit}/min)")
|
print(f"Initialized {name} provider with session-specific config (rate: {actual_rate_limit}/min)")
|
||||||
|
|
||||||
def __getstate__(self):
|
def __getstate__(self):
|
||||||
|
"""Prepare BaseProvider for pickling by excluding unpicklable objects."""
|
||||||
state = self.__dict__.copy()
|
state = self.__dict__.copy()
|
||||||
# Exclude the unpickleable '_local' attribute
|
# Exclude the unpickleable '_local' attribute and stop event
|
||||||
if '_local' in state:
|
unpicklable_attrs = ['_local', '_stop_event']
|
||||||
del state['_local']
|
for attr in unpicklable_attrs:
|
||||||
|
if attr in state:
|
||||||
|
del state[attr]
|
||||||
return state
|
return state
|
||||||
|
|
||||||
def __setstate__(self, state):
|
def __setstate__(self, state):
|
||||||
|
"""Restore BaseProvider after unpickling by reconstructing threading objects."""
|
||||||
self.__dict__.update(state)
|
self.__dict__.update(state)
|
||||||
# Re-initialize the '_local' attribute
|
# Re-initialize the '_local' attribute and stop event
|
||||||
self._local = threading.local()
|
self._local = threading.local()
|
||||||
|
self._stop_event = None
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def session(self):
|
def session(self):
|
||||||
|
Loading…
x
Reference in New Issue
Block a user