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

souleyez/integrations/wazuh/client.py

@@ -4,6 +4,7 @@ Wazuh API Client for SoulEyez Detection Validation.
 Connects to Wazuh Manager API (port 55000) for management operations
 and Wazuh Indexer API (port 9200) for querying alerts.
 """
+
 import requests
 from datetime import datetime, timedelta
 from typing import List, Dict, Optional, Any
@@ -24,7 +25,7 @@ class WazuhClient:
         verify_ssl: bool = False,
         indexer_url: Optional[str] = None,
         indexer_user: Optional[str] = None,
-        indexer_password: Optional[str] = None
+        indexer_password: Optional[str] = None,
     ):
         """
         Initialize Wazuh API client.
@@ -38,7 +39,7 @@ class WazuhClient:
             indexer_user: Indexer username (defaults to 'admin')
             indexer_password: Indexer password (defaults to manager password)
         """
-        self.api_url = api_url.rstrip('/')
+        self.api_url = api_url.rstrip("/")
         self.username = username
         self.password = password
         self.verify_ssl = verify_ssl
@@ -48,13 +49,13 @@ class WazuhClient:
         # Derive indexer URL from manager URL if not provided
         # Manager is on :55000, Indexer is on :9200
         if indexer_url:
-            self.indexer_url = indexer_url.rstrip('/')
+            self.indexer_url = indexer_url.rstrip("/")
         else:
             # Replace port 55000 with 9200
-            self.indexer_url = self.api_url.replace(':55000', ':9200')
+            self.indexer_url = self.api_url.replace(":55000", ":9200")
 
         # Indexer credentials (often different from Manager)
-        self.indexer_user = indexer_user or 'admin'
+        self.indexer_user = indexer_user or "admin"
         self.indexer_password = indexer_password or password
 
     def _get_token(self) -> str:
@@ -66,15 +67,12 @@ class WazuhClient:
         # Authenticate
         url = f"{self.api_url}/security/user/authenticate"
         response = requests.post(
-            url,
-            auth=(self.username, self.password),
-            verify=self.verify_ssl,
-            timeout=30
+            url, auth=(self.username, self.password), verify=self.verify_ssl, timeout=30
         )
         response.raise_for_status()
 
         data = response.json()
-        self._token = data.get('data', {}).get('token')
+        self._token = data.get("data", {}).get("token")
         # Tokens typically expire in 15 minutes, refresh at 10
         self._token_expiry = datetime.now() + timedelta(minutes=10)
 
@@ -88,7 +86,7 @@ class WazuhClient:
         method: str,
         endpoint: str,
         params: Optional[Dict] = None,
-        json_data: Optional[Dict] = None
+        json_data: Optional[Dict] = None,
     ) -> Dict[str, Any]:
         """Make authenticated API request."""
         token = self._get_token()
@@ -96,7 +94,7 @@ class WazuhClient:
 
         headers = {
             "Authorization": f"Bearer {token}",
-            "Content-Type": "application/json"
+            "Content-Type": "application/json",
         }
 
         response = requests.request(
@@ -106,7 +104,7 @@ class WazuhClient:
             params=params,
             json=json_data,
             verify=self.verify_ssl,
-            timeout=60
+            timeout=60,
         )
         response.raise_for_status()
         return response.json()
@@ -128,7 +126,9 @@ class WazuhClient:
             # Get hostname from agent 000 (manager)
             hostname = "unknown"
             try:
-                agent_result = self._request("GET", "/agents", params={"agents_list": "000"})
+                agent_result = self._request(
+                    "GET", "/agents", params={"agents_list": "000"}
+                )
                 agent_items = agent_result.get("data", {}).get("affected_items", [])
                 if agent_items:
                     hostname = agent_items[0].get("name", "unknown")
@@ -139,7 +139,7 @@ class WazuhClient:
                 "connected": True,
                 "version": info.get("version", "unknown"),
                 "cluster": data.get("cluster", {}).get("enabled", False),
-                "hostname": hostname
+                "hostname": hostname,
             }
         except requests.exceptions.ConnectionError as e:
             return {"connected": False, "error": f"Connection failed: {str(e)}"}
@@ -155,7 +155,7 @@ class WazuhClient:
         agent_ip: Optional[str] = None,
         rule_ids: Optional[List[int]] = None,
         limit: int = 100,
-        search_text: Optional[str] = None
+        search_text: Optional[str] = None,
     ) -> List[Dict[str, Any]]:
         """
         Query alerts from Wazuh Indexer (Elasticsearch-based).
