| 1 | import re |
| 2 | from typing import Dict |
| 3 | from typing import Union |
| 4 | |
| 5 | from loguru import logger |
| 6 | |
| 7 | from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealth |
| 8 | from app.connectors.wazuh_indexer.schema.monitoring import ClusterHealthResponse |
| 9 | from app.connectors.wazuh_indexer.schema.monitoring import CustomerIndicesSize |
| 10 | from app.connectors.wazuh_indexer.schema.monitoring import CustomerIndicesSizeResponse |
| 11 | from app.connectors.wazuh_indexer.schema.monitoring import IndicesStats |
| 12 | from app.connectors.wazuh_indexer.schema.monitoring import IndicesStatsResponse |
| 13 | from app.connectors.wazuh_indexer.schema.monitoring import NodeAllocation |
| 14 | from app.connectors.wazuh_indexer.schema.monitoring import NodeAllocationResponse |
| 15 | from app.connectors.wazuh_indexer.schema.monitoring import Shards |
| 16 | from app.connectors.wazuh_indexer.schema.monitoring import ShardsResponse |
| 17 | from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client |
| 18 | from app.connectors.wazuh_indexer.utils.universal import format_indices_stats |
| 19 | from app.connectors.wazuh_indexer.utils.universal import format_node_allocation |
| 20 | from app.connectors.wazuh_indexer.utils.universal import format_shards |
| 21 | |
| 22 | |
| 23 | async def cluster_healthcheck() -> Union[ClusterHealthResponse, Dict[str, str]]: |
| 24 | """ |
| 25 | Returns the cluster health of the Wazuh Indexer service. |
| 26 | |
| 27 | Returns: |
| 28 | ElasticsearchResponse: A Pydantic model containing the cluster health of the Wazuh Indexer service. |
| 29 | |
| 30 | Raises: |
| 31 | Exception: An exception is raised if the cluster health cannot be retrieved. |
| 32 | """ |
| 33 | logger.info("Collecting Wazuh Indexer healthcheck") |
| 34 | es_client = await create_wazuh_indexer_client("Wazuh-Indexer") |
| 35 | try: |
| 36 | cluster_health_data = es_client.cluster.health() |
| 37 | cluster_health_model = ClusterHealth(**cluster_health_data) |
| 38 | return ClusterHealthResponse( |
| 39 | cluster_health=cluster_health_model, |
| 40 | success=True, |
| 41 | message="Successfully collected Wazuh Indexer cluster health", |
| 42 | ) |
| 43 | except Exception as e: |
| 44 | e = f"Cluster health check failed with error: {e}" |
| 45 | raise Exception(str(e)) |
| 46 | |
| 47 | |
| 48 | async def node_allocation() -> Union[NodeAllocationResponse, Dict[str, bool]]: |
| 49 | """ |
| 50 | Returns the node allocation of the Wazuh Indexer service. |
| 51 | |
| 52 | Returns: |
| 53 | ElasticsearchResponse: A Pydantic model containing the node allocation of the Wazuh Indexer service. |
| 54 | |
| 55 | Raises: |
| 56 | Exception: An exception is raised if the node allocation cannot be retrieved. |
| 57 | """ |
| 58 | logger.info("Collecting Wazuh Indexer node allocation") |
| 59 | es_client = await create_wazuh_indexer_client("Wazuh-Indexer") |
| 60 | try: |
| 61 | raw_node_allocation_data = es_client.cat.allocation(format="json") |
| 62 | logger.info(raw_node_allocation_data) |
| 63 | |
| 64 | formatted_node_allocation_data = await format_node_allocation( |
| 65 | raw_node_allocation_data, |
| 66 | ) |
| 67 | |
| 68 | node_allocation_models = [NodeAllocation(**node) for node in formatted_node_allocation_data] |
| 69 | |
| 70 | return NodeAllocationResponse( |
| 71 | node_allocation=node_allocation_models, |
| 72 | success=True, |
| 73 | message="Successfully collected Wazuh Indexer node allocation", |
| 74 | ) |
| 75 | except Exception as e: |
| 76 | e = f"Node allocation check failed with error: {e}" |
| 77 | raise Exception(str(e)) |
| 78 | |
| 79 | |
| 80 | async def indices_stats() -> Union[IndicesStatsResponse, Dict[str, str]]: |
| 81 | """ |
| 82 | Returns the indices stats of the Wazuh Indexer service. |
| 83 | |
| 84 | Returns: |
| 85 | ElasticsearchResponse: A Pydantic model containing the indices stats of the Wazuh Indexer service. |
| 86 | |
| 87 | Raises: |
| 88 | Exception: An exception is raised if the indices stats cannot be retrieved. |
| 89 | """ |
| 90 | logger.info("Collecting Wazuh Indexer indices stats") |
| 91 | es_client = await create_wazuh_indexer_client("Wazuh-Indexer") |
| 92 | try: |
| 93 | raw_indices_stats_data = es_client.cat.indices(format="json") |
| 94 | |
| 95 | formatted_indices_stats_data = await format_indices_stats( |
| 96 | raw_indices_stats_data, |
| 97 | ) |
| 98 | |
| 99 | indices_stats_models = [IndicesStats(**index) for index in formatted_indices_stats_data] |
| 100 | |
| 101 | return IndicesStatsResponse( |
| 102 | indices_stats=indices_stats_models, |
| 103 | success=True, |
| 104 | message="Successfully collected Wazuh Indexer indices stats", |
| 105 | ) |
| 106 | except Exception as e: |
| 107 | e = f"Indices stats check failed with error: {e}" |
| 108 | raise Exception(str(e)) |
| 109 | |
| 110 | |
| 111 | async def shards() -> Union[ShardsResponse, Dict[str, str]]: |
| 112 | """ |
| 113 | Returns the shards of the Wazuh Indexer service. |
| 114 | |
| 115 | Returns: |
| 116 | ElasticsearchResponse: A Pydantic model containing the shards of the Wazuh Indexer service. |
| 117 | |
| 118 | Raises: |
| 119 | Exception: An exception is raised if the shards cannot be retrieved. |
| 120 | """ |
| 121 | logger.info("Collecting Wazuh Indexer shards") |
| 122 | es_client = await create_wazuh_indexer_client("Wazuh-Indexer") |
| 123 | try: |
| 124 | raw_shards_data = es_client.cat.shards(format="json") |
| 125 | |
| 126 | formatted_shards_data = await format_shards(raw_shards_data) |
| 127 | |
| 128 | shard_models = [Shards(**shard) for shard in formatted_shards_data] |
| 129 | |
| 130 | return ShardsResponse( |
| 131 | shards=shard_models, |
| 132 | success=True, |
| 133 | message="Successfully collected Wazuh Indexer shards", |
| 134 | ) |
| 135 | except Exception as e: |
| 136 | logger.error(f"Shards check failed with error: {e}") |
| 137 | e = f"Shards check failed with error: {e}" |
| 138 | raise Exception(str(e)) |
| 139 | |
| 140 | |
| 141 | async def output_shard_number_to_be_set_based_on_nodes() -> int: |
| 142 | """ |
| 143 | Retrieves the number of nodes in the Wazuh Indexer cluster. |
| 144 | Based on that number, it returns the number of shards to be set for the new index. |
| 145 | This is a 1:1 mapping between the number of nodes and the number of shards. |
| 146 | |
| 147 | Returns: |
| 148 | int: The number of shards to be set for the new index. |
| 149 | """ |
| 150 | es_client = await create_wazuh_indexer_client("Wazuh-Indexer") |
| 151 | try: |
| 152 | cluster_health_data = es_client.cluster.health() |
| 153 | cluster_health_model = ClusterHealth(**cluster_health_data) |
| 154 | return cluster_health_model.number_of_nodes |
| 155 | except Exception as e: |
| 156 | logger.error(f"Shards check failed with error: {e}") |
| 157 | e = f"Shards check failed with error: {e}" |
| 158 | raise Exception(str(e)) |
| 159 | |
| 160 | |
| 161 | def parse_size_to_bytes(size_str: str) -> int: |
| 162 | """ |
| 163 | Convert a human-readable size string to bytes. |
| 164 | Handles formats like '1.2gb', '500mb', '100kb', '1024b'. |
| 165 | """ |
| 166 | if not size_str or size_str == "Store size not found": |
| 167 | return 0 |
| 168 | |
| 169 | size_str = size_str.lower().strip() |
| 170 | |
| 171 | # Define multipliers |
| 172 | multipliers = { |
| 173 | "b": 1, |
| 174 | "kb": 1024, |
| 175 | "mb": 1024**2, |
| 176 | "gb": 1024**3, |
| 177 | "tb": 1024**4, |
| 178 | } |
| 179 | |
| 180 | # Match number and unit |
| 181 | match = re.match(r"^([\d.]+)\s*([a-z]+)$", size_str) |
| 182 | if match: |
| 183 | value = float(match.group(1)) |
| 184 | unit = match.group(2) |
| 185 | return int(value * multipliers.get(unit, 1)) |
| 186 | |
| 187 | # Try to parse as pure number (bytes) |
| 188 | try: |
| 189 | return int(float(size_str)) |
| 190 | except ValueError: |
| 191 | return 0 |
| 192 | |
| 193 | |
| 194 | def bytes_to_human_readable(size_bytes: int) -> str: |
| 195 | """Convert bytes to human-readable format.""" |
| 196 | for unit in ["b", "kb", "mb", "gb", "tb"]: |
| 197 | if abs(size_bytes) < 1024.0: |
| 198 | return f"{size_bytes:.2f}{unit}" |
| 199 | size_bytes /= 1024.0 |
| 200 | return f"{size_bytes:.2f}pb" |
| 201 | |
| 202 | |
| 203 | def extract_customer_from_index(index_name: str) -> str: |
| 204 | """ |
| 205 | Extract customer name from index name. |
| 206 | Pattern: after dash or underscore, before the next underscore or end. |
| 207 | Examples: |
| 208 | - wazuh-copilot_37 -> copilot |
| 209 | - dev-taylor_37 -> taylor |
| 210 | - wazuh-509dine2v_0 -> 509dine2v |
| 211 | """ |
| 212 | # Match pattern: prefix-customer_suffix or prefix_customer_suffix |
| 213 | match = re.match(r"^[^-_]+-([^_]+)_", index_name) |
| 214 | if match: |
| 215 | return match.group(1) |
| 216 | |
| 217 | # Fallback: try underscore as first separator |
| 218 | match = re.match(r"^[^_]+_([^_]+)_", index_name) |
| 219 | if match: |
| 220 | return match.group(1) |
| 221 | |
| 222 | return "unknown" |
| 223 | |
| 224 | |
| 225 | async def indices_size_per_customer() -> Union[CustomerIndicesSizeResponse, Dict[str, str]]: |
| 226 | """ |
| 227 | Returns the total indices size aggregated per customer. |
| 228 | |
| 229 | Returns: |
| 230 | CustomerIndicesSizeResponse: A Pydantic model containing the indices size per customer. |
| 231 | |
| 232 | Raises: |
| 233 | Exception: An exception is raised if the indices stats cannot be retrieved. |
| 234 | """ |
| 235 | logger.info("Collecting Wazuh Indexer indices size per customer") |
| 236 | es_client = await create_wazuh_indexer_client("Wazuh-Indexer") |
| 237 | try: |
| 238 | raw_indices_stats_data = es_client.cat.indices(format="json") |
| 239 | |
| 240 | formatted_indices_stats_data = await format_indices_stats(raw_indices_stats_data) |
| 241 | |
| 242 | # Aggregate by customer |
| 243 | customer_data: Dict[str, Dict] = {} |
| 244 | |
| 245 | for index_data in formatted_indices_stats_data: |
| 246 | index_name = index_data.get("index", "") |
| 247 | store_size = index_data.get("store_size", "0b") |
| 248 | |
| 249 | customer = extract_customer_from_index(index_name) |
| 250 | size_bytes = parse_size_to_bytes(store_size) |
| 251 | |
| 252 | if customer not in customer_data: |
| 253 | customer_data[customer] = { |
| 254 | "total_size_bytes": 0, |
| 255 | "index_count": 0, |
| 256 | "indices": [], |
| 257 | } |
| 258 | |
| 259 | customer_data[customer]["total_size_bytes"] += size_bytes |
| 260 | customer_data[customer]["index_count"] += 1 |
| 261 | customer_data[customer]["indices"].append(index_name) |
| 262 | |
| 263 | # Convert to response models |
| 264 | customer_sizes = [ |
| 265 | CustomerIndicesSize( |
| 266 | customer=customer, |
| 267 | total_size_bytes=data["total_size_bytes"], |
| 268 | total_size_human=bytes_to_human_readable(data["total_size_bytes"]), |
| 269 | index_count=data["index_count"], |
| 270 | indices=data["indices"], |
| 271 | ) |
| 272 | for customer, data in sorted(customer_data.items()) |
| 273 | ] |
| 274 | |
| 275 | return CustomerIndicesSizeResponse( |
| 276 | customer_sizes=customer_sizes, |
| 277 | success=True, |
| 278 | message="Successfully collected Wazuh Indexer indices size per customer", |
| 279 | ) |
| 280 | except Exception as e: |
| 281 | logger.error(f"Indices size per customer check failed with error: {e}") |
| 282 | raise Exception(f"Indices size per customer check failed with error: {e}") |