souleyez 2.43.29__py3-none-any.whl → 2.43.34__py3-none-any.whl

souleyez/ai/recommender.py

@@ -59,12 +59,14 @@ If no hosts are discovered yet, suggest reconnaissance actions.
         self.context_builder = ContextBuilder()
 
         # For backward compatibility, also set ollama attribute if using OllamaProvider
-        if hasattr(self.provider, 'service'):
+        if hasattr(self.provider, "service"):
             self.ollama = self.provider.service
         else:
             self.ollama = None
 
-    def suggest_next_step(self, engagement_id: int, target_host_ids: Optional[list] = None) -> Optional[Dict[str, Any]]:
+    def suggest_next_step(
+        self, engagement_id: int, target_host_ids: Optional[list] = None
+    ) -> Optional[Dict[str, Any]]:
         """
         Suggest the next best attack step for an engagement.
 
@@ -85,36 +87,40 @@ If no hosts are discovered yet, suggest reconnaissance actions.
         if not self.provider or not self.provider.is_available():
             logger.error("Cannot generate recommendation: AI provider not available")
             return {
-                'error': 'AI provider not available. Configure Ollama or Claude in Settings → AI Settings.',
-                'action': None,
-                'target': None,
-                'rationale': None,
-                'expected_outcome': None,
-                'risk_level': None
+                "error": "AI provider not available. Configure Ollama or Claude in Settings → AI Settings.",
+                "action": None,
+                "target": None,
+                "rationale": None,
+                "expected_outcome": None,
+                "risk_level": None,
             }
 
         # Build context from engagement data (with optional host filtering)
         try:
-            context = self.context_builder.build_context(engagement_id, target_host_ids=target_host_ids)
+            context = self.context_builder.build_context(
+                engagement_id, target_host_ids=target_host_ids
+            )
         except Exception as e:
             logger.error(f"Failed to build context: {e}")
             return {
-                'error': f'Failed to load engagement data: {e}',
-                'action': None,
-                'target': None,
-                'rationale': None,
-                'expected_outcome': None,
-                'risk_level': None
+                "error": f"Failed to load engagement data: {e}",
+                "action": None,
+                "target": None,
+                "rationale": None,
+                "expected_outcome": None,
+                "risk_level": None,
             }
 
         # Build prompt with state awareness
-        state_summary = self.context_builder.get_state_summary(engagement_id, target_host_ids=target_host_ids)
-        
+        state_summary = self.context_builder.get_state_summary(
+            engagement_id, target_host_ids=target_host_ids
+        )
+
         # Add targeting note if specific hosts selected
         targeting_note = ""
         if target_host_ids:
             targeting_note = f"\nNOTE: Only target the {len(target_host_ids)} selected host(s) listed above. Do not suggest actions against other hosts.\n"
-        
+
         prompt = f"""You are an expert penetration tester analyzing an engagement and suggesting the next most promising attack step.
 
 {state_summary}
@@ -158,47 +164,53 @@ Be specific and actionable. Focus on practical pentesting methodology.
 Continue from current state - don't restart from scratch."""
 
         # Generate recommendation
-        provider_type = getattr(self.provider, 'provider_type', 'unknown')
-        logger.info(f"Generating recommendation for engagement {engagement_id} using {provider_type}")
+        provider_type = getattr(self.provider, "provider_type", "unknown")
+        logger.info(
+            f"Generating recommendation for engagement {engagement_id} using {provider_type}"
+        )
         try:
-            response = self.provider.generate(prompt=prompt, max_tokens=2000, temperature=0.7)
+            response = self.provider.generate(
+                prompt=prompt, max_tokens=2000, temperature=0.7
+            )
             if not response:
                 logger.error("LLM returned empty response")
                 return {
-                    'error': 'LLM generation failed (empty response). Check AI provider configuration.',
-                    'action': None,
-                    'target': None,
-                    'rationale': None,
-                    'expected_outcome': None,
-                    'risk_level': None
+                    "error": "LLM generation failed (empty response). Check AI provider configuration.",
+                    "action": None,
+                    "target": None,
+                    "rationale": None,
+                    "expected_outcome": None,
+                    "risk_level": None,
                 }
         except Exception as e:
             logger.error(f"LLM generation failed: {e}")
             return {
-                'error': f'LLM generation failed: {e}',
-                'action': None,
-                'target': None,
-                'rationale': None,
-                'expected_outcome': None,
-                'risk_level': None
+                "error": f"LLM generation failed: {e}",
+                "action": None,
+                "target": None,
+                "rationale": None,
+                "expected_outcome": None,
+                "risk_level": None,
             }
 
         # Parse response
         try:
             recommendation = self._parse_response(response)
