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

souleyez/intelligence/exploit_suggestions.py

@@ -30,15 +30,15 @@ def _generate_exploit_identifier(exploit: Dict) -> str:
         Unique identifier string like "msf:exploit/unix/ftp/vsftpd_234_backdoor"
         or "cve:CVE-2011-2523" or "edb:49757" or "title:<hash>"
     """
-    if exploit.get('msf_module'):
+    if exploit.get("msf_module"):
         return f"msf:{exploit['msf_module']}"
-    elif exploit.get('cve'):
+    elif exploit.get("cve"):
         return f"cve:{exploit['cve']}"
-    elif exploit.get('edb_id'):
+    elif exploit.get("edb_id"):
         return f"edb:{exploit['edb_id']}"
     else:
         # Fallback: use title (should be rare)
-        title_hash = hash(exploit.get('title', 'unknown'))
+        title_hash = hash(exploit.get("title", "unknown"))
         return f"title:{title_hash}"
 
 
@@ -54,12 +54,12 @@ class ExploitSuggestionEngine:
         """
         self.version_extractor = ServiceVersionExtractor()
         self.knowledge_base = ExploitKnowledgeBase(use_searchsploit=use_searchsploit)
-    
+
     def generate_suggestions(self, engagement_id: int, host_id: int = None) -> Dict:
         """
         Generate exploit suggestions for engagement or specific host.
         Results cached for 30 seconds to improve dashboard performance.
-        
+
         Returns:
             {
                 'hosts': [
@@ -100,7 +100,7 @@ class ExploitSuggestionEngine:
             cached_result, cached_time = _SUGGESTION_CACHE[cache_key]
             if time.time() - cached_time < _CACHE_TIMEOUT:
                 return cached_result
-        
+
         # Cache miss or expired - do the analysis
         from souleyez.storage.hosts import HostManager
         from souleyez.storage.findings import FindingsManager
@@ -118,42 +118,46 @@ class ExploitSuggestionEngine:
             hosts = [hm.get_host(host_id)]
         else:
             hosts = hm.list_hosts(engagement_id)
-        
+
+        # Fetch all findings ONCE (not per-host - N+1 query fix)
+        all_findings = fm.list_findings(engagement_id)
+
         results = []
-        
+
         for host in hosts:
             if not host:
                 continue
 
-            h_id = host.get('id')
-            host_ip = host.get('ip') or host.get('ip_address')
+            h_id = host.get("id")
+            host_ip = host.get("ip") or host.get("ip_address")
 
-            # Get findings for this host
-            all_findings = fm.list_findings(engagement_id)
-            host_findings = [
-                f for f in all_findings
-                if f.get('host') == host_ip
-            ]
+            # Filter findings for this host
+            host_findings = [f for f in all_findings if f.get("host") == host_ip]
 
             # Get Wazuh vulnerabilities for this host
             host_wazuh_vulns = [
-                v for v in wazuh_vulns
-                if v.get('host_id') == h_id or v.get('agent_ip') == host_ip or v.get('host_ip') == host_ip
+                v
+                for v in wazuh_vulns
+                if v.get("host_id") == h_id
+                or v.get("agent_ip") == host_ip
+                or v.get("host_ip") == host_ip
             ]
 
             # Build set of CVEs found by scans (from findings)
             scan_cves = set()
             for finding in host_findings:
                 # Extract CVE from finding title or refs
-                title = finding.get('title', '')
-                refs = finding.get('refs', '')
+                title = finding.get("title", "")
+                refs = finding.get("refs", "")
                 for text in [title, refs]:
-                    cve_match = re.search(r'CVE-\d{4}-\d{4,}', text, re.IGNORECASE)
+                    cve_match = re.search(r"CVE-\d{4}-\d{4,}", text, re.IGNORECASE)
                     if cve_match:
                         scan_cves.add(cve_match.group(0).upper())
 
             # Build set of CVEs from Wazuh
-            wazuh_cves = {v.get('cve_id').upper() for v in host_wazuh_vulns if v.get('cve_id')}
+            wazuh_cves = {
+                v.get("cve_id").upper() for v in host_wazuh_vulns if v.get("cve_id")
+            }
 
             # Extract versions from all sources
             services = self.version_extractor.combine_sources(h_id, host_findings)
@@ -164,52 +168,48 @@ class ExploitSuggestionEngine:
 
             for svc in services:
                 exploits = self.knowledge_base.find_exploits(
-                    svc['service'],
-                    svc['version_number'],
-                    svc.get('product')
+                    svc["service"], svc["version_number"], svc.get("product")
                 )
 
                 # Enrich exploits with attempt status and detection source
                 for exploit in exploits:
                     exploit_id = _generate_exploit_identifier(exploit)
-                    exploit['identifier'] = exploit_id
+                    exploit["identifier"] = exploit_id
 
                     # Look up attempt status
                     status = exploit_attempts.get_attempt_status(
                         engagement_id=engagement_id,
                         host_id=h_id,
                         exploit_identifier=exploit_id,
-                        service_id=svc.get('service_id')
+                        service_id=svc.get("service_id"),
                     )
-                    exploit['attempt_status'] = status or 'not_tried'
+                    exploit["attempt_status"] = status or "not_tried"
 
                     # Set detection source
-                    cve = exploit.get('cve', '').upper() if exploit.get('cve') else None
+                    cve = exploit.get("cve", "").upper() if exploit.get("cve") else None
                     if cve:
                         seen_cves_in_services.add(cve)
                         in_scan = cve in scan_cves
                         in_wazuh = cve in wazuh_cves
                         if in_scan and in_wazuh:
-                            exploit['detection_source'] = 'both'
+                            exploit["detection_source"] = "both"
                         elif in_wazuh:
-                            exploit['detection_source'] = 'wazuh'
+                            exploit["detection_source"] = "wazuh"
                         else:
-                            exploit['detection_source'] = 'scan'
+                            exploit["detection_source"] = "scan"
                     else:
-                        exploit['detection_source'] = 'scan'
+                        exploit["detection_source"] = "scan"
 
-                techniques = self.knowledge_base.get_techniques(svc['service'])
+                techniques = self.knowledge_base.get_techniques(svc["service"])
 
-                enriched_services.append({
-                    **svc,
-                    'exploits': exploits,
-                    'techniques': techniques
-                })
+                enriched_services.append(
+                    {**svc, "exploits": exploits, "techniques": techniques}
+                )
 
             # Add Wazuh-only CVE exploits (not already covered by service exploits)
             wazuh_only_exploits = []
             for wazuh_vuln in host_wazuh_vulns:
-                cve_id = wazuh_vuln.get('cve_id')
+                cve_id = wazuh_vuln.get("cve_id")
                 if not cve_id:
                     continue
                 cve_upper = cve_id.upper()
@@ -220,69 +220,73 @@ class ExploitSuggestionEngine:
                 cve_exploits = self.knowledge_base.find_exploits_by_cve(cve_id)
                 for exploit in cve_exploits:
                     exploit_id = _generate_exploit_identifier(exploit)
-                    exploit['identifier'] = exploit_id
+                    exploit["identifier"] = exploit_id
 
                     # Look up attempt status
                     status = exploit_attempts.get_attempt_status(
                         engagement_id=engagement_id,
                         host_id=h_id,
-                        exploit_identifier=exploit_id
+                        exploit_identifier=exploit_id,
                     )
-                    exploit['attempt_status'] = status or 'not_tried'
+                    exploit["attempt_status"] = status or "not_tried"
 
                     # Set detection source
                     in_scan = cve_upper in scan_cves
                     if in_scan:
-                        exploit['detection_source'] = 'both'
+                        exploit["detection_source"] = "both"
                     else:
-                        exploit['detection_source'] = 'wazuh'
+                        exploit["detection_source"] = "wazuh"
 
                     # Add package context from Wazuh
-                    exploit['wazuh_package'] = wazuh_vuln.get('package_name')
-                    exploit['wazuh_version'] = wazuh_vuln.get('package_version')
+                    exploit["wazuh_package"] = wazuh_vuln.get("package_name")
+                    exploit["wazuh_version"] = wazuh_vuln.get("package_version")
 
                     wazuh_only_exploits.append(exploit)
                     seen_cves_in_services.add(cve_upper)
 
             # Add Wazuh-only exploits as a pseudo-service if any exist
             if wazuh_only_exploits:
-                enriched_services.append({
-                    'port': None,
-                    'service': 'wazuh_detected',
-                    'version': None,
-                    'product': None,
-                    'version_number': None,
-                    'exploits': wazuh_only_exploits,
-                    'techniques': [],
-                    'source': 'wazuh'
-                })
+                enriched_services.append(
+                    {
+                        "port": None,
+                        "service": "wazuh_detected",
+                        "version": None,
+                        "product": None,
+                        "version_number": None,
+                        "exploits": wazuh_only_exploits,
+                        "techniques": [],
+                        "source": "wazuh",
+                    }
+                )
 
             # Sort all exploits to prioritize 'both' detection source
             for svc in enriched_services:
-                if svc.get('exploits'):
-                    svc['exploits'] = self._prioritize_exploits(svc['exploits'])
+                if svc.get("exploits"):
+                    svc["exploits"] = self._prioritize_exploits(svc["exploits"])
 
             # Only include hosts with exploits or techniques
-            if any(s.get('exploits') or s.get('techniques') for s in enriched_services):
-                results.append({
-                    'host_id': h_id,
-                    'ip': host_ip,
-                    'hostname': host.get('hostname'),
-                    'services': enriched_services,
-                    'wazuh_cve_count': len(wazuh_cves)
-                })
-        
-        result = {'hosts': results}
-        
+            if any(s.get("exploits") or s.get("techniques") for s in enriched_services):
+                results.append(
+                    {
+                        "host_id": h_id,
+                        "ip": host_ip,
+                        "hostname": host.get("hostname"),
+                        "services": enriched_services,
+                        "wazuh_cve_count": len(wazuh_cves),
+                    }
+                )
+
+        result = {"hosts": results}
+
         # Cache the result
         _SUGGESTION_CACHE[cache_key] = (result, time.time())
-        
+
         return result
-    
+
     def get_host_summary(self, engagement_id: int, host_id: int) -> Dict:
         """
         Get exploit summary for a specific host.
