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

souleyez/ai/feedback_handler.py

@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
 class FeedbackHandler:
     """
     Automatically update database based on command execution results.
-    
+
     Implements the feedback loop:
     - Update credential status (valid/invalid)
     - Update host status (compromised/active)
@@ -32,63 +32,74 @@ class FeedbackHandler:
         engagement_id: int,
         parsed_result: Dict[str, Any],
         recommendation: Dict[str, Any],
-        command: str
+        command: str,
     ) -> Dict[str, Any]:
         """
         Apply feedback updates to database.
-        
+
         Args:
             engagement_id: Current engagement ID
             parsed_result: Parsed command result
             recommendation: Original AI recommendation
             command: Command that was executed
-            
+
         Returns:
             Dict describing what was updated
         """
         feedback = {
-            'hosts_updated': 0,
-            'credentials_updated': 0,
-            'services_added': 0,
-            'notes_added': []
+            "hosts_updated": 0,
+            "credentials_updated": 0,
+            "services_added": 0,
+            "notes_added": [],
         }
-        
+
         # Extract target IP
-        target = recommendation.get('target', '')
-        action = recommendation.get('action', '').lower()
-        
+        target = recommendation.get("target", "")
+        action = recommendation.get("action", "").lower()
+
         import re
-        ip_match = re.search(r'\b(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\b', target)
+
+        ip_match = re.search(r"\b(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\b", target)
         if not ip_match:
             logger.warning(f"Could not extract IP from target: {target}")
             return feedback
-        
+
         ip = ip_match.group(1)
-        
+
         # Get host from database
         host = self.host_mgr.get_host_by_ip(engagement_id, ip)
         if not host:
             logger.warning(f"Host {ip} not found in engagement {engagement_id}")
             return feedback
-        
+
         # Handle SSH results
-        if 'ssh' in command.lower() and parsed_result.get('credential_valid') is not None:
-            feedback.update(self._handle_ssh_feedback(
-                engagement_id, host, parsed_result, recommendation
-            ))
-        
+        if (
+            "ssh" in command.lower()
+            and parsed_result.get("credential_valid") is not None
+        ):
+            feedback.update(
+                self._handle_ssh_feedback(
+                    engagement_id, host, parsed_result, recommendation
+                )
+            )
+
         # Handle MySQL results
-        elif 'mysql' in command.lower() and parsed_result.get('credential_valid') is not None:
-            feedback.update(self._handle_mysql_feedback(
-                engagement_id, host, parsed_result, recommendation
-            ))
-        
+        elif (
+            "mysql" in command.lower()
+            and parsed_result.get("credential_valid") is not None
+        ):
+            feedback.update(
+                self._handle_mysql_feedback(
+                    engagement_id, host, parsed_result, recommendation
+                )
+            )
+
         # Handle nmap results
-        elif 'nmap' in command.lower() and parsed_result.get('open_ports'):
-            feedback.update(self._handle_nmap_feedback(
-                engagement_id, host, parsed_result
-            ))
-        
+        elif "nmap" in command.lower() and parsed_result.get("open_ports"):
+            feedback.update(
+                self._handle_nmap_feedback(engagement_id, host, parsed_result)
+            )
+
         return feedback
 
     def _handle_ssh_feedback(
@@ -96,61 +107,63 @@ class FeedbackHandler:
         engagement_id: int,
         host: Dict[str, Any],
         result: Dict[str, Any],
-        recommendation: Dict[str, Any]
+        recommendation: Dict[str, Any],
     ) -> Dict[str, Any]:
         """Handle feedback for SSH credential testing."""
-        feedback = {'hosts_updated': 0, 'credentials_updated': 0, 'notes_added': []}
-        
+        feedback = {"hosts_updated": 0, "credentials_updated": 0, "notes_added": []}
+
         # Extract credentials from recommendation
         import re
-        text = f"{recommendation.get('action', '')} {recommendation.get('rationale', '')}"
-        cred_match = re.search(r'(\w+):(\w+)', text)
-        
+
+        text = (
+            f"{recommendation.get('action', '')} {recommendation.get('rationale', '')}"
+        )
+        cred_match = re.search(r"(\w+):(\w+)", text)
+
         if not cred_match:
             logger.warning("Could not extract credentials from recommendation")
             return feedback
-        
+
         username = cred_match.group(1)
         password = cred_match.group(2)
-        
+
         # Find credential in database
         creds = self.creds_mgr.list_credentials(engagement_id)
         cred_id = None
         for cred in creds:
-            if cred.get('username') == username and cred.get('password') == password:
-                cred_id = cred.get('id')
+            if cred.get("username") == username and cred.get("password") == password:
+                cred_id = cred.get("id")
                 break
-        
+
         # Update credential status
         if cred_id:
-            status = 'valid' if result.get('credential_valid') else 'invalid'
-            notes = result.get('details', '')
-            
-            self.creds_mgr.update_credential_status(
-                cred_id,
-                status=status,
-                notes=notes
+            status = "valid" if result.get("credential_valid") else "invalid"
+            notes = result.get("details", "")
+
+            self.creds_mgr.update_credential_status(cred_id, status=status, notes=notes)
+            feedback["credentials_updated"] = 1
+            feedback["notes_added"].append(
+                f"Updated credential {username} status: {status}"
             )
-            feedback['credentials_updated'] = 1
-            feedback['notes_added'].append(f"Updated credential {username} status: {status}")
             logger.info(f"Updated credential {cred_id} status to {status}")
-        
+
         # Update host status if credentials were valid
-        if result.get('credential_valid'):
-            access_level = result.get('access_level', 'user')
-            status = 'compromised'
+        if result.get("credential_valid"):
+            access_level = result.get("access_level", "user")
+            status = "compromised"
             notes = f"Gained {access_level} access via SSH with {username}:{password}"
-            
+
             self.host_mgr.update_host_status(
-                host['id'],
-                status=status,
-                access_level=access_level,
-                notes=notes
+                host["id"], status=status, access_level=access_level, notes=notes
+            )
+            feedback["hosts_updated"] = 1
+            feedback["notes_added"].append(
+                f"Updated host {host['ip_address']}: {status}, access={access_level}"
             )
-            feedback['hosts_updated'] = 1
-            feedback['notes_added'].append(f"Updated host {host['ip_address']}: {status}, access={access_level}")
-            logger.info(f"Updated host {host['id']} to {status} with {access_level} access")
-        
+            logger.info(
+                f"Updated host {host['id']} to {status} with {access_level} access"
+            )
+
         return feedback
 
     def _handle_mysql_feedback(
@@ -158,106 +171,108 @@ class FeedbackHandler:
         engagement_id: int,
         host: Dict[str, Any],
         result: Dict[str, Any],
-        recommendation: Dict[str, Any]
+        recommendation: Dict[str, Any],
     ) -> Dict[str, Any]:
         """Handle feedback for MySQL credential testing."""
-        feedback = {'hosts_updated': 0, 'credentials_updated': 0, 'notes_added': []}
-        
+        feedback = {"hosts_updated": 0, "credentials_updated": 0, "notes_added": []}
+
         # Extract credentials
         import re
-        text = f"{recommendation.get('action', '')} {recommendation.get('rationale', '')}"
-        cred_match = re.search(r'(\w+):(\w+)', text)
-        
+
+        text = (
+            f"{recommendation.get('action', '')} {recommendation.get('rationale', '')}"
+        )
+        cred_match = re.search(r"(\w+):(\w+)", text)
+
         if not cred_match:
             return feedback
-        
+
         username = cred_match.group(1)
         password = cred_match.group(2)
-        
+
         # Find/update credential
         creds = self.creds_mgr.list_credentials(engagement_id)
         cred_id = None
         for cred in creds:
-            if (cred.get('username') == username and 
-                cred.get('password') == password and
-                cred.get('service') == 'mysql'):
-                cred_id = cred.get('id')
+            if (
+                cred.get("username") == username
+                and cred.get("password") == password
+                and cred.get("service") == "mysql"
+            ):
+                cred_id = cred.get("id")
                 break
-        
+
         if cred_id:
-            status = 'valid' if result.get('credential_valid') else 'invalid'
-            notes = result.get('details', '')
-            
-            self.creds_mgr.update_credential_status(
-                cred_id,
-                status=status,
-                notes=notes
+            status = "valid" if result.get("credential_valid") else "invalid"
+            notes = result.get("details", "")
+
+            self.creds_mgr.update_credential_status(cred_id, status=status, notes=notes)
+            feedback["credentials_updated"] = 1
+            feedback["notes_added"].append(
+                f"Updated MySQL credential {username} status: {status}"
             )
-            feedback['credentials_updated'] = 1
-            feedback['notes_added'].append(f"Updated MySQL credential {username} status: {status}")
-        
+
         # Update host notes if connection was successful
-        if result.get('credential_valid'):
+        if result.get("credential_valid"):
             notes = f"MySQL access confirmed with {username}:{password}"
-            if result.get('databases'):
+            if result.get("databases"):
                 notes += f" - Databases: {', '.join(result['databases'][:5])}"
-            
+
             # Only update if not already compromised with higher access
-            current_access = host.get('access_level', 'none')
-            if current_access == 'none':
-                self.host_mgr.update_host_status(
-                    host['id'],
-                    notes=notes
+            current_access = host.get("access_level", "none")
+            if current_access == "none":
+                self.host_mgr.update_host_status(host["id"], notes=notes)
+                feedback["hosts_updated"] = 1
+                feedback["notes_added"].append(
+                    f"Added MySQL access notes to host {host['ip_address']}"
                 )
-                feedback['hosts_updated'] = 1
-                feedback['notes_added'].append(f"Added MySQL access notes to host {host['ip_address']}")
-        
+
         return feedback
 
     def _handle_nmap_feedback(
-        self,
-        engagement_id: int,
-        host: Dict[str, Any],
-        result: Dict[str, Any]
+        self, engagement_id: int, host: Dict[str, Any], result: Dict[str, Any]
     ) -> Dict[str, Any]:
         """Handle feedback for nmap scans."""
-        feedback = {'hosts_updated': 0, 'services_added': 0, 'notes_added': []}
-        
+        feedback = {"hosts_updated": 0, "services_added": 0, "notes_added": []}
+
         # Add discovered services
-        open_ports = result.get('open_ports', [])
-        services_info = result.get('services', {})
-        
+        open_ports = result.get("open_ports", [])
+        services_info = result.get("services", {})
+
         for port in open_ports:
-            service_info = services_info.get(str(port), 'unknown')
-            
+            service_info = services_info.get(str(port), "unknown")
+
             # Try to parse service name and version
             import re
-            service_match = re.match(r'(\S+)(?:\s+(.+))?', service_info)
+
+            service_match = re.match(r"(\S+)(?:\s+(.+))?", service_info)
             if service_match:
                 service_name = service_match.group(1)
-                service_version = service_match.group(2) if service_match.group(2) else None
+                service_version = (
+                    service_match.group(2) if service_match.group(2) else None
+                )
             else:
-                service_name = 'unknown'
+                service_name = "unknown"
                 service_version = None
-            
+
             # Add service to database
             try:
                 self.host_mgr.add_service(
-                    host['id'],
+                    host["id"],
                     port=port,
-                    protocol='tcp',
+                    protocol="tcp",
                     service_name=service_name,
                     service_version=service_version,
-                    state='open'
+                    state="open",
                 )
-                feedback['services_added'] += 1
+                feedback["services_added"] += 1
                 logger.info(f"Added service {service_name} on port {port}")
             except Exception as e:
                 logger.warning(f"Failed to add service {port}: {e}")
-        
-        if feedback['services_added'] > 0:
-            feedback['notes_added'].append(
+
+        if feedback["services_added"] > 0:
+            feedback["notes_added"].append(
                 f"Added {feedback['services_added']} services to host {host['ip_address']}"
             )
-        
+
         return feedback

souleyez/ai/llm_factory.py

@@ -4,6 +4,7 @@ souleyez.ai.llm_factory - Factory for creating LLM providers
 This module provides a factory pattern for creating LLM providers
 based on user configuration.
 """