-            logger.info(f"Generated recommendation: {recommendation.get('action', 'unknown')}")
+            logger.info(
+                f"Generated recommendation: {recommendation.get('action', 'unknown')}"
+            )
             return recommendation
         except Exception as e:
             logger.error(f"Failed to parse LLM response: {e}")
             logger.debug(f"Raw response: {response}")
             return {
-                'error': f'Failed to parse LLM response: {e}',
-                'action': None,
-                'target': None,
-                'rationale': None,
-                'expected_outcome': None,
-                'risk_level': None,
-                'raw_response': response
+                "error": f"Failed to parse LLM response: {e}",
+                "action": None,
+                "target": None,
+                "rationale": None,
+                "expected_outcome": None,
+                "risk_level": None,
+                "raw_response": response,
             }
 
     def _parse_response(self, response: str) -> Dict[str, Any]:
@@ -222,25 +234,37 @@ Continue from current state - don't restart from scratch."""
             ValueError: If response format is invalid
         """
         # Extract fields using regex
-        action_match = re.search(r'ACTION:\s*(.+?)(?=\n[A-Z]+:|$)', response, re.DOTALL | re.IGNORECASE)
-        target_match = re.search(r'TARGET:\s*(.+?)(?=\n[A-Z]+:|$)', response, re.DOTALL | re.IGNORECASE)
-        rationale_match = re.search(r'RATIONALE:\s*(.+?)(?=\n[A-Z]+:|$)', response, re.DOTALL | re.IGNORECASE)
-        expected_match = re.search(r'EXPECTED:\s*(.+?)(?=\n[A-Z]+:|$)', response, re.DOTALL | re.IGNORECASE)
-        risk_match = re.search(r'RISK:\s*(.+?)(?=\n[A-Z]+:|$)', response, re.DOTALL | re.IGNORECASE)
+        action_match = re.search(
+            r"ACTION:\s*(.+?)(?=\n[A-Z]+:|$)", response, re.DOTALL | re.IGNORECASE
+        )
+        target_match = re.search(
+            r"TARGET:\s*(.+?)(?=\n[A-Z]+:|$)", response, re.DOTALL | re.IGNORECASE
+        )
+        rationale_match = re.search(
+            r"RATIONALE:\s*(.+?)(?=\n[A-Z]+:|$)", response, re.DOTALL | re.IGNORECASE
+        )
+        expected_match = re.search(
+            r"EXPECTED:\s*(.+?)(?=\n[A-Z]+:|$)", response, re.DOTALL | re.IGNORECASE
+        )
+        risk_match = re.search(
+            r"RISK:\s*(.+?)(?=\n[A-Z]+:|$)", response, re.DOTALL | re.IGNORECASE
+        )
 
         # Validate all fields present
-        if not all([action_match, target_match, rationale_match, expected_match, risk_match]):
+        if not all(
+            [action_match, target_match, rationale_match, expected_match, risk_match]
+        ):
             missing = []
             if not action_match:
-                missing.append('ACTION')
+                missing.append("ACTION")
             if not target_match:
-                missing.append('TARGET')
+                missing.append("TARGET")
             if not rationale_match:
-                missing.append('RATIONALE')
+                missing.append("RATIONALE")
             if not expected_match:
-                missing.append('EXPECTED')
+                missing.append("EXPECTED")
             if not risk_match:
-                missing.append('RISK')
+                missing.append("RISK")
             raise ValueError(f"Missing required fields: {', '.join(missing)}")
 
         # Extract and clean values
@@ -248,33 +272,35 @@ Continue from current state - don't restart from scratch."""
         target = target_match.group(1).strip()
         rationale = rationale_match.group(1).strip()
         expected = expected_match.group(1).strip()
-        
+
         # Extract risk level (grab only first word, handle multi-line responses)
         risk_text = risk_match.group(1).strip()
