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.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:
 | 
			
		||||
        """
 | 
			
		||||
 | 
			
		||||
@ -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}')
 | 
			
		||||
 | 
			
		||||
@ -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."""
 | 
			
		||||
 | 
			
		||||
@ -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):
 | 
			
		||||
 | 
			
		||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user