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

souleyez/ai/chain_advisor.py

@@ -18,53 +18,56 @@ logger = logging.getLogger(__name__)
 
 class ChainAdvisorMode(Enum):
     """AI chain recommendation modes."""
-    OFF = 'off'           # No AI involvement
-    SUGGEST = 'suggest'   # Show recommendations, don't auto-queue
-    AUTO = 'auto'         # Auto-queue AI recommendations
+
+    OFF = "off"  # No AI involvement
+    SUGGEST = "suggest"  # Show recommendations, don't auto-queue
+    AUTO = "auto"  # Auto-queue AI recommendations
 
 
 @dataclass
 class AIChainRecommendation:
     """Single tool recommendation from LLM."""
+
     tool: str
     target: str
     args: List[str] = field(default_factory=list)
-    priority: int = 5          # 1-10
-    rationale: str = ''
-    confidence: float = 0.7    # 0.0-1.0
-    expected_outcome: str = ''
-    risk_level: str = 'medium'  # low/medium/high
+    priority: int = 5  # 1-10
+    rationale: str = ""
+    confidence: float = 0.7  # 0.0-1.0
+    expected_outcome: str = ""
+    risk_level: str = "medium"  # low/medium/high
 
 
 @dataclass
 class AIChainAnalysis:
     """Complete AI analysis of scan results."""
+
     recommendations: List[AIChainRecommendation]
-    summary: str = ''
+    summary: str = ""
     static_rules_applied: List[str] = field(default_factory=list)
     analysis_time_ms: int = 0
-    provider: str = ''
+    provider: str = ""
     error: Optional[str] = None
 
 
 # Available tools for recommendations (must match actual plugins in souleyez/plugins/)
 AVAILABLE_TOOLS = {
-    'nmap': 'Network scanning, service detection, OS fingerprinting',
-    'nuclei': 'Vulnerability scanning with templates (CVEs, misconfigs)',
-    'nikto': 'Web server vulnerability scanner',
-    'gobuster': 'Directory/file brute-forcing',
-    'ffuf': 'Web fuzzing (parameters, directories)',
-    'sqlmap': 'SQL injection detection and exploitation',
-    'wpscan': 'WordPress vulnerability scanner',
-    'hydra': 'Brute-force password cracking',
-    'enum4linux': 'SMB/NetBIOS enumeration',
-    'smbmap': 'SMB share enumeration',
-    'crackmapexec': 'Windows/AD enumeration and exploitation',
-    'dnsrecon': 'DNS enumeration and zone transfers',
-    'theharvester': 'OSINT - emails, subdomains, IPs',
-    'dalfox': 'XSS vulnerability scanner',
-    'searchsploit': 'Exploit database search',
-    'bloodhound': 'Active Directory attack path mapping',
+    "nmap": "Network scanning, service detection, OS fingerprinting",
+    "nuclei": "Vulnerability scanning with templates (CVEs, misconfigs)",
+    "nikto": "Web server vulnerability scanner",
+    "gobuster": "Directory/file brute-forcing",
+    "ffuf": "Web fuzzing (parameters, directories)",
+    "sqlmap": "SQL injection detection and exploitation",
+    "wpscan": "WordPress vulnerability scanner",
+    "hydra": "Brute-force password cracking",
+    "enum4linux": "SMB/NetBIOS enumeration",
+    "smbmap": "SMB share enumeration",
+    "crackmapexec": "Windows/AD enumeration and exploitation",
+    "dnsrecon": "DNS enumeration and zone transfers",
+    "theharvester": "OSINT - emails, subdomains, IPs",
+    "dalfox": "XSS vulnerability scanner",
+    "searchsploit": "Exploit database search",
+    "bloodhound": "Active Directory attack path mapping",
 }
 
 
@@ -121,7 +124,9 @@ Guidelines:
 - Focus on high-value services (databases, admin panels, etc.)
 """
 
-    def __init__(self, provider=None, mode: ChainAdvisorMode = ChainAdvisorMode.SUGGEST):
+    def __init__(
+        self, provider=None, mode: ChainAdvisorMode = ChainAdvisorMode.SUGGEST
+    ):
         """
         Initialize chain advisor.
 
@@ -153,7 +158,7 @@ Guidelines:
         target: str,
         parse_results: Dict[str, Any],
         static_commands: List[Dict[str, Any]],
-        engagement_id: Optional[int] = None
+        engagement_id: Optional[int] = None,
     ) -> AIChainAnalysis:
         """
         Analyze scan results and recommend next tools.
@@ -169,13 +174,13 @@ Guidelines:
             AIChainAnalysis with recommendations, or error if failed
         """
         import time
