Create cluster.py
taylor_socfortress committed
Jul 10, 2023 at 16:45 UTC
907288eab808193ce65d784fbd01d06b32255af1
1 file changed
+181
backend/app/services/WazuhIndexer/cluster.py
new
+181
@@ -0,0 +1,181 @@
1
+from typing import Dict
2
+import requests
3
+from elasticsearch7 import Elasticsearch
4
+from loguru import logger
5
+from app.services.WazuhIndexer.universal import UniversalService
6
+
7
+
8
+class ClusterService:
9
+ """
10
+ A service class that encapsulates the logic for pulling indices from the Wazuh-Indexer.
11
+ """
12
+
13
+ def __init__(self):
14
+ self._collect_wazuhindexer_details()
15
+ self._initialize_es_client()
16
+
17
+ def _collect_wazuhindexer_details(self):
18
+ self.connector_url, self.connector_username, self.connector_password = UniversalService().collect_wazuhindexer_details("Wazuh-Indexer")
19
+
20
+ def _initialize_es_client(self):
21
+ self.es = Elasticsearch(
22
+ [self.connector_url],
23
+ http_auth=(self.connector_username, self.connector_password),
24
+ verify_certs=False,
25
+ timeout=15,
26
+ max_retries=10,
27
+ retry_on_timeout=False,
28
+ )
29
+
30
+ def _are_details_collected(self) -> bool:
31
+ return all([self.connector_url, self.connector_username, self.connector_password])
32
+
33
+ def collect_node_allocation(self) -> Dict[str, object]:
34
+ """
35
+ Collects the node allocation from the Wazuh-Indexer.
36
+
37
+ Returns:
38
+ dict: A dictionary containing the success status, a message and potentially the index allocation.
39
+ """
40
+ if not self._are_details_collected():
41
+ return {
42
+ "message": "Failed to collect Wazuh-Indexer details",
43
+ "success": False,
44
+ }
45
+
46
+ index_summary = self._collect_node_allocation()
47
+ if not index_summary["success"]:
48
+ return index_summary
49
+
50
+ return {
51
+ "message": "Successfully collected node allocation",
52
+ "success": True,
53
+ "node_allocation": index_summary["node_allocation"],
54
+ }
55
+
56
+ def _collect_node_allocation(self) -> Dict[str, object]:
57
+ """
58
+ Collects the node allocation from the Wazuh-Indexer.
59
+
60
+ Returns:
61
+ dict: A dictionary containing the success status, a message and potentially the index allocation.
62
+ """
63
+ try:
64
+ node_allocation = self.es.cat.allocation(format="json")
65
+ node_allocation_list = self._format_node_allocation(node_allocation)
66
+ return {
67
+ "message": "Successfully collected node allocation",
68
+ "success": True,
69
+ "node_allocation": node_allocation_list,
70
+ }
71
+ except Exception as e:
72
+ logger.error(f"Failed to collect node allocation: {e}")
73
+ return {"message": "Failed to collect node allocation", "success": False}
74
+
75
+ def _format_node_allocation(self, node_allocation):
76
+ return [
77
+ {
78
+ "disk_used": node["disk.used"],
79
+ "disk_available": node["disk.avail"],
80
+ "disk_total": node["disk.total"],
81
+ "disk_percent": node["disk.percent"],
82
+ "node": node["node"],
83
+ }
84
+ for node in node_allocation
85
+ ]
86
+
87
+ def collect_cluster_health(self) -> Dict[str, object]:
88
+ """
89
+ Collects the cluster health from the Wazuh-Indexer.
90
+
91
+ Returns:
92
+ dict: A dictionary containing the success status, a message and potentially the cluster health.
93
+ """
94
+ if not self._are_details_collected():
95
+ return {
96
+ "message": "Failed to collect Wazuh-Indexer details",
97
+ "success": False,
98
+ }
99
+
100
+ index_summary = self._collect_cluster_health()
101
+ if not index_summary["success"]:
102
+ return index_summary
103
+
104
+ return {
105
+ "message": "Successfully collected cluster health",
106
+ "success": True,
107
+ "cluster_health": index_summary["cluster_health"],
108
+ }
109
+
110
+ def _collect_cluster_health(self) -> Dict[str, object]:
111
+ """
112
+ Collects the cluster health from the Wazuh-Indexer.
113
+
114
+ Returns:
115
+ dict: A dictionary containing the success status, a message and potentially the cluster health.
116
+ """
117
+ try:
118
+ cluster_health = self.es.cluster.health()
119
+ return {
120
+ "message": "Successfully collected cluster health",
121
+ "success": True,
122
+ "cluster_health": cluster_health,
123
+ }
124
+ except Exception as e:
125
+ logger.error(f"Failed to collect cluster health: {e}")
126
+ return {"message": "Failed to collect cluster health", "success": False}
127
+
128
+ def collect_shards(self) -> Dict[str, object]:
129
+ """
130
+ Collects the shards from the Wazuh-Indexer.
131
+
132
+ Returns:
133
+ dict: A dictionary containing the success status, a message and potentially the shards.
134
+ """
135
+ if not self._are_details_collected():
136
+ return {
137
+ "message": "Failed to collect Wazuh-Indexer details",
138
+ "success": False,
139
+ }
140
+
141
+ index_summary = self._collect_shards()
142
+ if not index_summary["success"]:
143
+ return index_summary
144
+
145
+ return {
146
+ "message": "Successfully collected shards",
147
+ "success": True,
148
+ "shards": index_summary["shards"],
149
+ }
150
+
151
+ def _collect_shards(self) -> Dict[str, object]:
152
+ """
153
+ Collects the shards from the Wazuh-Indexer.
154
+
155
+ Returns:
156
+ dict: A dictionary containing the success status, a message and potentially the shards.
157
+ """
158
+ try:
159
+ shards = self.es.cat.shards(format="json")
160
+ shards_list = self._format_shards(shards)
161
+ return {
162
+ "message": "Successfully collected shards",
163
+ "success": True,
164
+ "shards": shards_list,
165
+ }
166
+ except Exception as e:
167
+ logger.error(f"Failed to collect shards: {e}")
168
+ return {"message": "Failed to collect shards", "success": False}
169
+
170
+ def _format_shards(self, shards):
171
+ return [
172
+ {
173
+ "index": shard["index"],
174
+ "shard": shard["shard"],
175
+ "state": shard["state"],
176
+ "size": shard["store"],
177
+ "node": shard["node"],
178
+ }
179
+ for shard in shards
180
+ ]
181
+