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

souleyez/integrations/wazuh/config.py

@@ -4,6 +4,7 @@ SIEM Configuration Management.
 Stores SIEM connection settings with encrypted credentials.
 Supports multiple SIEM platforms: Wazuh, Splunk, Elastic, Sentinel.
 """
+
 import json
 from typing import Optional, Dict, Any, List
 from pathlib import Path
@@ -11,7 +12,7 @@ from souleyez.storage.database import get_db
 from souleyez.storage.crypto import get_crypto_manager
 
 # Supported SIEM types (Open Source first, then Commercial)
-SIEM_TYPES = ['wazuh', 'elastic', 'splunk', 'sentinel', 'google_secops']
+SIEM_TYPES = ["wazuh", "elastic", "splunk", "sentinel", "google_secops"]
 
 
 class WazuhConfig:
@@ -22,7 +23,9 @@ class WazuhConfig:
     """
 
     @staticmethod
-    def get_config(engagement_id: int, siem_type: str = None) -> Optional[Dict[str, Any]]:
+    def get_config(
+        engagement_id: int, siem_type: str = None
+    ) -> Optional[Dict[str, Any]]:
         """
         Get SIEM config for an engagement.
 
@@ -40,35 +43,44 @@ class WazuhConfig:
         # Check if new columns exist (migration 025+)
         cursor.execute("PRAGMA table_info(wazuh_config)")
         columns = [col[1] for col in cursor.fetchall()]
-        has_new_columns = 'siem_type' in columns
+        has_new_columns = "siem_type" in columns
 
         # Query with or without new columns
         if has_new_columns:
             if siem_type:
-                cursor.execute("""
+                cursor.execute(
+                    """
                     SELECT api_url, api_user, api_password, indexer_url, indexer_user,
                            indexer_password, verify_ssl, enabled, siem_type, config_json
                     FROM wazuh_config
                     WHERE engagement_id = ? AND siem_type = ?
-                """, (engagement_id, siem_type))
+                """,
+                    (engagement_id, siem_type),
+                )
             else:
                 # Get most recently updated config (the "current" selected SIEM)
                 # Not filtering by enabled - user may have selected but not configured yet
-                cursor.execute("""
+                cursor.execute(
+                    """
                     SELECT api_url, api_user, api_password, indexer_url, indexer_user,
                            indexer_password, verify_ssl, enabled, siem_type, config_json
                     FROM wazuh_config
                     WHERE engagement_id = ?
                     ORDER BY updated_at DESC
                     LIMIT 1
-                """, (engagement_id,))
+                """,
+                    (engagement_id,),
+                )
         else:
-            cursor.execute("""
+            cursor.execute(
+                """
                 SELECT api_url, api_user, api_password, indexer_url, indexer_user,
                        indexer_password, verify_ssl, enabled
                 FROM wazuh_config
                 WHERE engagement_id = ?
-            """, (engagement_id,))
+            """,
+                (engagement_id,),
+            )
 
         row = cursor.fetchone()
         if not row:
@@ -107,7 +119,7 @@ class WazuhConfig:
                 # Decrypt any encrypted fields in extra config
                 crypto = get_crypto_manager()
                 if crypto and extra_config:
-                    for key in ['password', 'client_secret', 'api_key', 'token']:
+                    for key in ["password", "client_secret", "api_key", "token"]:
                         if key in extra_config and extra_config[key]:
                             try:
                                 extra_config[key] = crypto.decrypt(extra_config[key])
@@ -118,9 +130,9 @@ class WazuhConfig:
                 pass
 
         # Map Wazuh fields to generic names for SIEMFactory compatibility
-        if config['siem_type'] == 'wazuh':
-            config['username'] = config.get('api_user')
-            config['password'] = config.get('api_password')
+        if config["siem_type"] == "wazuh":
+            config["username"] = config.get("api_user")
+            config["password"] = config.get("api_password")
 
         return config
 
@@ -130,14 +142,12 @@ class WazuhConfig:
         cursor.execute("PRAGMA table_info(wazuh_config)")
         columns = [col[1] for col in cursor.fetchall()]
 
-        if 'siem_type' not in columns:
+        if "siem_type" not in columns:
             cursor.execute(
                 "ALTER TABLE wazuh_config ADD COLUMN siem_type TEXT DEFAULT 'wazuh'"
             )
-        if 'config_json' not in columns:
-            cursor.execute(
-                "ALTER TABLE wazuh_config ADD COLUMN config_json TEXT"
-            )
+        if "config_json" not in columns:
+            cursor.execute("ALTER TABLE wazuh_config ADD COLUMN config_json TEXT")
         conn.commit()
 
     @staticmethod
@@ -152,7 +162,7 @@ class WazuhConfig:
         verify_ssl: bool = False,
         enabled: bool = True,
         siem_type: str = "wazuh",
-        config_json: Optional[Dict[str, Any]] = None
+        config_json: Optional[Dict[str, Any]] = None,
     ) -> bool:
         """
         Save SIEM config for an engagement.