-        
+
         Returns:
             {
                 'total_exploits': 5,
@@ -294,56 +298,62 @@ class ExploitSuggestionEngine:
             }
         """
         suggestions = self.generate_suggestions(engagement_id, host_id)
-        
-        if not suggestions['hosts']:
+
+        if not suggestions["hosts"]:
             return {
-                'total_exploits': 0,
-                'critical': 0,
-                'high': 0,
-                'medium': 0,
-                'low': 0,
-                'services_with_exploits': 0,
-                'top_exploits': []
+                "total_exploits": 0,
+                "critical": 0,
+                "high": 0,
+                "medium": 0,
+                "low": 0,
+                "services_with_exploits": 0,
+                "top_exploits": [],
             }
-        
-        host_data = suggestions['hosts'][0]
+
+        host_data = suggestions["hosts"][0]
         all_exploits = []
         services_with_exploits = 0
-        
-        for svc in host_data['services']:
-            if svc.get('exploits'):
+
+        for svc in host_data["services"]:
+            if svc.get("exploits"):
                 services_with_exploits += 1
-                all_exploits.extend(svc['exploits'])
-        
+                all_exploits.extend(svc["exploits"])
+
         # Count by severity
         severity_counts = {
-            'critical': len([e for e in all_exploits if e.get('severity') == 'critical']),
-            'high': len([e for e in all_exploits if e.get('severity') == 'high']),
-            'medium': len([e for e in all_exploits if e.get('severity') == 'medium']),
-            'low': len([e for e in all_exploits if e.get('severity') == 'low'])
+            "critical": len(
+                [e for e in all_exploits if e.get("severity") == "critical"]
+            ),
+            "high": len([e for e in all_exploits if e.get("severity") == "high"]),
+            "medium": len([e for e in all_exploits if e.get("severity") == "medium"]),
+            "low": len([e for e in all_exploits if e.get("severity") == "low"]),
         }
-        
+
         # Get top 5 exploits
         top_exploits = sorted(
             all_exploits,
             key=lambda x: (
-                {'critical': 4, 'high': 3, 'medium': 2, 'low': 1}.get(x.get('severity', 'low'), 0),
-                {'exact': 3, 'range': 2, 'generic': 1}.get(x.get('match_type', 'generic'), 0)
+                {"critical": 4, "high": 3, "medium": 2, "low": 1}.get(
+                    x.get("severity", "low"), 0
+                ),
+                {"exact": 3, "range": 2, "generic": 1}.get(
+                    x.get("match_type", "generic"), 0
+                ),
             ),
-            reverse=True
+            reverse=True,
         )[:5]
-        
+
         return {
-            'total_exploits': len(all_exploits),
+            "total_exploits": len(all_exploits),
             **severity_counts,
-            'services_with_exploits': services_with_exploits,
-            'top_exploits': top_exploits
+            "services_with_exploits": services_with_exploits,
+            "top_exploits": top_exploits,
         }
-    
+
     def get_engagement_summary(self, engagement_id: int) -> Dict:
         """
         Get exploit summary for entire engagement.