+
 import logging
 from typing import Optional
 
@@ -35,20 +36,24 @@ class LLMFactory:
 
         if provider_type is None:
             # Read from config
-            provider_str = get('ai.provider', 'ollama')
+            provider_str = get("ai.provider", "ollama")
             try:
                 provider_type = LLMProviderType(provider_str)
             except ValueError:
-                logger.warning(f"Unknown provider '{provider_str}', defaulting to ollama")
+                logger.warning(
+                    f"Unknown provider '{provider_str}', defaulting to ollama"
+                )
                 provider_type = LLMProviderType.OLLAMA
 
         if provider_type == LLMProviderType.CLAUDE:
             from .claude_provider import ClaudeProvider
-            model = get('ai.claude_model')
+
+            model = get("ai.claude_model")
             return ClaudeProvider(model=model)
         else:
             from .ollama_provider import OllamaProvider
-            model = get('ai.ollama_model') or get('settings.ollama_model')
+
+            model = get("ai.ollama_model") or get("settings.ollama_model")
             return OllamaProvider(model=model)
 
     @staticmethod
@@ -103,33 +108,35 @@ class LLMFactory:
         # Check Ollama
         try:
             from .ollama_provider import OllamaProvider
+
             ollama = OllamaProvider()
-            results['ollama'] = {
-                'available': ollama.is_available(),
-                'status': ollama.get_status(),
-                'configured': get('ai.provider', 'ollama') == 'ollama'
+            results["ollama"] = {
+                "available": ollama.is_available(),
+                "status": ollama.get_status(),
+                "configured": get("ai.provider", "ollama") == "ollama",
             }
         except Exception as e:
-            results['ollama'] = {
-                'available': False,
-                'error': str(e),
-                'configured': get('ai.provider', 'ollama') == 'ollama'
+            results["ollama"] = {
+                "available": False,
+                "error": str(e),
+                "configured": get("ai.provider", "ollama") == "ollama",
             }
 
         # Check Claude
         try:
             from .claude_provider import ClaudeProvider
+
             claude = ClaudeProvider()
-            results['claude'] = {
-                'available': claude.is_available(),
-                'status': claude.get_status(),
-                'configured': get('ai.provider', 'ollama') == 'claude'
+            results["claude"] = {
+                "available": claude.is_available(),
+                "status": claude.get_status(),
+                "configured": get("ai.provider", "ollama") == "claude",
             }
         except Exception as e:
-            results['claude'] = {
-                'available': False,
-                'error': str(e),
-                'configured': get('ai.provider', 'ollama') == 'claude'
+            results["claude"] = {
+                "available": False,
+                "error": str(e),
+                "configured": get("ai.provider", "ollama") == "claude",
             }
 
         return results

souleyez/ai/llm_provider.py

@@ -4,6 +4,7 @@ souleyez.ai.llm_provider - Abstract LLM provider interface
 This module defines the abstract base class for LLM providers,
 enabling support for multiple backends (Ollama, Claude, etc.)
 """
+
 from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Optional, Dict, Any
@@ -14,6 +15,7 @@ logger = logging.getLogger(__name__)
 
 class LLMProviderType(Enum):
     """Supported LLM provider types."""
+
     OLLAMA = "ollama"
     CLAUDE = "claude"
 
@@ -54,7 +56,7 @@ class LLMProvider(ABC):
         prompt: str,
         system_prompt: Optional[str] = None,
         max_tokens: int = 4096,
-        temperature: float = 0.3
+        temperature: float = 0.3,
     ) -> Optional[str]:
         """
         Generate text from a prompt.
@@ -86,7 +88,7 @@ class LLMProvider(ABC):
         system_prompt: Optional[str] = None,
         max_tokens: int = 4096,
         temperature: float = 0.3,
-        fallback: Optional[str] = None
+        fallback: Optional[str] = None,
     ) -> str:
         """
         Generate text with a fallback value if generation fails.

souleyez/ai/ollama_provider.py

@@ -3,6 +3,7 @@ souleyez.ai.ollama_provider - Ollama LLM provider implementation
 
 Wraps the existing OllamaService to implement the LLMProvider interface.
 """
+
 import logging
 from typing import Optional, Dict, Any
 
@@ -20,11 +21,7 @@ class OllamaProvider(LLMProvider):
     No API key required, but Ollama must be installed and running.
     """
 
-    def __init__(
-        self,
-        model: Optional[str] = None,
-        endpoint: Optional[str] = None
-    ):
+    def __init__(self, model: Optional[str] = None, endpoint: Optional[str] = None):
         """
         Initialize Ollama provider.
 
@@ -71,7 +68,7 @@ class OllamaProvider(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 Ollama.
@@ -108,7 +105,7 @@ class OllamaProvider(LLMProvider):
             dict: Status including connection, models, availability
         """
         status = self._service.get_status()
-        status['provider_type'] = self.provider_type.value
-        status['provider'] = 'Ollama'
-        status['provider_name'] = 'Ollama (Local)'
+        status["provider_type"] = self.provider_type.value
+        status["provider"] = "Ollama"
+        status["provider_name"] = "Ollama (Local)"
         return status