-        risk = risk_text.split()[0].lower() if risk_text else 'medium'
+        risk = risk_text.split()[0].lower() if risk_text else "medium"
 
         # Validate risk level
-        if risk not in ['low', 'medium', 'high']:
+        if risk not in ["low", "medium", "high"]:
             logger.warning(f"Invalid risk level '{risk}', defaulting to 'medium'")
-            risk = 'medium'
+            risk = "medium"
 
         return {
-            'error': None,
-            'action': action,
-            'target': target,
-            'rationale': rationale,
-            'expected_outcome': expected,
-            'risk_level': risk
+            "error": None,
+            "action": action,
+            "target": target,
+            "rationale": rationale,
+            "expected_outcome": expected,
+            "risk_level": risk,
         }
 
-    def generate_chain(self, engagement_id: int, num_steps: int = 5) -> Optional[Dict[str, Any]]:
+    def generate_chain(
+        self, engagement_id: int, num_steps: int = 5
+    ) -> Optional[Dict[str, Any]]:
         """
         Generate a multi-step attack chain.
-        
+
         Args:
             engagement_id: ID of engagement to analyze
             num_steps: Number of steps to generate
-            
+
         Returns:
             dict: Chain with list of steps, or error dict if failed
         """
@@ -282,8 +308,8 @@ Continue from current state - don't restart from scratch."""
         if not self.provider or not self.provider.is_available():
             logger.error("Cannot generate chain: AI provider not available")
             return {
-                'error': 'AI provider not available. Configure Ollama or Claude in Settings → AI Settings.',
-                'steps': []
+                "error": "AI provider not available. Configure Ollama or Claude in Settings → AI Settings.",
+                "steps": [],
             }
 
         # Build context from engagement data
@@ -291,62 +317,61 @@ Continue from current state - don't restart from scratch."""
             context = self.context_builder.build_context(engagement_id)
         except Exception as e:
             logger.error(f"Failed to build context: {e}")
-            return {
-                'error': f'Failed to load engagement data: {e}',
-                'steps': []
-            }
+            return {"error": f"Failed to load engagement data: {e}", "steps": []}
 
         # Create multi-step prompt
         prompt = self._build_chain_prompt(context, num_steps, engagement_id)
 
         # Get AI response
-        provider_type = getattr(self.provider, 'provider_type', 'unknown')
-        logger.info(f"Generating {num_steps}-step chain for engagement {engagement_id} using {provider_type}")
+        provider_type = getattr(self.provider, "provider_type", "unknown")
+        logger.info(
+            f"Generating {num_steps}-step chain for engagement {engagement_id} using {provider_type}"
+        )
         try:
-            response = self.provider.generate(prompt=prompt, max_tokens=3000, temperature=0.7)
+            response = self.provider.generate(
+                prompt=prompt, max_tokens=3000, temperature=0.7
+            )
             if not response:
                 logger.error("LLM returned empty response")
-                return {
-                    'error': 'LLM generation failed (empty response)',
-                    'steps': []
-                }
+                return {"error": "LLM generation failed (empty response)", "steps": []}
         except Exception as e:
             logger.error(f"LLM generation failed: {e}")
-            return {
-                'error': f'LLM generation failed: {e}',
-                'steps': []
-            }
+            return {"error": f"LLM generation failed: {e}", "steps": []}
 
         # Parse multi-step response
         try:
             steps = self._parse_chain_response(response)
             logger.info(f"Generated {len(steps)}-step attack chain")
-            
+
             # DEBUG: Log if we got empty steps
             if not steps:
-                logger.warning(f"Parsed 0 steps from response. Raw response:\n{response}")
-            
+                logger.warning(
+                    f"Parsed 0 steps from response. Raw response:\n{response}"
+                )
+
             return {
-                'error': None,
-                'engagement_id': engagement_id,
-                'num_steps': len(steps),
-                'steps': steps
+                "error": None,
+                "engagement_id": engagement_id,
+                "num_steps": len(steps),
+                "steps": steps,
             }
         except Exception as e:
             logger.error(f"Failed to parse chain response: {e}")
             logger.debug(f"Raw response: {response}")
             return {
-                'error': f'Failed to parse chain response: {e}',
-                'steps': [],
-                'raw_response': response
+                "error": f"Failed to parse chain response: {e}",
+                "steps": [],
+                "raw_response": response,
             }
 