@@ -178,9 +178,9 @@ class WazuhClient:
         if start_time or end_time:
             time_range = {}
             if start_time:
-                time_range["gte"] = start_time.strftime('%Y-%m-%dT%H:%M:%S')
+                time_range["gte"] = start_time.strftime("%Y-%m-%dT%H:%M:%S")
             if end_time:
-                time_range["lte"] = end_time.strftime('%Y-%m-%dT%H:%M:%S')
+                time_range["lte"] = end_time.strftime("%Y-%m-%dT%H:%M:%S")
             must_clauses.append({"range": {"timestamp": time_range}})
 
         # Rule ID filter
@@ -189,28 +189,26 @@ class WazuhClient:
 
         # Agent IP filter (check various IP fields)
         if agent_ip:
-            must_clauses.append({
-                "multi_match": {
-                    "query": agent_ip,
-                    "fields": ["agent.ip", "data.srcip", "data.src_ip", "src_ip"]
+            must_clauses.append(
+                {
+                    "multi_match": {
+                        "query": agent_ip,
+                        "fields": ["agent.ip", "data.srcip", "data.src_ip", "src_ip"],
+                    }
                 }
-            })
+            )
 
         # Free text search
         if search_text:
-            must_clauses.append({
-                "query_string": {"query": f"*{search_text}*"}
-            })
+            must_clauses.append({"query_string": {"query": f"*{search_text}*"}})
 
         # Build the query
         query = {
             "size": limit,
             "sort": [{"timestamp": {"order": "desc"}}],
             "query": {
-                "bool": {
-                    "must": must_clauses if must_clauses else [{"match_all": {}}]
-                }
-            }
+                "bool": {"must": must_clauses if must_clauses else [{"match_all": {}}]}
+            },
         }
 
         try:
@@ -220,7 +218,7 @@ class WazuhClient:
                 auth=(self.indexer_user, self.indexer_password),
                 json=query,
                 verify=self.verify_ssl,
-                timeout=30
+                timeout=30,
             )
             response.raise_for_status()
             data = response.json()
@@ -239,11 +237,7 @@ class WazuhClient:
             raise
 
     def get_alerts_by_src_ip(
-        self,
-        src_ip: str,
-        start_time: datetime,
-        end_time: datetime,
-        limit: int = 100
+        self, src_ip: str, start_time: datetime, end_time: datetime, limit: int = 100
     ) -> List[Dict[str, Any]]:
         """
         Get alerts where source IP matches (attacker IP).
@@ -252,10 +246,7 @@ class WazuhClient:
         we look for alerts triggered by our attack source IP.
         """
         return self.get_alerts(
-            start_time=start_time,
-            end_time=end_time,
-            agent_ip=src_ip,
-            limit=limit
+            start_time=start_time, end_time=end_time, agent_ip=src_ip, limit=limit
         )
 
     def get_rules(self, rule_ids: Optional[List[int]] = None) -> List[Dict[str, Any]]:
@@ -267,7 +258,9 @@ class WazuhClient:
         result = self._request("GET", "/rules", params=params)
         return result.get("data", {}).get("affected_items", [])
 
