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

souleyez/ai/report_context.py

@@ -3,6 +3,7 @@ souleyez.ai.report_context - Build context for AI report generation
 
 Prepares engagement data in formats suitable for LLM prompt templates.
 """
+
 import logging
 from typing import Dict, Any, List, Optional
 
@@ -48,44 +49,44 @@ class ReportContextBuilder:
         creds = self.cm.list_credentials(engagement_id)
 
         # Count findings by severity
-        severity_counts = {'critical': 0, 'high': 0, 'medium': 0, 'low': 0, 'info': 0}
+        severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
         for f in findings:
-            sev = f.get('severity', 'info').lower()
+            sev = f.get("severity", "info").lower()
             if sev in severity_counts:
                 severity_counts[sev] += 1
             else:
-                severity_counts['info'] += 1
+                severity_counts["info"] += 1
 
         # Count hosts with critical/high findings
         hosts_with_issues = set()
         for f in findings:
-            if f.get('severity', '').lower() in ['critical', 'high']:
-                if f.get('host_id'):
-                    hosts_with_issues.add(f['host_id'])
+            if f.get("severity", "").lower() in ["critical", "high"]:
+                if f.get("host_id"):
+                    hosts_with_issues.add(f["host_id"])
 
         # Build top findings summary
         top_findings = self._format_top_findings(findings, limit=5)
 
         # Engagement type and duration
-        eng_type = engagement.get('type', 'Penetration Test')
+        eng_type = engagement.get("type", "Penetration Test")
         duration = self._calculate_duration(engagement)
-        scope = engagement.get('description', 'Network and application assessment')
+        scope = engagement.get("description", "Network and application assessment")
 
         return {
-            'engagement_name': engagement.get('name', 'Unknown'),
-            'engagement_type': eng_type,
-            'duration': duration,
-            'scope_summary': scope[:200] if scope else 'Full scope assessment',
-            'total_findings': len(findings),
-            'critical_count': severity_counts['critical'],
-            'high_count': severity_counts['high'],
-            'medium_count': severity_counts['medium'],
-            'low_count': severity_counts['low'],
-            'info_count': severity_counts['info'],
-            'total_hosts': len(hosts),
-            'compromised_hosts': len(hosts_with_issues),
-            'credentials_count': len(creds),
-            'top_findings': top_findings,
+            "engagement_name": engagement.get("name", "Unknown"),
+            "engagement_type": eng_type,
+            "duration": duration,
+            "scope_summary": scope[:200] if scope else "Full scope assessment",
+            "total_findings": len(findings),
+            "critical_count": severity_counts["critical"],
+            "high_count": severity_counts["high"],
+            "medium_count": severity_counts["medium"],
+            "low_count": severity_counts["low"],
+            "info_count": severity_counts["info"],
+            "total_hosts": len(hosts),
+            "compromised_hosts": len(hosts_with_issues),
+            "credentials_count": len(creds),
+            "top_findings": top_findings,
         }
 
     def build_finding_context(self, finding: Dict[str, Any]) -> Dict[str, Any]:
@@ -99,35 +100,35 @@ class ReportContextBuilder:
             dict: Context variables for FINDING_ENHANCEMENT_PROMPT
         """
         # Get host info if available
-        host_ip = 'N/A'
-        hostname = ''
-        if finding.get('host_id'):
-            host = self.hm.get_host(finding['host_id'])
+        host_ip = "N/A"
+        hostname = ""
+        if finding.get("host_id"):
+            host = self.hm.get_host(finding["host_id"])
             if host:
-                host_ip = host.get('ip_address', 'N/A')
-                hostname = host.get('hostname', '')
+                host_ip = host.get("ip_address", "N/A")
+                hostname = host.get("hostname", "")
 
         # Extract CVE/CWE
-        cve = finding.get('refs', '') or 'N/A'
+        cve = finding.get("refs", "") or "N/A"
         if isinstance(cve, list):
-            cve = ', '.join(cve[:3])  # Limit to first 3
+            cve = ", ".join(cve[:3])  # Limit to first 3
 
         # Clean evidence (truncate if too long)
-        evidence = finding.get('evidence', '')
+        evidence = finding.get("evidence", "")
         if len(evidence) > 1000:
-            evidence = evidence[:1000] + '\n... [truncated]'
+            evidence = evidence[:1000] + "\n... [truncated]"
 
         return {
-            'title': finding.get('title', 'Unknown Finding'),
-            'severity': finding.get('severity', 'Unknown').upper(),
-            'host': host_ip,
-            'hostname': hostname,
-            'port': finding.get('port', 'N/A'),
-            'service': finding.get('service', 'Unknown'),
-            'tool': finding.get('tool', 'Manual'),
-            'description': finding.get('description', 'No description provided'),
-            'cve': cve,
-            'evidence': evidence or 'No evidence recorded',
+            "title": finding.get("title", "Unknown Finding"),
+            "severity": finding.get("severity", "Unknown").upper(),
+            "host": host_ip,
+            "hostname": hostname,
+            "port": finding.get("port", "N/A"),
+            "service": finding.get("service", "Unknown"),
+            "tool": finding.get("tool", "Manual"),
+            "description": finding.get("description", "No description provided"),
+            "cve": cve,
+            "evidence": evidence or "No evidence recorded",
         }
 
     def build_remediation_context(self, engagement_id: int) -> Dict[str, Any]:
@@ -145,9 +146,9 @@ class ReportContextBuilder:
         creds = self.cm.list_credentials(engagement_id)
 
         # Count by severity
-        severity_counts = {'critical': 0, 'high': 0, 'medium': 0, 'low': 0, 'info': 0}
+        severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
         for f in findings:
-            sev = f.get('severity', 'info').lower()
+            sev = f.get("severity", "info").lower()
             if sev in severity_counts:
                 severity_counts[sev] += 1
 
@@ -158,61 +159,59 @@ class ReportContextBuilder:
         top_vulns = self._format_top_findings(findings, limit=10, include_medium=True)
 
         return {
-            'findings_summary': findings_summary,
-            'total_hosts': len(hosts),
-            'critical_count': severity_counts['critical'],
-            'high_count': severity_counts['high'],
-            'medium_count': severity_counts['medium'],
-            'creds_count': len(creds),
-            'top_vulnerabilities': top_vulns,
+            "findings_summary": findings_summary,
+            "total_hosts": len(hosts),
+            "critical_count": severity_counts["critical"],
+            "high_count": severity_counts["high"],
+            "medium_count": severity_counts["medium"],
+            "creds_count": len(creds),
+            "top_vulnerabilities": top_vulns,
         }
 
     def _format_top_findings(
-        self,
-        findings: List[Dict],
-        limit: int = 5,
-        include_medium: bool = False
+        self, findings: List[Dict], limit: int = 5, include_medium: bool = False
     ) -> str:
         """Format top critical/high findings as text."""
-        severity_order = {'critical': 0, 'high': 1, 'medium': 2, 'low': 3, 'info': 4}
+        severity_order = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4}
 
         # Filter to critical/high (and optionally medium)
         max_severity = 2 if include_medium else 1
         priority_findings = [
-            f for f in findings
-            if severity_order.get(f.get('severity', 'info').lower(), 4) <= max_severity
+            f
+            for f in findings
+            if severity_order.get(f.get("severity", "info").lower(), 4) <= max_severity
         ]
 
         # Sort by severity
         priority_findings.sort(
-            key=lambda f: severity_order.get(f.get('severity', 'info').lower(), 4)
+            key=lambda f: severity_order.get(f.get("severity", "info").lower(), 4)
         )
 
         # Format
         lines = []
         for f in priority_findings[:limit]:
-            sev = f.get('severity', 'unknown').upper()
-            title = f.get('title', 'Unknown')
-            desc = f.get('description', '')[:100]
+            sev = f.get("severity", "unknown").upper()
+            title = f.get("title", "Unknown")
+            desc = f.get("description", "")[:100]
             lines.append(f"- [{sev}] {title}")
             if desc:
                 lines.append(f"  {desc}...")
 
-        return '\n'.join(lines) if lines else 'No critical or high severity findings.'
+        return "\n".join(lines) if lines else "No critical or high severity findings."
 
     def _format_findings_by_severity(self, findings: List[Dict]) -> str:
         """Format all findings grouped by severity."""
-        by_severity = {'critical': [], 'high': [], 'medium': [], 'low': [], 'info': []}
+        by_severity = {"critical": [], "high": [], "medium": [], "low": [], "info": []}
 
         for f in findings:
-            sev = f.get('severity', 'info').lower()
+            sev = f.get("severity", "info").lower()
             if sev in by_severity:
                 by_severity[sev].append(f)
             else:
-                by_severity['info'].append(f)
+                by_severity["info"].append(f)
 
         lines = []
-        for sev in ['critical', 'high', 'medium', 'low']:
+        for sev in ["critical", "high", "medium", "low"]:
             items = by_severity[sev]
             if items:
                 lines.append(f"\n{sev.upper()} ({len(items)}):")
@@ -221,20 +220,22 @@ class ReportContextBuilder:
                 if len(items) > 5:
                     lines.append(f"  ... and {len(items) - 5} more")
 
-        return '\n'.join(lines) if lines else 'No findings recorded.'
+        return "\n".join(lines) if lines else "No findings recorded."
 
     def _calculate_duration(self, engagement: Dict) -> str:
         """Calculate engagement duration from dates."""
         from datetime import datetime
 
-        start = engagement.get('start_date')
-        end = engagement.get('end_date')
+        start = engagement.get("start_date")
+        end = engagement.get("end_date")
 
         if not start:
-            return 'Duration not specified'
+            return "Duration not specified"
 
         try:
-            start_dt = datetime.fromisoformat(start) if isinstance(start, str) else start
+            start_dt = (
+                datetime.fromisoformat(start) if isinstance(start, str) else start
+            )
             if end:
                 end_dt = datetime.fromisoformat(end) if isinstance(end, str) else end
                 days = (end_dt - start_dt).days
@@ -242,4 +243,4 @@ class ReportContextBuilder:
             else:
                 return "Ongoing"
         except Exception:
-            return 'Duration not specified'
+            return "Duration not specified"

souleyez/ai/report_service.py

@@ -4,6 +4,7 @@ souleyez.ai.report_service - AI-powered report generation service
 Provides methods for generating AI-enhanced report sections using
 configured LLM providers (Claude or Ollama).
 """
+
 import logging
 import re
 import concurrent.futures
@@ -83,13 +84,11 @@ class AIReportService:
     def get_provider_info(self) -> Dict[str, Any]:
         """Get information about the current provider."""
         if not self.provider:
-            return {'available': False, 'error': 'No provider configured'}
+            return {"available": False, "error": "No provider configured"}
         return self.provider.get_status()
 
     def generate_executive_summary(
-        self,
-        engagement_id: int,
-        max_tokens: int = 1500
+        self, engagement_id: int, max_tokens: int = 1500
     ) -> Optional[str]:
         """
         Generate AI-powered executive summary.
@@ -118,13 +117,15 @@ class AIReportService:
                     prompt=prompt,
                     system_prompt=REPORT_SYSTEM_PROMPT,
                     max_tokens=max_tokens,
-                    temperature=0.3
+                    temperature=0.3,
                 )
 
             result = _run_with_timeout(_generate)
 
             if result:
-                logger.info(f"Generated executive summary for engagement {engagement_id}")
+                logger.info(
+                    f"Generated executive summary for engagement {engagement_id}"
+                )
             return result
 
         except Exception as e:
@@ -132,9 +133,7 @@ class AIReportService:
             return None
 
     def enhance_finding(
-        self,
-        finding: Dict[str, Any],
-        max_tokens: int = 800
+        self, finding: Dict[str, Any], max_tokens: int = 800
     ) -> Optional[Dict[str, str]]:
         """
         Enhance a single finding with business context.
@@ -160,7 +159,7 @@ class AIReportService:
                     prompt=prompt,
                     system_prompt=REPORT_SYSTEM_PROMPT,
                     max_tokens=max_tokens,
-                    temperature=0.3
+                    temperature=0.3,
                 )
 
             result = _run_with_timeout(_generate)
@@ -174,9 +173,7 @@ class AIReportService:
             return None
 
     def generate_remediation_plan(
-        self,
-        engagement_id: int,
-        max_tokens: int = 2500
+        self, engagement_id: int, max_tokens: int = 2500
     ) -> Optional[str]:
         """
         Generate prioritized remediation plan.
@@ -205,23 +202,22 @@ class AIReportService:
                     prompt=prompt,
                     system_prompt=REPORT_SYSTEM_PROMPT,
                     max_tokens=max_tokens,
-                    temperature=0.3
+                    temperature=0.3,
                 )
 
             result = _run_with_timeout(_generate)
 
             if result:
-                logger.info(f"Generated remediation plan for engagement {engagement_id}")
+                logger.info(
+                    f"Generated remediation plan for engagement {engagement_id}"
+                )
             return result
 
         except Exception as e:
             logger.error(f"Failed to generate remediation plan: {e}")
             return None
 
-    def generate_risk_rating(
-        self,
-        engagement_id: int
-    ) -> Optional[Dict[str, str]]:
+    def generate_risk_rating(self, engagement_id: int) -> Optional[Dict[str, str]]:
         """
         Generate overall risk rating.
 
@@ -251,7 +247,7 @@ Credentials Compromised: {context.get('credentials_count', 0)}
                     prompt=prompt,
                     system_prompt=REPORT_SYSTEM_PROMPT,
                     max_tokens=200,
-                    temperature=0.2
+                    temperature=0.2,
                 )
 
             result = _run_with_timeout(_generate)
@@ -265,10 +261,7 @@ Credentials Compromised: {context.get('credentials_count', 0)}
             return None
 
     def generate_all_content(
-        self,
-        engagement_id: int,
-        enhance_findings: bool = True,
-        max_findings: int = 10
+        self, engagement_id: int, enhance_findings: bool = True, max_findings: int = 10
     ) -> Dict[str, Any]:
         """
         Generate all AI content for a report.
@@ -284,37 +277,39 @@ Credentials Compromised: {context.get('credentials_count', 0)}
         from souleyez.storage.findings import FindingsManager
 
         content = {
-            'executive_summary': None,
-            'remediation_plan': None,
-            'risk_rating': None,
-            'enhanced_findings': {},
-            'provider': None,
-            'errors': []
+            "executive_summary": None,
+            "remediation_plan": None,
+            "risk_rating": None,
+            "enhanced_findings": {},
+            "provider": None,
+            "errors": [],
         }
 
         if not self.is_available():
-            content['errors'].append('AI provider not available')
+            content["errors"].append("AI provider not available")
             return content
 
-        content['provider'] = self.provider.provider_type.value
+        content["provider"] = self.provider.provider_type.value
 
         # Generate executive summary
         try:
-            content['executive_summary'] = self.generate_executive_summary(engagement_id)
+            content["executive_summary"] = self.generate_executive_summary(
+                engagement_id
+            )
         except Exception as e:
-            content['errors'].append(f'Executive summary: {e}')
+            content["errors"].append(f"Executive summary: {e}")
 
         # Generate remediation plan
         try:
-            content['remediation_plan'] = self.generate_remediation_plan(engagement_id)
+            content["remediation_plan"] = self.generate_remediation_plan(engagement_id)
         except Exception as e:
-            content['errors'].append(f'Remediation plan: {e}')
+            content["errors"].append(f"Remediation plan: {e}")
 
         # Generate risk rating
         try:
-            content['risk_rating'] = self.generate_risk_rating(engagement_id)
+            content["risk_rating"] = self.generate_risk_rating(engagement_id)
         except Exception as e:
-            content['errors'].append(f'Risk rating: {e}')
+            content["errors"].append(f"Risk rating: {e}")
 
         # Enhance individual findings (top critical/high only)
         if enhance_findings:
@@ -323,38 +318,48 @@ Credentials Compromised: {context.get('credentials_count', 0)}
                 findings = fm.list_findings(engagement_id)
 
                 # Sort by severity and take top N
-                severity_order = {'critical': 0, 'high': 1, 'medium': 2, 'low': 3, 'info': 4}
+                severity_order = {
+                    "critical": 0,
+                    "high": 1,
+                    "medium": 2,
+                    "low": 3,
+                    "info": 4,
+                }
                 priority_findings = sorted(
-                    [f for f in findings if f.get('severity', '').lower() in ['critical', 'high']],
-                    key=lambda f: severity_order.get(f.get('severity', 'info').lower(), 4)
+                    [
+                        f
+                        for f in findings
+                        if f.get("severity", "").lower() in ["critical", "high"]
+                    ],
+                    key=lambda f: severity_order.get(
+                        f.get("severity", "info").lower(), 4
+                    ),
                 )[:max_findings]
 
                 for finding in priority_findings:
                     try:
                         enhanced = self.enhance_finding(finding)
                         if enhanced:
-                            content['enhanced_findings'][finding['id']] = enhanced
+                            content["enhanced_findings"][finding["id"]] = enhanced
                     except Exception as e:
-                        logger.warning(f"Failed to enhance finding {finding.get('id')}: {e}")
+                        logger.warning(
+                            f"Failed to enhance finding {finding.get('id')}: {e}"
+                        )
 
             except Exception as e:
-                content['errors'].append(f'Finding enhancement: {e}')
+                content["errors"].append(f"Finding enhancement: {e}")
 
         return content
 
     def _parse_finding_enhancement(self, response: str) -> Dict[str, str]:
         """Parse LLM response into structured finding enhancement."""
-        result = {
-            'business_impact': '',
-            'attack_scenario': '',
-            'risk_context': ''
-        }
+        result = {"business_impact": "", "attack_scenario": "", "risk_context": ""}
 
         # Try to extract sections
         sections = {
-            'business_impact': r'BUSINESS IMPACT[:\s]*(.+?)(?=ATTACK SCENARIO|RISK CONTEXT|$)',
-            'attack_scenario': r'ATTACK SCENARIO[:\s]*(.+?)(?=RISK CONTEXT|BUSINESS IMPACT|$)',
-            'risk_context': r'RISK CONTEXT[:\s]*(.+?)(?=BUSINESS IMPACT|ATTACK SCENARIO|$)'
+            "business_impact": r"BUSINESS IMPACT[:\s]*(.+?)(?=ATTACK SCENARIO|RISK CONTEXT|$)",
+            "attack_scenario": r"ATTACK SCENARIO[:\s]*(.+?)(?=RISK CONTEXT|BUSINESS IMPACT|$)",
+            "risk_context": r"RISK CONTEXT[:\s]*(.+?)(?=BUSINESS IMPACT|ATTACK SCENARIO|$)",
         }
 
         for key, pattern in sections.items():
@@ -364,19 +369,23 @@ Credentials Compromised: {context.get('credentials_count', 0)}
 
         # If parsing failed, use the whole response as business impact
         if not any(result.values()):
-            result['business_impact'] = response.strip()
+            result["business_impact"] = response.strip()
 
         return result
 
     def _parse_risk_rating(self, response: str) -> Dict[str, str]:
         """Parse risk rating response."""
-        result = {'rating': 'UNKNOWN', 'justification': response}
+        result = {"rating": "UNKNOWN", "justification": response}
 
         # Look for rating pattern
-        match = re.search(r'RATING:\s*(CRITICAL|HIGH|MODERATE|LOW)\s*[-–]\s*(.+)', response, re.IGNORECASE)
+        match = re.search(
+            r"RATING:\s*(CRITICAL|HIGH|MODERATE|LOW)\s*[-–]\s*(.+)",
+            response,
+            re.IGNORECASE,
+        )
         if match:
-            result['rating'] = match.group(1).upper()
-            result['justification'] = match.group(2).strip()
+            result["rating"] = match.group(1).upper()
+            result["justification"] = match.group(2).strip()
 
         return result
 
@@ -404,18 +413,27 @@ Credentials Compromised: {context.get('credentials_count', 0)}
         remediation_input = len(REMEDIATION_PLAN_PROMPT.format(**remediation_ctx)) // 4
         remediation_output = 800
 
-        critical_high = len([f for f in findings if f.get('severity', '').lower() in ['critical', 'high']])
+        critical_high = len(
+            [
+                f
+                for f in findings
+                if f.get("severity", "").lower() in ["critical", "high"]
+            ]
+        )
         findings_input = critical_high * 300
         findings_output = critical_high * 200
 
         return {
-            'executive_summary': executive_input + executive_output,
-            'remediation_plan': remediation_input + remediation_output,
-            'findings_enhancement': findings_input + findings_output,
-            'total_estimated': (
-                executive_input + executive_output +
-                remediation_input + remediation_output +
-                findings_input + findings_output
+            "executive_summary": executive_input + executive_output,
+            "remediation_plan": remediation_input + remediation_output,
+            "findings_enhancement": findings_input + findings_output,
+            "total_estimated": (
+                executive_input
+                + executive_output
+                + remediation_input
+                + remediation_output
+                + findings_input
+                + findings_output
             ),
-            'findings_to_enhance': min(critical_high, 10)
+            "findings_to_enhance": min(critical_high, 10),
         }