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

souleyez/ai/ollama_service.py

@@ -5,6 +5,7 @@ This module provides the connection and management layer for Ollama,
 enabling SoulEyez to generate AI-powered attack path recommendations
 without sending sensitive engagement data to the cloud.
 """
+
 import logging
 from typing import Optional, Dict, Any
 
@@ -12,6 +13,7 @@ from typing import Optional, Dict, Any
 try:
     import ollama
     from ollama import Client
+
     OLLAMA_AVAILABLE = True
 except ImportError:
     ollama = None
@@ -35,8 +37,12 @@ class OllamaService:
     DEFAULT_MODEL = "llama3.1:8b"  # 8B parameter model for good quality reasoning
     CONNECTION_TIMEOUT = 30  # seconds
 
-    def __init__(self, endpoint: Optional[str] = None, model: Optional[str] = None,
-                 allow_failback: bool = True):
+    def __init__(
+        self,
+        endpoint: Optional[str] = None,
+        model: Optional[str] = None,
+        allow_failback: bool = True,
+    ):
         """
         Initialize Ollama service.
 
@@ -52,15 +58,16 @@ class OllamaService:
 
         # Read endpoint from config if not provided
         if endpoint is None:
-            endpoint = get('ai.ollama_url', self.DEFAULT_ENDPOINT)
+            endpoint = get("ai.ollama_url", self.DEFAULT_ENDPOINT)
 
         self._original_endpoint = endpoint
         self._failback_used = False
         self._allow_failback = allow_failback
-        self._ollama_mode = get('ai.ollama_mode', 'local')
+        self._ollama_mode = get("ai.ollama_mode", "local")
 
         # Security: Validate endpoint is localhost or VM host gateway
         from souleyez.core.network_utils import is_valid_ollama_host
+
         is_valid, reason = is_valid_ollama_host(endpoint)
         if not is_valid:
             logger.error(f"Ollama endpoint blocked: {reason}")
@@ -80,7 +87,7 @@ class OllamaService:
 
         # Read model from config if not provided
         if model is None:
-            model = get('ai.ollama_model', self.DEFAULT_MODEL)
+            model = get("ai.ollama_model", self.DEFAULT_MODEL)
 
         self.model = model
 
@@ -100,14 +107,16 @@ class OllamaService:
             return False
 
         # Only failback if we're in remote mode and not already on localhost
-        if self._ollama_mode != 'remote':
+        if self._ollama_mode != "remote":
             return False
 
         if self.endpoint == self.DEFAULT_ENDPOINT:
             return False
 
-        logger.warning(f"Remote Ollama at {self._original_endpoint} unreachable, "
-                       f"falling back to localhost")
+        logger.warning(
+            f"Remote Ollama at {self._original_endpoint} unreachable, "
+            f"falling back to localhost"
+        )
 
         # Try localhost
         try:
@@ -183,7 +192,9 @@ class OllamaService:
             if self._model_available:
                 logger.info(f"Model '{model_name}' is available")
             else:
-                logger.warning(f"Model '{model_name}' not found. Available: {available_models}")
+                logger.warning(
+                    f"Model '{model_name}' not found. Available: {available_models}"
+                )
 
             return self._model_available
         except Exception as e:
@@ -216,7 +227,9 @@ class OllamaService:
             logger.error(f"Failed to pull model: {e}")
             return False
 