-    def get_rule_file_content(self, filename: str, relative_dirname: str = None) -> Optional[str]:
+    def get_rule_file_content(
+        self, filename: str, relative_dirname: str = None
+    ) -> Optional[str]:
         """
         Get the content of a rule file.
 
@@ -297,13 +290,13 @@ class WazuhClient:
                 headers=headers,
                 params=params,
                 verify=self.verify_ssl,
-                timeout=30
+                timeout=30,
             )
 
             if response.status_code == 200:
                 # Check if response is JSON or raw text
-                content_type = response.headers.get('Content-Type', '')
-                if 'application/json' in content_type:
+                content_type = response.headers.get("Content-Type", "")
+                if "application/json" in content_type:
                     # JSON response - extract content from data
                     data = response.json()
                     affected_items = data.get("data", {}).get("affected_items", [])
@@ -311,7 +304,7 @@ class WazuhClient:
                         item = affected_items[0]
                         # The content might be in a 'content' field
                         if isinstance(item, dict):
-                            return item.get('content') or item.get('data')
+                            return item.get("content") or item.get("data")
                         return item if isinstance(item, str) else None
                 else:
                     # Raw text response
@@ -320,7 +313,9 @@ class WazuhClient:
         except Exception:
             return None
 
-    def get_rule_xml(self, rule_id: int, filename: str, relative_dirname: str = None) -> Optional[str]:
+    def get_rule_xml(
+        self, rule_id: int, filename: str, relative_dirname: str = None
+    ) -> Optional[str]:
         """
         Extract a specific rule's XML from a rule file.
 
@@ -357,10 +352,7 @@ class WazuhClient:
     # =========================================================================
 
     def get_agent_vulnerabilities(
-        self,
-        agent_id: str,
-        severity: Optional[str] = None,
-        limit: int = 1000
+        self, agent_id: str, severity: Optional[str] = None, limit: int = 1000
     ) -> List[Dict[str, Any]]:
         """
         Fetch vulnerabilities for a specific agent from Wazuh Indexer.
@@ -375,9 +367,7 @@ class WazuhClient:
         Returns:
             List of vulnerability dictionaries
         """
-        must_clauses = [
-            {"term": {"agent.id": agent_id}}
-        ]
+        must_clauses = [{"term": {"agent.id": agent_id}}]
 
         if severity:
             must_clauses.append({"term": {"vulnerability.severity": severity}})
@@ -385,11 +375,7 @@ class WazuhClient:
         query = {
             "size": limit,
             "sort": [{"vulnerability.severity": {"order": "desc"}}],
-            "query": {
-                "bool": {
-                    "must": must_clauses
-                }
-            }
+            "query": {"bool": {"must": must_clauses}},
         }
 
         try:
@@ -398,13 +384,15 @@ class WazuhClient:
                 auth=(self.indexer_user, self.indexer_password),
                 json=query,
                 verify=self.verify_ssl,
-                timeout=60
+                timeout=60,
             )
             response.raise_for_status()
             data = response.json()
 
             hits = data.get("hits", {}).get("hits", [])
-            return [self._normalize_vulnerability(hit.get("_source", {})) for hit in hits]
+            return [
+                self._normalize_vulnerability(hit.get("_source", {})) for hit in hits
+            ]
 
         except requests.exceptions.ConnectionError:
             return []
@@ -417,7 +405,7 @@ class WazuhClient:
         self,
         severity: Optional[str] = None,
         agent_ids: Optional[List[str]] = None,
