update wazuh_indexer_allocation table (#64)
taylor_socfortress committed
Jul 25, 2023 at 11:16 UTC
b51abbf3e663d7aa076a3b6c1f3aa8b20d2b4c30
2 files changed
+73
backend/app/routes/wazuhindexer.py
+14
@@ -1,5 +1,7 @@
1
from flask import Blueprint
2
+from loguru import logger
3
4
+from app import db
5
from app.services.WazuhIndexer.cluster import ClusterService
6
from app.services.WazuhIndexer.index import IndexService
7
@@ -43,6 +45,18 @@ def get_node_allocation():
45
"""
46
service = ClusterService()
47
indices = service.collect_node_allocation()
48
+ if indices["success"]:
49
+ for allocation in indices["node_allocation"]:
50
+ processed_data = service._preprocess_allocation_data(allocation)
51
+ if processed_data is None:
52
+ continue
53
+ wazuh_indexer_allocation = service._create_wazuh_indexer_allocation(*processed_data)
54
+ try:
55
+ db.session.add(wazuh_indexer_allocation)
56
+ db.session.commit()
57
+ except Exception as e:
58
+ db.session.rollback()
59
+ logger.error(f"Failed to add Wazuh-Indexer allocation to database: {e}")
60
return indices
61
62
backend/app/services/WazuhIndexer/cluster.py
+59
@@ -1,9 +1,12 @@
1
from typing import Dict
2
+from typing import Optional
3
+from typing import Tuple
4
5
# import requests
6
from elasticsearch7 import Elasticsearch
7
from loguru import logger
8
9
+from app.models.wazuh_indexer import WazuhIndexerAllocation
10
from app.services.WazuhIndexer.universal import UniversalService
11
12
@@ -221,3 +224,59 @@ class ClusterService:
224
}
225
for shard in shards
226
]
227
+
228
+ def _preprocess_allocation_data(self, allocation: Dict[str, str]) -> Optional[Tuple[str, float, float, float, float]]:
229
+ """
230
+ Preprocess the allocation data by stripping units and converting values to float.
231
+
232
+ Args:
233
+ allocation (Dict[str, str]): A dictionary containing the allocation data for a node.
234
+
235
+ Returns:
236
+ Optional[Tuple[str, float, float, float, float]]: A tuple containing the node name and the preprocessed disk usage
237
+ data. Returns None if the node is 'UNASSIGNED'.
238
+ """
239
+ # Skip if the node is 'UNASSIGNED'
240
+ if allocation["node"] == "UNASSIGNED":
241
+ return None
242
+
243
+ node = allocation["node"]
244
+
245
+ # Strip 'gb' from the disk fields and convert to float
246
+ disk_used = float(allocation["disk_used"].strip("gb"))
247
+ disk_available = float(allocation["disk_available"].strip("gb"))
248
+ disk_total = float(allocation["disk_total"].strip("gb"))
249
+
250
+ # Strip '%' from the disk_percent field and convert to float
251
+ disk_percent = float(allocation["disk_percent"].strip("%"))
252
+
253
+ return node, disk_used, disk_available, disk_total, disk_percent
254
+
255
+ def _create_wazuh_indexer_allocation(
256
+ self,
257
+ node: str,
258
+ disk_used: float,
259
+ disk_available: float,
260
+ disk_total: float,
261
+ disk_percent: float,
262
+ ) -> WazuhIndexerAllocation:
263
+ """
264
+ Create a WazuhIndexerAllocation instance.
265
+
266
+ Args:
267
+ node (str): The name of the node.
268
+ disk_used (float): The amount of disk used by the node.
269
+ disk_available (float): The amount of available disk space on the node.
270
+ disk_total (float): The total amount of disk space on the node.
271
+ disk_percent (float): The percentage of disk used on the node.
272
+
273
+ Returns:
274
+ WazuhIndexerAllocation: A WazuhIndexerAllocation instance containing the provided data.
275
+ """
276
+ return WazuhIndexerAllocation(
277
+ node=node,
278
+ disk_used=disk_used,
279
+ disk_available=disk_available,
280
+ disk_total=disk_total,
281
+ disk_percent=disk_percent,
282
+ )