+
         start_time = time.time()
 
         # Check mode
         if self.mode == ChainAdvisorMode.OFF:
             return AIChainAnalysis(
-                recommendations=[],
-                error='AI chain analysis is disabled'
+                recommendations=[], error="AI chain analysis is disabled"
             )
 
         # Ensure provider is available
@@ -183,12 +188,13 @@ Guidelines:
         if not self.provider or not self.provider.is_available():
             logger.debug("AI provider not available for chain analysis")
             return AIChainAnalysis(
-                recommendations=[],
-                error='AI provider not available'
+                recommendations=[], error="AI provider not available"
             )
 
         # Build prompt
-        prompt = self._build_analysis_prompt(tool, target, parse_results, static_commands)
+        prompt = self._build_analysis_prompt(
+            tool, target, parse_results, static_commands
+        )
 
         # Generate recommendations
         try:
@@ -196,14 +202,13 @@ Guidelines:
                 prompt=prompt,
                 system_prompt=self.SYSTEM_PROMPT,
                 max_tokens=2000,
-                temperature=0.3  # Lower for consistent tool recommendations
+                temperature=0.3,  # Lower for consistent tool recommendations
             )
 
             if not response:
                 logger.warning("AI returned empty chain analysis")
                 return AIChainAnalysis(
-                    recommendations=[],
-                    error='Empty response from AI provider'
+                    recommendations=[], error="Empty response from AI provider"
                 )
 
             # Parse recommendations
@@ -211,98 +216,103 @@ Guidelines:
 
             # Filter out low-confidence recommendations
             min_confidence = 0.6
-            recommendations = [r for r in recommendations if r.confidence >= min_confidence]
+            recommendations = [
+                r for r in recommendations if r.confidence >= min_confidence
+            ]
 
             # Filter duplicates with static commands
             recommendations = self._filter_duplicates(recommendations, static_commands)
 
             elapsed_ms = int((time.time() - start_time) * 1000)
-            provider_name = getattr(self.provider, 'provider_type', 'unknown')
+            provider_name = getattr(self.provider, "provider_type", "unknown")
 
             return AIChainAnalysis(
                 recommendations=recommendations,
                 summary=f"AI suggested {len(recommendations)} additional tools",
-                static_rules_applied=[cmd.get('description', '') for cmd in static_commands],
+                static_rules_applied=[
+                    cmd.get("description", "") for cmd in static_commands
+                ],
                 analysis_time_ms=elapsed_ms,
-                provider=str(provider_name)
+                provider=str(provider_name),
             )
 
         except Exception as e:
             logger.error(f"AI chain analysis failed: {e}")
-            return AIChainAnalysis(
-                recommendations=[],
-                error=str(e)
-            )
+            return AIChainAnalysis(recommendations=[], error=str(e))
 
     def _build_analysis_prompt(
         self,
         tool: str,
         target: str,
         parse_results: Dict[str, Any],
-        static_commands: List[Dict[str, Any]]
+        static_commands: List[Dict[str, Any]],
     ) -> str:
         """Build LLM prompt for chain analysis."""
         # Format results summary
         results_lines = []
 
         # Add hosts
-        hosts = parse_results.get('hosts', [])
+        hosts = parse_results.get("hosts", [])
         if hosts:
             results_lines.append(f"Hosts discovered: {len(hosts)}")
             for host in hosts[:5]:  # Limit to first 5
-                ip = host.get('ip', 'unknown')
-                os_name = host.get('os', 'unknown')
+                ip = host.get("ip", "unknown")
+                os_name = host.get("os", "unknown")
                 results_lines.append(f"  - {ip} (OS: {os_name})")
 
         # Add services
-        services = parse_results.get('services', [])
+        services = parse_results.get("services", [])
         if services:
             results_lines.append(f"\nServices found: {len(services)}")
             for svc in services[:10]:  # Limit to first 10
-                ip = svc.get('ip', target)
-                port = svc.get('port', '?')
-                name = svc.get('service_name', svc.get('service', 'unknown'))
-                product = svc.get('product', '')
-                version = svc.get('version', '')
-                ver_str = f" {product} {version}".strip() if product or version else ''
+                ip = svc.get("ip", target)
+                port = svc.get("port", "?")
+                name = svc.get("service_name", svc.get("service", "unknown"))
+                product = svc.get("product", "")
+                version = svc.get("version", "")
+                ver_str = f" {product} {version}".strip() if product or version else ""
                 results_lines.append(f"  - {ip}:{port} - {name}{ver_str}")
 
         # Add findings