-    def generate(self, prompt: str, model_name: Optional[str] = None, timeout: int = 120) -> Optional[str]:
+    def generate(
+        self, prompt: str, model_name: Optional[str] = None, timeout: int = 120
+    ) -> Optional[str]:
         """
         Generate a response from the LLM.
 
@@ -255,13 +268,10 @@ class OllamaService:
                 signal.alarm(timeout)
 
                 try:
-                    response = self.client.generate(
-                        model=model_name,
-                        prompt=prompt
-                    )
+                    response = self.client.generate(model=model_name, prompt=prompt)
                     signal.alarm(0)
                     signal.signal(signal.SIGALRM, old_handler)
-                    return response.get('response', '')
+                    return response.get("response", "")
                 except TimeoutError:
                     signal.alarm(0)
                     signal.signal(signal.SIGALRM, old_handler)
@@ -274,11 +284,8 @@ class OllamaService:
                     return None
             else:
                 # In a thread - can't use signals, just call directly
-                response = self.client.generate(
-                    model=model_name,
-                    prompt=prompt
-                )
-                return response.get('response', '')
+                response = self.client.generate(model=model_name, prompt=prompt)
+                return response.get("response", "")
         except Exception as e:
             logger.error(f"Generation failed: {e}")
             return None
@@ -291,36 +298,36 @@ class OllamaService:
             dict: Status information including connection, models, etc.
         """
         status = {
-            'endpoint': self.endpoint,
-            'original_endpoint': self._original_endpoint,
-            'mode': self._ollama_mode,
-            'failback_used': self._failback_used,
-            'connected': False,
-            'models': [],
-            'default_model': self.DEFAULT_MODEL,
-            'configured_model': self.model,
-            'model_available': False,
-            'error': None
+            "endpoint": self.endpoint,
+            "original_endpoint": self._original_endpoint,
+            "mode": self._ollama_mode,
+            "failback_used": self._failback_used,
+            "connected": False,
+            "models": [],
+            "default_model": self.DEFAULT_MODEL,
+            "configured_model": self.model,
+            "model_available": False,
+            "error": None,
         }
 
         # Check connection
         if not self.check_connection():
-            status['error'] = 'Cannot connect to Ollama. Is it running?'
+            status["error"] = "Cannot connect to Ollama. Is it running?"
             return status
 
-        status['connected'] = True
-        status['endpoint'] = self.endpoint  # Update in case failback changed it
-        status['failback_used'] = self._failback_used
+        status["connected"] = True
+        status["endpoint"] = self.endpoint  # Update in case failback changed it
+        status["failback_used"] = self._failback_used
 
         # Get available models
         try:
             models_response = self.client.list()
-            status['models'] = [m.model for m in models_response.models]
+            status["models"] = [m.model for m in models_response.models]
         except Exception as e:
-            status['error'] = f'Error listing models: {e}'
+            status["error"] = f"Error listing models: {e}"
             return status
 
         # Check if default model is available
-        status['model_available'] = self.check_model()
+        status["model_available"] = self.check_model()
 
         return status

souleyez/ai/path_scorer.py

@@ -18,52 +18,47 @@ class PathScorer:
 
     # Scoring weights (configurable)
     WEIGHTS = {
-        'success': 0.4,
-        'impact': 0.3,
-        'stealth': 0.2,
-        'complexity': 0.1  # Negative weight
+        "success": 0.4,
+        "impact": 0.3,
+        "stealth": 0.2,
+        "complexity": 0.1,  # Negative weight
     }
 
     # Success probability scoring
     SUCCESS_SCORES = {
-        'has_valid_creds': 30,
-        'known_vulnerability': 25,
-        'common_service': 15,  # SSH, RDP, HTTP
-        'uncommon_service': 5
+        "has_valid_creds": 30,
+        "known_vulnerability": 25,
+        "common_service": 15,  # SSH, RDP, HTTP
+        "uncommon_service": 5,
     }
 
     # Impact scoring
     IMPACT_SCORES = {
-        'domain_controller': 40,
-        'database': 30,
-        'file_server': 20,
-        'workstation': 10,
-        'unknown': 5
+        "domain_controller": 40,
+        "database": 30,
+        "file_server": 20,
+        "workstation": 10,
+        "unknown": 5,
     }
 
     # Stealth scoring
     STEALTH_SCORES = {
-        'credential_reuse': 30,  # No exploits
-        'quiet_exploit': 20,
-        'noisy_exploit': 10,
-        'unknown': 15  # Default to middle ground
+        "credential_reuse": 30,  # No exploits
+        "quiet_exploit": 20,
+        "noisy_exploit": 10,
+        "unknown": 15,  # Default to middle ground
     }
 
     # Complexity penalty
-    COMPLEXITY_PENALTIES = {
-        1: 0,
-        2: 5,
-        3: 5,
-        4: 10,
-        5: 10,
-        6: 20
-    }
+    COMPLEXITY_PENALTIES = {1: 0, 2: 5, 3: 5, 4: 10, 5: 10, 6: 20}
 
     def __init__(self):
         """Initialize path scorer."""
         pass
 