-    def _build_chain_prompt(self, context: str, num_steps: int, engagement_id: int) -> str:
+    def _build_chain_prompt(
+        self, context: str, num_steps: int, engagement_id: int
+    ) -> str:
         """Build prompt for multi-step attack chain with state awareness."""
-        
+
         # Get current state summary
         state_summary = self.context_builder.get_state_summary(engagement_id)
-        
+
         return f"""You are an expert penetration tester. Analyze this engagement and create a {num_steps}-step attack chain.
 
 {state_summary}
@@ -388,19 +413,21 @@ Be specific and actionable. Think like a professional pentester. Continue from c
         logger.debug(f"Parsing chain response: {response[:200]}...")
 
         # Split by "STEP N:" pattern (handle both with and without leading newline)
-        step_blocks = re.split(r'STEP \d+:', response)
-        
+        step_blocks = re.split(r"STEP \d+:", response)
+
         logger.debug(f"Split into {len(step_blocks)} blocks")
 
-        for i, block in enumerate(step_blocks[1:], 1):  # Skip first block (before first STEP)
+        for i, block in enumerate(
+            step_blocks[1:], 1
+        ):  # Skip first block (before first STEP)
             step = {
-                'step_number': i,
-                'action': self._extract_field(block, 'ACTION'),
-                'target': self._extract_field(block, 'TARGET'),
-                'rationale': self._extract_field(block, 'RATIONALE'),
-                'expected': self._extract_field(block, 'EXPECTED'),
-                'risk': self._extract_risk_field(block),
-                'dependencies': self._extract_field(block, 'DEPENDENCIES')
+                "step_number": i,
+                "action": self._extract_field(block, "ACTION"),
+                "target": self._extract_field(block, "TARGET"),
+                "rationale": self._extract_field(block, "RATIONALE"),
+                "expected": self._extract_field(block, "EXPECTED"),
+                "risk": self._extract_risk_field(block),
+                "dependencies": self._extract_field(block, "DEPENDENCIES"),
             }
             steps.append(step)
 
@@ -408,7 +435,7 @@ Be specific and actionable. Think like a professional pentester. Continue from c
 
     def _extract_field(self, text: str, field_name: str) -> str:
         """Extract a field value from text block."""
-        pattern = rf'{field_name}:\s*(.+?)(?=\n[A-Z]+:|$)'
+        pattern = rf"{field_name}:\s*(.+?)(?=\n[A-Z]+:|$)"
         match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
         if match:
             return match.group(1).strip()
@@ -416,17 +443,23 @@ Be specific and actionable. Think like a professional pentester. Continue from c
 
     def _extract_risk_field(self, text: str) -> str:
         """Extract and validate risk field (handles multi-line responses)."""
-        risk_text = self._extract_field(text, 'RISK')
-        risk = risk_text.split()[0].upper() if risk_text and risk_text != "Unknown" else 'MEDIUM'
-        
+        risk_text = self._extract_field(text, "RISK")
+        risk = (
+            risk_text.split()[0].upper()
+            if risk_text and risk_text != "Unknown"
+            else "MEDIUM"
+        )
+
         # Validate risk level
-        if risk not in ['LOW', 'MEDIUM', 'HIGH']:
+        if risk not in ["LOW", "MEDIUM", "HIGH"]:
             logger.warning(f"Invalid risk level '{risk}', defaulting to 'MEDIUM'")
-            risk = 'MEDIUM'
-        
+            risk = "MEDIUM"
+
         return risk
 
-    def suggest_multiple_paths(self, engagement_id: int, num_paths: int = 3) -> Dict[str, Any]:
+    def suggest_multiple_paths(
+        self, engagement_id: int, num_paths: int = 3
+    ) -> Dict[str, Any]:
         """
         Generate multiple alternative attack paths and rank them.
 
@@ -444,8 +477,8 @@ Be specific and actionable. Think like a professional pentester. Continue from c
         if not self.provider or not self.provider.is_available():
             logger.error("Cannot generate paths: AI provider not available")
             return {
-                'error': 'AI provider not available. Configure Ollama or Claude in Settings → AI Settings.',
-                'paths': []
+                "error": "AI provider not available. Configure Ollama or Claude in Settings → AI Settings.",
+                "paths": [],
             }
 
         # Build context