-        
+
         Returns:
             {
                 'total_hosts': 5,
@@ -355,45 +365,55 @@ class ExploitSuggestionEngine:
             }
         """
         suggestions = self.generate_suggestions(engagement_id)
-        
-        total_hosts = len(suggestions['hosts'])
+
+        total_hosts = len(suggestions["hosts"])
         all_exploits = []
-        
+
         host_exploit_counts = []
-        
-        for host in suggestions['hosts']:
+
+        for host in suggestions["hosts"]:
             host_exploits = []
-            for svc in host['services']:
-                if svc.get('exploits'):
-                    host_exploits.extend(svc['exploits'])
-            
+            for svc in host["services"]:
+                if svc.get("exploits"):
+                    host_exploits.extend(svc["exploits"])
+
             all_exploits.extend(host_exploits)
-            
+
             if host_exploits:
-                host_exploit_counts.append({
-                    'host_id': host['host_id'],
-                    'ip': host['ip'],
-                    'hostname': host.get('hostname'),
-                    'exploit_count': len(host_exploits),
-                    'critical_count': len([e for e in host_exploits if e.get('severity') == 'critical'])
-                })
-        
+                host_exploit_counts.append(
+                    {
+                        "host_id": host["host_id"],
+                        "ip": host["ip"],
+                        "hostname": host.get("hostname"),
+                        "exploit_count": len(host_exploits),
+                        "critical_count": len(
+                            [
+                                e
+                                for e in host_exploits
+                                if e.get("severity") == "critical"
+                            ]
+                        ),
+                    }
+                )
+
         # Sort hosts by exploit count
         top_vulnerable = sorted(
             host_exploit_counts,
-            key=lambda x: (x['critical_count'], x['exploit_count']),
-            reverse=True
+            key=lambda x: (x["critical_count"], x["exploit_count"]),
+            reverse=True,
         )[:10]
-        
+
         return {
-            'total_hosts': total_hosts,
-            'hosts_with_exploits': len(host_exploit_counts),
-            'total_exploits': len(all_exploits),
-            'critical': len([e for e in all_exploits if e.get('severity') == 'critical']),
-            'high': len([e for e in all_exploits if e.get('severity') == 'high']),
-            'medium': len([e for e in all_exploits if e.get('severity') == 'medium']),
-            'low': len([e for e in all_exploits if e.get('severity') == 'low']),
-            'top_vulnerable_hosts': top_vulnerable
+            "total_hosts": total_hosts,
+            "hosts_with_exploits": len(host_exploit_counts),
+            "total_exploits": len(all_exploits),
+            "critical": len(
+                [e for e in all_exploits if e.get("severity") == "critical"]
+            ),
+            "high": len([e for e in all_exploits if e.get("severity") == "high"]),
+            "medium": len([e for e in all_exploits if e.get("severity") == "medium"]),
+            "low": len([e for e in all_exploits if e.get("severity") == "low"]),
+            "top_vulnerable_hosts": top_vulnerable,
         }
 
     def _prioritize_exploits(self, exploits: List[Dict]) -> List[Dict]:
@@ -411,14 +431,14 @@ class ExploitSuggestionEngine:
         Returns:
             Sorted list of exploits
         """
-        source_order = {'both': 3, 'wazuh': 2, 'scan': 1}
-        severity_order = {'critical': 5, 'high': 4, 'medium': 3, 'low': 2, 'info': 1}
-        match_order = {'exact': 3, 'range': 2, 'generic': 1}
+        source_order = {"both": 3, "wazuh": 2, "scan": 1}
+        severity_order = {"critical": 5, "high": 4, "medium": 3, "low": 2, "info": 1}
+        match_order = {"exact": 3, "range": 2, "generic": 1}
 
         def sort_key(exploit):
-            source_score = source_order.get(exploit.get('detection_source', 'scan'), 0)
-            severity_score = severity_order.get(exploit.get('severity', 'low'), 0)
-            match_score = match_order.get(exploit.get('match_type', 'generic'), 0)
+            source_score = source_order.get(exploit.get("detection_source", "scan"), 0)
+            severity_score = severity_order.get(exploit.get("severity", "low"), 0)
+            match_score = match_order.get(exploit.get("match_type", "generic"), 0)
             return (source_score, severity_score, match_score)
 
         return sorted(exploits, key=sort_key, reverse=True)