@@ -199,15 +209,18 @@ class WazuhConfig:
             try:
                 crypto = get_crypto_manager()
                 if crypto:
-                    for key in ['password', 'client_secret', 'api_key', 'token']:
+                    for key in ["password", "client_secret", "api_key", "token"]:
                         if key in encrypted_config and encrypted_config[key]:
-                            encrypted_config[key] = crypto.encrypt(encrypted_config[key])
+                            encrypted_config[key] = crypto.encrypt(
+                                encrypted_config[key]
+                            )
             except Exception:
                 pass
             config_json_str = json.dumps(encrypted_config)
 
         # Upsert config - keyed by (engagement_id, siem_type) for multi-SIEM support
-        cursor.execute("""
+        cursor.execute(
+            """
             INSERT INTO wazuh_config (
                 engagement_id, api_url, api_user, api_password, indexer_url,
                 indexer_user, indexer_password, verify_ssl, enabled,
@@ -225,21 +238,28 @@ class WazuhConfig:
                 enabled = excluded.enabled,
                 config_json = excluded.config_json,
                 updated_at = CURRENT_TIMESTAMP
-        """, (
-            engagement_id, api_url, api_user, encrypted_api_password,
-            indexer_url, indexer_user or 'admin', encrypted_indexer_password,
-            verify_ssl, enabled, siem_type, config_json_str
-        ))
+        """,
+            (
+                engagement_id,
+                api_url,
+                api_user,
+                encrypted_api_password,
+                indexer_url,
+                indexer_user or "admin",
+                encrypted_indexer_password,
+                verify_ssl,
+                enabled,
+                siem_type,
+                config_json_str,
+            ),
+        )
 
         conn.commit()
         return True
 
     @staticmethod
     def save_siem_config(
-        engagement_id: int,
-        siem_type: str,
-        config: Dict[str, Any],
-        enabled: bool = True
+        engagement_id: int, siem_type: str, config: Dict[str, Any], enabled: bool = True
     ) -> bool:
         """
         Save configuration for any SIEM type.
@@ -252,29 +272,29 @@ class WazuhConfig:
             config: SIEM-specific configuration dict
             enabled: Whether integration is enabled
         """
-        if siem_type == 'wazuh':
+        if siem_type == "wazuh":
             # Use existing Wazuh fields for backwards compatibility
             return WazuhConfig.save_config(
                 engagement_id=engagement_id,
-                api_url=config.get('api_url', ''),
-                api_user=config.get('username', config.get('api_user', '')),
-                api_password=config.get('password', config.get('api_password', '')),
-                indexer_url=config.get('indexer_url'),
-                indexer_user=config.get('indexer_user'),
-                indexer_password=config.get('indexer_password'),
-                verify_ssl=config.get('verify_ssl', False),
+                api_url=config.get("api_url", ""),
+                api_user=config.get("username", config.get("api_user", "")),
+                api_password=config.get("password", config.get("api_password", "")),
+                indexer_url=config.get("indexer_url"),
+                indexer_user=config.get("indexer_user"),
+                indexer_password=config.get("indexer_password"),
+                verify_ssl=config.get("verify_ssl", False),
                 enabled=enabled,
-                siem_type='wazuh'
+                siem_type="wazuh",
             )
         else:
             # For other SIEMs, store config in config_json
             return WazuhConfig.save_config(
                 engagement_id=engagement_id,
-                api_url=config.get('api_url', config.get('elasticsearch_url', '')),
-                verify_ssl=config.get('verify_ssl', False),
+                api_url=config.get("api_url", config.get("elasticsearch_url", "")),
+                verify_ssl=config.get("verify_ssl", False),
                 enabled=enabled,
                 siem_type=siem_type,
-                config_json=config
+                config_json=config,
             )
 
     @staticmethod
@@ -292,10 +312,12 @@ class WazuhConfig:
         if siem_type:
             cursor.execute(
                 "DELETE FROM wazuh_config WHERE engagement_id = ? AND siem_type = ?",
-                (engagement_id, siem_type)
+                (engagement_id, siem_type),
             )
         else:
-            cursor.execute("DELETE FROM wazuh_config WHERE engagement_id = ?", (engagement_id,))
+            cursor.execute(
+                "DELETE FROM wazuh_config WHERE engagement_id = ?", (engagement_id,)
+            )
         conn.commit()
         return cursor.rowcount > 0
 
@@ -320,28 +342,34 @@ class WazuhConfig:
         cursor.execute("PRAGMA table_info(wazuh_config)")
         columns = [col[1] for col in cursor.fetchall()]
 
-        if 'siem_type' not in columns:
+        if "siem_type" not in columns:
             # Old schema - only one config possible
-            cursor.execute("""
+            cursor.execute(
+                """
                 SELECT 'wazuh' as siem_type, enabled, api_url
                 FROM wazuh_config
                 WHERE engagement_id = ?
-            """, (engagement_id,))
+            """,
+                (engagement_id,),
+            )
         else:
-            cursor.execute("""
+            cursor.execute(
+                """
                 SELECT siem_type, enabled, api_url, updated_at
                 FROM wazuh_config
                 WHERE engagement_id = ?
                 ORDER BY siem_type
-            """, (engagement_id,))
+            """,
+                (engagement_id,),
+            )
 
         rows = cursor.fetchall()
         return [
             {
-                'siem_type': row[0],
-                'enabled': bool(row[1]),
-                'api_url': row[2],
-                'updated_at': row[3] if len(row) > 3 else None
+                "siem_type": row[0],
+                "enabled": bool(row[1]),
+                "api_url": row[2],
+                "updated_at": row[3] if len(row) > 3 else None,
             }
             for row in rows
         ]
@@ -373,8 +401,8 @@ class WazuhConfig:
         """
         config = WazuhConfig.get_config(engagement_id)
         if config:
-            return config.get('siem_type', 'wazuh')
-        return 'wazuh'
+            return config.get("siem_type", "wazuh")
+        return "wazuh"
 
     @staticmethod
     def set_current_siem(engagement_id: int, siem_type: str) -> bool:
@@ -395,7 +423,7 @@ class WazuhConfig:
         cursor = conn.cursor()
         cursor.execute(
             "UPDATE wazuh_config SET updated_at = CURRENT_TIMESTAMP WHERE engagement_id = ? AND siem_type = ?",
-            (engagement_id, siem_type)
+            (engagement_id, siem_type),
         )
         conn.commit()
         return cursor.rowcount > 0

souleyez/integrations/wazuh/host_mapper.py

@@ -17,11 +17,7 @@ class WazuhHostMapper:
     def __init__(self):
         self.db = get_db()
 
-    def map_agent_to_host(
-        self,
-        engagement_id: int,
-        agent_ip: str
-    ) -> Optional[int]:
+    def map_agent_to_host(self, engagement_id: int, agent_ip: str) -> Optional[int]:
         """
         Find SoulEyez host_id matching an agent IP.
 
@@ -42,7 +38,7 @@ class WazuhHostMapper:
         result = self.db.execute_one(query, (engagement_id, agent_ip))
 
         if result:
-            return result['id']
+            return result["id"]
 
         return None
 
@@ -66,7 +62,7 @@ class WazuhHostMapper:
 
         mapping = {}
         for agent in agents:
-            agent_ip = agent.get('agent_ip')
+            agent_ip = agent.get("agent_ip")
             if agent_ip:
                 host_id = self.map_agent_to_host(engagement_id, agent_ip)
                 mapping[agent_ip] = host_id
@@ -78,10 +74,7 @@ class WazuhHostMapper:
         return mapping
 
     def _update_vuln_host_mapping(
-        self,
-        engagement_id: int,
-        agent_ip: str,
-        host_id: int
+        self, engagement_id: int, agent_ip: str, host_id: int
     ) -> int:
         """
         Update all vulnerabilities from an agent IP with the host_id.
@@ -103,7 +96,7 @@ class WazuhHostMapper:
             WHERE engagement_id = ? AND agent_ip = ? AND host_id = ?
         """
         result = self.db.execute_one(count_query, (engagement_id, agent_ip, host_id))
-        return result.get('count', 0) if result else 0
+        return result.get("count", 0) if result else 0
 
     def get_unmapped_agents(self, engagement_id: int) -> List[Dict[str, any]]:
         """
@@ -167,18 +160,15 @@ class WazuhHostMapper:
 
         if result:
             return {
-                'mapped': result.get('mapped', 0) or 0,
-                'unmapped': result.get('unmapped', 0) or 0,
-                'total': result.get('total', 0) or 0
+                "mapped": result.get("mapped", 0) or 0,
+                "unmapped": result.get("unmapped", 0) or 0,
+                "total": result.get("total", 0) or 0,
             }
 
-        return {'mapped': 0, 'unmapped': 0, 'total': 0}
+        return {"mapped": 0, "unmapped": 0, "total": 0}
 
     def manual_map(
-        self,
-        engagement_id: int,
-        agent_ip: str,
-        host_id: int
+        self, engagement_id: int, agent_ip: str, host_id: int
     ) -> Tuple[bool, int]:
         """
         Manually map an agent IP to a host.
@@ -225,7 +215,7 @@ class WazuhHostMapper:
             WHERE engagement_id = ? AND agent_ip = ? AND host_id IS NULL
         """
         result = self.db.execute_one(count_query, (engagement_id, agent_ip))
-        return result.get('count', 0) if result else 0
+        return result.get("count", 0) if result else 0
 
     def suggest_mappings(self, engagement_id: int) -> List[Dict[str, any]]:
         """
@@ -250,38 +240,40 @@ class WazuhHostMapper:
         suggestions = []
 
         for agent in unmapped:
-            agent_ip = agent.get('agent_ip')
+            agent_ip = agent.get("agent_ip")
             if not agent_ip:
                 continue
 
             # Try to find hosts in same subnet
-            agent_parts = agent_ip.split('.')
+            agent_parts = agent_ip.split(".")
             if len(agent_parts) != 4:
                 continue
 
-            agent_subnet = '.'.join(agent_parts[:3])
+            agent_subnet = ".".join(agent_parts[:3])
 
             for host in hosts:
-                host_ip = host.get('ip_address')
+                host_ip = host.get("ip_address")
                 if not host_ip:
                     continue
 
-                host_parts = host_ip.split('.')
+                host_parts = host_ip.split(".")
                 if len(host_parts) != 4:
                     continue
 
-                host_subnet = '.'.join(host_parts[:3])
+                host_subnet = ".".join(host_parts[:3])
 
                 # Same subnet = possible match
                 if agent_subnet == host_subnet:
-                    suggestions.append({
-                        'agent_ip': agent_ip,
-                        'agent_name': agent.get('agent_name'),
-                        'suggested_host_id': host['id'],
-                        'suggested_host_ip': host_ip,
-                        'suggested_host_name': host.get('hostname'),
-                        'confidence': 'medium',
-                        'reason': 'Same subnet'
-                    })
+                    suggestions.append(
+                        {
+                            "agent_ip": agent_ip,
+                            "agent_name": agent.get("agent_name"),
+                            "suggested_host_id": host["id"],
+                            "suggested_host_ip": host_ip,
+                            "suggested_host_name": host.get("hostname"),
+                            "confidence": "medium",
+                            "reason": "Same subnet",
+                        }
+                    )
 
         return suggestions