@@ -453,31 +486,29 @@ Be specific and actionable. Think like a professional pentester. Continue from c
             context = self.context_builder.build_context(engagement_id)
         except Exception as e:
             logger.error(f"Failed to build context: {e}")
-            return {
-                'error': f'Failed to load engagement data: {e}',
-                'paths': []
-            }
+            return {"error": f"Failed to load engagement data: {e}", "paths": []}
 
         # Build multi-path prompt
         prompt = self._build_multi_path_prompt(context, num_paths)
 
         # Generate recommendations
-        provider_type = getattr(self.provider, 'provider_type', 'unknown')
-        logger.info(f"Generating {num_paths} alternative paths for engagement {engagement_id} using {provider_type}")
+        provider_type = getattr(self.provider, "provider_type", "unknown")
+        logger.info(
+            f"Generating {num_paths} alternative paths for engagement {engagement_id} using {provider_type}"
+        )
         try:
-            response = self.provider.generate(prompt=prompt, max_tokens=2500, temperature=0.7)
+            response = self.provider.generate(
+                prompt=prompt, max_tokens=2500, temperature=0.7
+            )
             if not response:
                 logger.error("LLM returned empty response")
                 return {
-                    'error': 'AI generation failed. Check provider configuration.',
-                    'paths': []
+                    "error": "AI generation failed. Check provider configuration.",
+                    "paths": [],
                 }
         except Exception as e:
             logger.error(f"LLM generation failed: {e}")
-            return {
-                'error': f'LLM generation failed: {e}',
-                'paths': []
-            }
+            return {"error": f"LLM generation failed: {e}", "paths": []}
 
         # Parse response
         try:
@@ -486,6 +517,7 @@ Be specific and actionable. Think like a professional pentester. Continue from c
 
             # Score and rank paths
             from .path_scorer import PathScorer
+
             scorer = PathScorer()
 
             # Get engagement data for scoring
@@ -493,16 +525,16 @@ Be specific and actionable. Think like a professional pentester. Continue from c
             ranked_paths = scorer.rank_paths(paths, engagement_data)
 
             return {
-                'error': None,
-                'engagement_id': engagement_id,
-                'paths': ranked_paths
+                "error": None,
+                "engagement_id": engagement_id,
+                "paths": ranked_paths,
             }
         except Exception as e:
             logger.error(f"Failed to parse/score paths: {e}")
             return {
-                'error': f'Failed to process paths: {e}',
-                'paths': [],
-                'raw_response': response
+                "error": f"Failed to process paths: {e}",
+                "paths": [],
+                "raw_response": response,
             }
 
     def _build_multi_path_prompt(self, context: str, num_paths: int) -> str:
@@ -537,15 +569,15 @@ Be specific and actionable. Think like a professional pentester considering mult
         paths = []
 
         # Split by "PATH N:" pattern
-        path_blocks = re.split(r'PATH \d+:', response)
+        path_blocks = re.split(r"PATH \d+:", response)
 
         for i, block in enumerate(path_blocks[1:], 1):  # Skip first block
             path = {
-                'action': self._extract_field(block, 'ACTION'),
-                'target': self._extract_field(block, 'TARGET'),
-                'rationale': self._extract_field(block, 'RATIONALE'),
-                'expected': self._extract_field(block, 'EXPECTED'),
-                'risk_level': self._extract_risk_field(block)
+                "action": self._extract_field(block, "ACTION"),
+                "target": self._extract_field(block, "TARGET"),
+                "rationale": self._extract_field(block, "RATIONALE"),
+                "expected": self._extract_field(block, "EXPECTED"),
+                "risk_level": self._extract_risk_field(block),
             }
             paths.append(path)
 
@@ -564,8 +596,8 @@ Be specific and actionable. Think like a professional pentester considering mult
         fm = FindingsManager()
 
         return {
-            'engagement': em.get_by_id(engagement_id),
-            'hosts': hm.list_hosts(engagement_id),
-            'credentials': cm.list_credentials(engagement_id),
-            'findings': fm.list_findings(engagement_id)
+            "engagement": em.get_by_id(engagement_id),
+            "hosts": hm.list_hosts(engagement_id),
+            "credentials": cm.list_credentials(engagement_id),
+            "findings": fm.list_findings(engagement_id),
         }