-    def score_path(self, path: Dict[str, Any], engagement_data: Dict[str, Any]) -> Dict[str, Any]:
+    def score_path(
+        self, path: Dict[str, Any], engagement_data: Dict[str, Any]
+    ) -> Dict[str, Any]:
         """
         Score a single attack path.
 
@@ -85,93 +80,111 @@ class PathScorer:
 
         # Calculate weighted total
         total_score = (
-            (success_score * self.WEIGHTS['success']) +
-            (impact_score * self.WEIGHTS['impact']) +
-            (stealth_score * self.WEIGHTS['stealth']) -
-            (complexity_penalty * self.WEIGHTS['complexity'])
+            (success_score * self.WEIGHTS["success"])
+            + (impact_score * self.WEIGHTS["impact"])
+            + (stealth_score * self.WEIGHTS["stealth"])
+            - (complexity_penalty * self.WEIGHTS["complexity"])
         )
 
         return {
-            'path': path,
-            'scores': {
-                'success': success_score,
-                'impact': impact_score,
-                'stealth': stealth_score,
-                'complexity': complexity_penalty
+            "path": path,
+            "scores": {
+                "success": success_score,
+                "impact": impact_score,
+                "stealth": stealth_score,
+                "complexity": complexity_penalty,
             },
-            'total_score': round(total_score, 2),
-            'rank': None  # Set by rank_paths()
+            "total_score": round(total_score, 2),
+            "rank": None,  # Set by rank_paths()
         }
 
-    def _score_success(self, path: Dict[str, Any], engagement_data: Dict[str, Any]) -> float:
+    def _score_success(
+        self, path: Dict[str, Any], engagement_data: Dict[str, Any]
+    ) -> float:
         """Score success probability based on available resources."""
         score = 0
-        action = path.get('action', '').lower()
-        target = path.get('target', '').lower()
+        action = path.get("action", "").lower()
+        target = path.get("target", "").lower()
 
         # Check for valid credentials
-        creds = engagement_data.get('credentials', [])
-        valid_creds = [c for c in creds if c.get('status') == 'valid']
-        if valid_creds and ('credential' in action or 'login' in action or 'auth' in action):
-            score += self.SUCCESS_SCORES['has_valid_creds']
+        creds = engagement_data.get("credentials", [])
+        valid_creds = [c for c in creds if c.get("status") == "valid"]
+        if valid_creds and (
+            "credential" in action or "login" in action or "auth" in action
+        ):
+            score += self.SUCCESS_SCORES["has_valid_creds"]
 
         # Check for known vulnerabilities
-        findings = engagement_data.get('findings', [])
-        critical_vulns = [f for f in findings if f.get('severity') in ['critical', 'high']]
-        if critical_vulns and ('exploit' in action or 'vulnerability' in action):
-            score += self.SUCCESS_SCORES['known_vulnerability']
+        findings = engagement_data.get("findings", [])
+        critical_vulns = [
+            f for f in findings if f.get("severity") in ["critical", "high"]
+        ]
+        if critical_vulns and ("exploit" in action or "vulnerability" in action):
+            score += self.SUCCESS_SCORES["known_vulnerability"]
 
         # Check service type
-        common_services = ['ssh', 'rdp', 'http', 'https', 'mysql', 'smb']
+        common_services = ["ssh", "rdp", "http", "https", "mysql", "smb"]
         if any(svc in target or svc in action for svc in common_services):
-            score += self.SUCCESS_SCORES['common_service']
+            score += self.SUCCESS_SCORES["common_service"]
         else:
-            score += self.SUCCESS_SCORES['uncommon_service']
+            score += self.SUCCESS_SCORES["uncommon_service"]
 
         return min(score, 100)  # Cap at 100
 
-    def _score_impact(self, path: Dict[str, Any], engagement_data: Dict[str, Any]) -> float:
+    def _score_impact(
+        self, path: Dict[str, Any], engagement_data: Dict[str, Any]
+    ) -> float:
         """Score potential impact of successful path."""
-        target = path.get('target', '').lower()
-        action = path.get('action', '').lower()
+        target = path.get("target", "").lower()
+        action = path.get("action", "").lower()
 
         # Check target type
-        if 'domain controller' in target or 'dc' in target or 'ad' in target:
-            return self.IMPACT_SCORES['domain_controller']
+        if "domain controller" in target or "dc" in target or "ad" in target:
+            return self.IMPACT_SCORES["domain_controller"]
 
-        if 'database' in target or 'mysql' in target or 'postgres' in target or 'sql' in target:
-            return self.IMPACT_SCORES['database']
+        if (
+            "database" in target
+            or "mysql" in target
+            or "postgres" in target
+            or "sql" in target
+        ):
+            return self.IMPACT_SCORES["database"]
 
-        if 'file server' in target or 'smb' in target or 'share' in target:
-            return self.IMPACT_SCORES['file_server']
+        if "file server" in target or "smb" in target or "share" in target:
+            return self.IMPACT_SCORES["file_server"]
 
         # Check for privilege escalation (high impact)
-        if 'privilege' in action or 'escalat' in action or 'root' in action or 'admin' in action:
-            return self.IMPACT_SCORES['database']  # High impact
+        if (
+            "privilege" in action
+            or "escalat" in action
+            or "root" in action
+            or "admin" in action
+        ):
+            return self.IMPACT_SCORES["database"]  # High impact
 
         # Default to workstation
-        return self.IMPACT_SCORES['workstation']
+        return self.IMPACT_SCORES["workstation"]
 
     def _score_stealth(self, path: Dict[str, Any]) -> float:
         """Score stealth/detection likelihood."""
-        action = path.get('action', '').lower()
+        action = path.get("action", "").lower()
 
         # Credential reuse is stealthiest
-        if 'credential' in action and 'exploit' not in action and 'brute' not in action:
-            return self.STEALTH_SCORES['credential_reuse']
+        if "credential" in action and "exploit" not in action and "brute" not in action:
+            return self.STEALTH_SCORES["credential_reuse"]
 
         # Known quiet exploits
-        quiet_keywords = ['pass-the-hash', 'token', 'kerberos', 'golden ticket']
+        quiet_keywords = ["pass-the-hash", "token", "kerberos", "golden ticket"]
         if any(keyword in action for keyword in quiet_keywords):
-            return self.STEALTH_SCORES['quiet_exploit']
+            return self.STEALTH_SCORES["quiet_exploit"]
 
         # Noisy exploits
-        noisy_keywords = ['brute', 'scan', 'spray', 'exploit', 'buffer overflow']
+        noisy_keywords = ["brute", "scan", "spray", "exploit", "buffer overflow"]
         if any(keyword in action for keyword in noisy_keywords):
-            return self.STEALTH_SCORES['noisy_exploit']
+            return self.STEALTH_SCORES["noisy_exploit"]
 
         # Default
-        return self.STEALTH_SCORES['unknown']
+        return self.STEALTH_SCORES["unknown"]
 
     def _score_complexity(self, path: Dict[str, Any]) -> float:
         """Score complexity (simpler = better, so this is a penalty)."""
@@ -179,7 +192,7 @@ class PathScorer:
         # Multi-step paths get penalties based on length
 
         # Check if this is a multi-step path
-        steps = path.get('steps', [])
+        steps = path.get("steps", [])
         if steps:
             num_steps = len(steps)
         else:
@@ -191,7 +204,9 @@ class PathScorer:
         else:
             return self.COMPLEXITY_PENALTIES[6]
 
-    def rank_paths(self, paths: List[Dict[str, Any]], engagement_data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def rank_paths(
+        self, paths: List[Dict[str, Any]], engagement_data: Dict[str, Any]
+    ) -> List[Dict[str, Any]]:
         """
         Score and rank multiple paths.
 
@@ -209,10 +224,10 @@ class PathScorer:
             scored_paths.append(scored_path)
 
         # Sort by total score (descending)
-        scored_paths.sort(key=lambda x: x['total_score'], reverse=True)
+        scored_paths.sort(key=lambda x: x["total_score"], reverse=True)
 
         # Assign ranks
         for i, scored_path in enumerate(scored_paths, 1):
-            scored_path['rank'] = i
+            scored_path["rank"] = i
 
         return scored_paths