-        findings = parse_results.get('findings', [])
+        findings = parse_results.get("findings", [])
         if findings:
             results_lines.append(f"\nFindings: {len(findings)}")
             for finding in findings[:5]:  # Limit to first 5
-                title = finding.get('title', 'unknown')
-                severity = finding.get('severity', 'info')
+                title = finding.get("title", "unknown")
+                severity = finding.get("severity", "info")
                 results_lines.append(f"  - [{severity.upper()}] {title}")
 
         # Add credentials
-        credentials = parse_results.get('credentials', [])
+        credentials = parse_results.get("credentials", [])
         if credentials:
             results_lines.append(f"\nCredentials found: {len(credentials)}")
 
-        results_summary = '\n'.join(results_lines) if results_lines else 'No structured results parsed.'
+        results_summary = (
+            "\n".join(results_lines)
+            if results_lines
+            else "No structured results parsed."
+        )
 
         # Format static commands
         static_lines = []
         for cmd in static_commands:
-            tool_name = cmd.get('tool', 'unknown')
-            desc = cmd.get('description', '')
+            tool_name = cmd.get("tool", "unknown")
+            desc = cmd.get("description", "")
             static_lines.append(f"  - {tool_name}: {desc}")
-        static_summary = '\n'.join(static_lines) if static_lines else '  (none)'
+        static_summary = "\n".join(static_lines) if static_lines else "  (none)"
 
         # Format available tools
         tools_lines = []
         for tool_name, desc in AVAILABLE_TOOLS.items():
             tools_lines.append(f"  - {tool_name}: {desc}")
-        tools_summary = '\n'.join(tools_lines)
+        tools_summary = "\n".join(tools_lines)
 
         return self.ANALYSIS_TEMPLATE.format(
             tool=tool,
             target=target,
             results_summary=results_summary,
             static_commands=static_summary,
-            available_tools=tools_summary
+            available_tools=tools_summary,
         )
 
     def _parse_recommendations(self, response: str) -> List[AIChainRecommendation]:
@@ -310,11 +320,11 @@ Guidelines:
         recommendations = []
 
         # Check for no recommendations response
-        if 'NO_RECOMMENDATIONS' in response.upper():
+        if "NO_RECOMMENDATIONS" in response.upper():
             return []
 
         # Split by separator
-        sections = response.split('---')
+        sections = response.split("---")
 
         for section in sections:
             section = section.strip()
@@ -331,17 +341,23 @@ Guidelines:
 
         return recommendations
 
-    def _parse_single_recommendation(self, section: str) -> Optional[AIChainRecommendation]:
+    def _parse_single_recommendation(
+        self, section: str
+    ) -> Optional[AIChainRecommendation]:
         """Parse a single recommendation section."""
         # Extract fields
-        tool_match = re.search(r'TOOL:\s*(.+?)(?=\n|$)', section, re.IGNORECASE)
-        target_match = re.search(r'TARGET:\s*(.+?)(?=\n|$)', section, re.IGNORECASE)
-        args_match = re.search(r'ARGS:\s*(.+?)(?=\n|$)', section, re.IGNORECASE)
-        priority_match = re.search(r'PRIORITY:\s*(\d+)', section, re.IGNORECASE)
-        rationale_match = re.search(r'RATIONALE:\s*(.+?)(?=\n[A-Z]+:|$)', section, re.DOTALL | re.IGNORECASE)
-        confidence_match = re.search(r'CONFIDENCE:\s*([\d.]+)', section, re.IGNORECASE)
-        expected_match = re.search(r'EXPECTED:\s*(.+?)(?=\n[A-Z]+:|$)', section, re.DOTALL | re.IGNORECASE)
-        risk_match = re.search(r'RISK:\s*(\w+)', section, re.IGNORECASE)
+        tool_match = re.search(r"TOOL:\s*(.+?)(?=\n|$)", section, re.IGNORECASE)
+        target_match = re.search(r"TARGET:\s*(.+?)(?=\n|$)", section, re.IGNORECASE)
+        args_match = re.search(r"ARGS:\s*(.+?)(?=\n|$)", section, re.IGNORECASE)
+        priority_match = re.search(r"PRIORITY:\s*(\d+)", section, re.IGNORECASE)
+        rationale_match = re.search(
+            r"RATIONALE:\s*(.+?)(?=\n[A-Z]+:|$)", section, re.DOTALL | re.IGNORECASE
+        )
+        confidence_match = re.search(r"CONFIDENCE:\s*([\d.]+)", section, re.IGNORECASE)
+        expected_match = re.search(
+            r"EXPECTED:\s*(.+?)(?=\n[A-Z]+:|$)", section, re.DOTALL | re.IGNORECASE
+        )
+        risk_match = re.search(r"RISK:\s*(\w+)", section, re.IGNORECASE)
 
         # Require at least tool and target
         if not tool_match or not target_match:
@@ -359,9 +375,9 @@ Guidelines:
         args = []
         if args_match:
             args_str = args_match.group(1).strip()
-            if args_str and args_str.lower() not in ['none', 'n/a', '-']:
+            if args_str and args_str.lower() not in ["none", "n/a", "-"]:
                 # Split by comma or space
-                args = [a.strip() for a in re.split(r'[,\s]+', args_str) if a.strip()]
+                args = [a.strip() for a in re.split(r"[,\s]+", args_str) if a.strip()]
 
         # Parse priority (default 5)
         priority = 5
@@ -380,21 +396,21 @@ Guidelines:
                 pass
 
         # Parse rationale
-        rationale = ''
+        rationale = ""
         if rationale_match:
             rationale = rationale_match.group(1).strip()
 
         # Parse expected
-        expected = ''
+        expected = ""
         if expected_match:
             expected = expected_match.group(1).strip()
 
         # Parse risk
-        risk = 'medium'
+        risk = "medium"
         if risk_match:
             risk = risk_match.group(1).strip().lower()
-            if risk not in ['low', 'medium', 'high']:
-                risk = 'medium'
+            if risk not in ["low", "medium", "high"]:
+                risk = "medium"
 
         return AIChainRecommendation(
             tool=tool,
@@ -404,20 +420,20 @@ Guidelines:
             rationale=rationale,
             confidence=confidence,
             expected_outcome=expected,
-            risk_level=risk
+            risk_level=risk,
         )
 
     def _filter_duplicates(
         self,
         recommendations: List[AIChainRecommendation],
-        static_commands: List[Dict[str, Any]]
+        static_commands: List[Dict[str, Any]],
     ) -> List[AIChainRecommendation]:
         """Remove AI recommendations that duplicate static rules."""
         # Build set of (tool, target) pairs from static commands
         static_pairs = set()
         for cmd in static_commands:
-            tool = cmd.get('tool', '').lower()
-            target = cmd.get('target', '').lower()
+            tool = cmd.get("tool", "").lower()
+            target = cmd.get("target", "").lower()
             static_pairs.add((tool, target))
 
         # Filter recommendations
@@ -427,14 +443,14 @@ Guidelines:
             if pair not in static_pairs:
                 filtered.append(rec)
             else:
-                logger.debug(f"Filtered duplicate AI recommendation: {rec.tool} -> {rec.target}")
+                logger.debug(
+                    f"Filtered duplicate AI recommendation: {rec.tool} -> {rec.target}"
+                )
 
         return filtered
 
     def to_chain_commands(
-        self,
-        recommendations: List[AIChainRecommendation],
-        context: Dict[str, Any]
+        self, recommendations: List[AIChainRecommendation], context: Dict[str, Any]
     ) -> List[Dict[str, Any]]:
         """
         Convert AI recommendations to chain command format.
@@ -449,17 +465,17 @@ Guidelines:
         commands = []
         for rec in recommendations:
             cmd = {
-                'tool': rec.tool,
-                'target': rec.target,
-                'args': rec.args,
-                'priority': rec.priority,
-                'description': f"[AI] {rec.rationale[:100]}",
-                'source': 'ai_advisor',
-                'confidence': rec.confidence,
-                'risk_level': rec.risk_level
+                "tool": rec.tool,
+                "target": rec.target,
+                "args": rec.args,
+                "priority": rec.priority,
+                "description": f"[AI] {rec.rationale[:100]}",
+                "source": "ai_advisor",
+                "confidence": rec.confidence,
+                "risk_level": rec.risk_level,
             }
             commands.append(cmd)
 
         # Sort by priority (highest first)
-        commands.sort(key=lambda c: c['priority'], reverse=True)
+        commands.sort(key=lambda c: c["priority"], reverse=True)
         return commands

souleyez/ai/claude_provider.py

@@ -4,6 +4,7 @@ souleyez.ai.claude_provider - Anthropic Claude API provider implementation
 Provides Claude LLM integration for high-quality AI report generation.
 Requires API key stored securely via CryptoManager.
 """
+
 import logging
 from typing import Optional, Dict, Any
 
@@ -14,6 +15,7 @@ logger = logging.getLogger(__name__)
 # Check if anthropic package is available
 try:
     import anthropic