-        limit: int = 5000
+        limit: int = 5000,
     ) -> List[Dict[str, Any]]:
         """
         Fetch all vulnerabilities from Wazuh Indexer.
@@ -442,13 +430,11 @@ class WazuhClient:
             "size": limit,
             "sort": [
                 {"vulnerability.severity": {"order": "desc"}},
-                {"agent.id": {"order": "asc"}}
+                {"agent.id": {"order": "asc"}},
             ],
             "query": {
-                "bool": {
-                    "must": must_clauses if must_clauses else [{"match_all": {}}]
-                }
-            }
+                "bool": {"must": must_clauses if must_clauses else [{"match_all": {}}]}
+            },
         }
 
         try:
@@ -457,13 +443,15 @@ class WazuhClient:
                 auth=(self.indexer_user, self.indexer_password),
                 json=query,
                 verify=self.verify_ssl,
-                timeout=120
+                timeout=120,
             )
             response.raise_for_status()
             data = response.json()
 
             hits = data.get("hits", {}).get("hits", [])
-            return [self._normalize_vulnerability(hit.get("_source", {})) for hit in hits]
+            return [
+                self._normalize_vulnerability(hit.get("_source", {})) for hit in hits
+            ]
 
         except requests.exceptions.ConnectionError:
             return []
@@ -491,12 +479,10 @@ class WazuhClient:
                         "by_severity": {
                             "terms": {"field": "vulnerability.severity", "size": 10}
                         }
-                    }
+                    },
                 },
-                "total": {
-                    "value_count": {"field": "vulnerability.id"}
-                }
-            }
+                "total": {"value_count": {"field": "vulnerability.id"}},
+            },
         }
 
         try:
@@ -505,7 +491,7 @@ class WazuhClient:
                 auth=(self.indexer_user, self.indexer_password),
                 json=query,
                 verify=self.verify_ssl,
-                timeout=30
+                timeout=30,
             )
             response.raise_for_status()
             data = response.json()
@@ -523,15 +509,17 @@ class WazuhClient:
                 agent_id = bucket["key"]
                 agent_breakdown[agent_id] = {
                     "total": bucket["doc_count"],
-                    "by_severity": {}
+                    "by_severity": {},
                 }
                 for sev_bucket in bucket.get("by_severity", {}).get("buckets", []):
-                    agent_breakdown[agent_id]["by_severity"][sev_bucket["key"]] = sev_bucket["doc_count"]
+                    agent_breakdown[agent_id]["by_severity"][sev_bucket["key"]] = (
+                        sev_bucket["doc_count"]
+                    )
 
             return {
                 "total": aggs.get("total", {}).get("value", 0),
                 "by_severity": severity_counts,
-                "by_agent": agent_breakdown
+                "by_agent": agent_breakdown,
             }
 
         except requests.exceptions.ConnectionError:
@@ -551,25 +539,19 @@ class WazuhClient:
                 f"{self.indexer_url}/wazuh-states-vulnerabilities-*/_count",
                 auth=(self.indexer_user, self.indexer_password),
                 verify=self.verify_ssl,
-                timeout=10
+                timeout=10,
             )
 
             if response.status_code == 200:
                 data = response.json()
-                return {
-                    "accessible": True,
-                    "count": data.get("count", 0)
-                }
+                return {"accessible": True, "count": data.get("count", 0)}
             elif response.status_code == 404:
                 return {
                     "accessible": False,
-                    "error": "Vulnerability index not found. Ensure vulnerability detection is enabled in Wazuh."
+                    "error": "Vulnerability index not found. Ensure vulnerability detection is enabled in Wazuh.",
                 }
             else:
-                return {
-                    "accessible": False,
-                    "error": f"HTTP {response.status_code}"
-                }
+                return {"accessible": False, "error": f"HTTP {response.status_code}"}
 
         except requests.exceptions.ConnectionError as e:
             return {"accessible": False, "error": f"Connection failed: {str(e)}"}
@@ -609,7 +591,6 @@ class WazuhClient:
             "agent_id": agent.get("id"),
             "agent_name": agent.get("name"),
             "agent_ip": agent.get("ip"),
-
             # Vulnerability info
             "cve_id": vuln.get("id"),
             "name": vuln.get("title") or vuln.get("description", "")[:100],
@@ -617,18 +598,14 @@ class WazuhClient:
             "cvss_score": cvss_score,
             "cvss_version": cvss_version,
             "published_date": vuln.get("published"),
-
             # Package info
             "package_name": package.get("name"),
             "package_version": package.get("version"),
             "package_architecture": package.get("architecture"),
-
             # References
             "reference_urls": vuln.get("reference", []),
-
             # Detection time
             "detection_time": vuln.get("detected_at"),
-
             # Raw data for debugging
-            "raw_data": doc
+            "raw_data": doc,
         }