+
     ANTHROPIC_AVAILABLE = True
 except ImportError:
     anthropic = None
@@ -30,11 +32,7 @@ class ClaudeProvider(LLMProvider):
 
     DEFAULT_MODEL = "claude-sonnet-4-20250514"
 
-    def __init__(
-        self,
-        api_key: Optional[str] = None,
-        model: Optional[str] = None
-    ):
+    def __init__(self, api_key: Optional[str] = None, model: Optional[str] = None):
         """
         Initialize Claude provider.
 
@@ -54,7 +52,7 @@ class ClaudeProvider(LLMProvider):
         from souleyez.config import get
         from souleyez.storage.crypto import get_crypto_manager
 
-        encrypted_key = get('ai.claude_api_key')
+        encrypted_key = get("ai.claude_api_key")
         if not encrypted_key:
             logger.debug("No Claude API key configured")
             return None
@@ -74,7 +72,8 @@ class ClaudeProvider(LLMProvider):
     def _load_model(self) -> str:
         """Load model from config."""
         from souleyez.config import get
-        return get('ai.claude_model', self.DEFAULT_MODEL)
+
+        return get("ai.claude_model", self.DEFAULT_MODEL)
 
     def _get_client(self):
         """Get or create Anthropic client."""
@@ -86,7 +85,9 @@ class ClaudeProvider(LLMProvider):
             return None
 
         if not ANTHROPIC_AVAILABLE:
-            logger.error("anthropic package not installed. Install with: pip install anthropic")
+            logger.error(
+                "anthropic package not installed. Install with: pip install anthropic"
+            )
             return None
 
         try:
@@ -131,7 +132,7 @@ class ClaudeProvider(LLMProvider):
         prompt: str,
         system_prompt: Optional[str] = None,
         max_tokens: int = 4096,
-        temperature: float = 0.3
+        temperature: float = 0.3,
     ) -> Optional[str]:
         """
         Generate text using Claude API.
@@ -159,7 +160,7 @@ class ClaudeProvider(LLMProvider):
                 "model": self._model,
                 "max_tokens": max_tokens,
                 "temperature": temperature,
-                "messages": messages
+                "messages": messages,
             }
 
             if system_prompt:
@@ -198,33 +199,33 @@ class ClaudeProvider(LLMProvider):
         has_key = api_key is not None
 
         status = {
-            'provider_type': self.provider_type.value,
-            'provider': 'Claude',
-            'provider_name': 'Claude (Anthropic)',
-            'model': self._model,
-            'api_key_configured': has_key,
-            'package_installed': ANTHROPIC_AVAILABLE,
-            'connected': False,
-            'error': None
+            "provider_type": self.provider_type.value,
+            "provider": "Claude",
+            "provider_name": "Claude (Anthropic)",
+            "model": self._model,
+            "api_key_configured": has_key,
+            "package_installed": ANTHROPIC_AVAILABLE,
+            "connected": False,
+            "error": None,
         }
 
         if not ANTHROPIC_AVAILABLE:
-            status['error'] = 'anthropic package not installed'
+            status["error"] = "anthropic package not installed"
             return status
 
         if not has_key:
-            status['error'] = 'No API key configured'
+            status["error"] = "No API key configured"
             return status
 
         # Test connection
         try:
             client = self._get_client()
             if client:
-                status['connected'] = True
+                status["connected"] = True
             else:
-                status['error'] = 'Failed to create client'
+                status["error"] = "Failed to create client"
         except Exception as e:
-            status['error'] = str(e)
+            status["error"] = str(e)
 
         return status
 
@@ -254,9 +255,9 @@ def set_claude_api_key(api_key: str) -> bool:
             return False
 
         cfg = read_config()
-        if 'ai' not in cfg:
-            cfg['ai'] = {}
-        cfg['ai']['claude_api_key'] = encrypted_key
+        if "ai" not in cfg:
+            cfg["ai"] = {}
+        cfg["ai"]["claude_api_key"] = encrypted_key
         write_config(cfg)
 
         logger.info("Claude API key stored successfully")
@@ -277,8 +278,8 @@ def clear_claude_api_key() -> bool:
 
     try:
         cfg = read_config()
-        if 'ai' in cfg and 'claude_api_key' in cfg['ai']:
-            del cfg['ai']['claude_api_key']
+        if "ai" in cfg and "claude_api_key" in cfg["ai"]:
+            del cfg["ai"]["claude_api_key"]
             write_config(cfg)
             logger.info("Claude API key removed")
         return True