main
py 656 lines 22.7 KB
Raw
1 import re
2 from datetime import datetime
3 from datetime import timedelta
4 from typing import Any
5 from typing import Dict
6 from typing import Iterable
7 from typing import Tuple
8
9 from elasticsearch7 import AsyncElasticsearch
10 from elasticsearch7 import Elasticsearch
11 from fastapi import HTTPException
12 from loguru import logger
13
14 from app.connectors.utils import get_connector_info_from_db
15 from app.connectors.wazuh_indexer.schema.indices import IndexConfigModel
16 from app.connectors.wazuh_indexer.schema.indices import Indices
17 from app.db.db_session import get_db_session
18
19
20 async def verify_wazuh_indexer_credentials(
21 attributes: Dict[str, Any],
22 ) -> Dict[str, Any]:
23 """
24 Verifies the connection to Wazuh Indexer service.
25
26 Returns:
27 dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
28 """
29 logger.info(
30 f"Verifying the wazuh-indexer connection to {attributes['connector_url']}",
31 )
32 try:
33 es = Elasticsearch(
34 [attributes["connector_url"]],
35 http_auth=(
36 attributes["connector_username"],
37 attributes["connector_password"],
38 ),
39 verify_certs=False,
40 timeout=15,
41 max_retries=10,
42 retry_on_timeout=False,
43 )
44 es.cluster.health()
45 logger.debug("Wazuh Indexer connection successful")
46 return {
47 "connectionSuccessful": True,
48 "message": "Wazuh Indexer connection successful",
49 }
50 except Exception as e:
51 logger.error(
52 f"Connection to {attributes['connector_url']} failed with error: {e}",
53 )
54 return {
55 "connectionSuccessful": False,
56 "message": f"Connection to {attributes['connector_url']} failed with error: {e}",
57 }
58
59
60 async def verify_wazuh_indexer_connection(connector_name: str) -> str:
61 """
62 Returns the authentication token for the Wazuh Indexer service.
63
64 Returns:
65 str: Authentication token for the Wazuh Indexer service.
66 """
67 async with get_db_session() as session: # This will correctly enter the context manager
68 attributes = await get_connector_info_from_db(connector_name, session)
69 logger.info(
70 f"Verifying the wazuh-indexer connection to {attributes['connector_url']}",
71 )
72 if attributes is None:
73 logger.error("No Wazuh Indexer connector found in the database")
74 return None
75 return await verify_wazuh_indexer_credentials(attributes)
76
77
78 async def create_wazuh_indexer_client(connector_name: str = "Wazuh-Indexer") -> Elasticsearch:
79 """
80 Returns an Elasticsearch client for the Wazuh Indexer service.
81
82 Returns:
83 Elasticsearch: Elasticsearch client for the Wazuh Indexer service.
84 """
85 # attributes = get_connector_info_from_db(connector_name)
86 async with get_db_session() as session: # This will correctly enter the context manager
87 attributes = await get_connector_info_from_db(connector_name, session)
88 if attributes is None:
89 raise HTTPException(
90 status_code=500,
91 detail=f"No {connector_name} connector found in the database",
92 )
93 if attributes["connector_url"] == "https://127.1.1.1:9200":
94 raise HTTPException(
95 status_code=500,
96 detail=f"Please update the {connector_name} connector URL",
97 )
98 try:
99 return Elasticsearch(
100 [attributes["connector_url"]],
101 http_auth=(
102 attributes["connector_username"],
103 attributes["connector_password"],
104 ),
105 verify_certs=False,
106 timeout=15,
107 max_retries=10,
108 retry_on_timeout=False,
109 )
110 except Exception as e:
111 raise HTTPException(
112 status_code=500,
113 detail=f"Failed to create Elasticsearch client: {e}",
114 )
115
116
117 async def create_wazuh_indexer_client_async(connector_name: str = "Wazuh-Indexer") -> AsyncElasticsearch:
118 """
119 Returns an Elasticsearch client for the Wazuh Indexer service.
120
121 Returns:
122 Elasticsearch: Elasticsearch client for the Wazuh Indexer service.
123 """
124 # attributes = get_connector_info_from_db(connector_name)
125 async with get_db_session() as session: # This will correctly enter the context manager
126 attributes = await get_connector_info_from_db(connector_name, session)
127 if attributes is None:
128 raise HTTPException(
129 status_code=500,
130 detail=f"No {connector_name} connector found in the database",
131 )
132 if attributes["connector_url"] == "https://127.1.1.1:9200":
133 raise HTTPException(
134 status_code=500,
135 detail=f"Please update the {connector_name} connector URL",
136 )
137 try:
138 return AsyncElasticsearch(
139 [attributes["connector_url"]],
140 http_auth=(
141 attributes["connector_username"],
142 attributes["connector_password"],
143 ),
144 verify_certs=False,
145 timeout=15,
146 max_retries=10,
147 retry_on_timeout=False,
148 )
149 except Exception as e:
150 raise HTTPException(
151 status_code=500,
152 detail=f"Failed to create Elasticsearch client: {e}",
153 )
154
155
156 async def format_node_allocation(node_allocation):
157 """
158 Format the node allocation details into a list of dictionaries. Each dictionary contains disk used, disk available, total disk, disk
159 usage percentage, and node name.
160
161 Args:
162 node_allocation: Node allocation details from Elasticsearch.
163
164 Returns:
165 list: A list of dictionaries containing formatted node allocation details.
166 """
167 return [
168 {
169 "disk_used": node["disk.used"],
170 "disk_available": node["disk.avail"],
171 "disk_total": node["disk.total"],
172 "disk_percent": node["disk.percent"],
173 "node": node["node"],
174 }
175 for node in node_allocation
176 ]
177
178
179 async def format_indices_stats(indices_stats):
180 """
181 Format the indices stats details into a list of dictionaries. Each dictionary contains the index name, the number of documents in the index,
182 the size of the index, and the number of shards in the index.
183
184 Args:
185 indices_stats: Indices stats details from Elasticsearch.
186
187 Returns:
188 list: A list of dictionaries containing formatted indices stats details.
189 """
190 return [
191 {
192 "index": index["index"],
193 "docs_count": index["docs.count"],
194 "store_size": index["store.size"],
195 "replica_count": index["rep"],
196 "health": index["health"],
197 }
198 for index in indices_stats
199 ]
200
201
202 async def format_shards(shards):
203 """
204 Format the shards details into a list of dictionaries. Each dictionary contains the index name, the shard number, the shard state, the shard
205 size, and the node name.
206
207 Args:
208 shards: Shards details from Elasticsearch.
209
210 Returns:
211 list: A list of dictionaries containing formatted shards details.
212 """
213 return [
214 {
215 "index": shard["index"],
216 "shard": shard["shard"],
217 "state": shard["state"],
218 "size": shard["store"],
219 "node": shard["node"],
220 }
221 for shard in shards
222 ]
223
224
225 # async def collect_indices() -> Indices:
226 # """
227 # Collects the indices from Elasticsearch.
228
229 # Returns:
230 # dict: A dictionary containing the indices, shards, and indices stats.
231 # """
232 # logger.info("Collecting indices from Elasticsearch")
233 # es = await create_wazuh_indexer_client("Wazuh-Indexer")
234 # try:
235 # indices_dict = es.indices.get_alias("*", expand_wildcards="open")
236 # indices_list = list(indices_dict.keys())
237 # # Check if the index is valid
238 # index_config = IndexConfigModel()
239 # indices_list = [index for index in indices_list if index_config.is_valid_index(index)]
240 # return Indices(
241 # indices_list=indices_list,
242 # success=True,
243 # message="Indices collected successfully",
244 # )
245 # except Exception as e:
246 # logger.error(f"Failed to collect indices: {e}")
247 # raise HTTPException(status_code=500, detail=f"Failed to collect indices: {e}")
248
249
250 async def collect_indices(all_indices: bool = False) -> Indices:
251 """
252 Collects the indices from Elasticsearch.
253
254 Args:
255 all_indices (bool, optional): If True, all indices are listed. If False, only valid indices are listed. Defaults to False.
256
257 Returns:
258 dict: A dictionary containing the indices, shards, and indices stats.
259 """
260 logger.info("Collecting indices from Elasticsearch")
261 es = await create_wazuh_indexer_client("Wazuh-Indexer")
262 try:
263 indices_dict = es.indices.get_alias("*", expand_wildcards="open")
264 indices_list = list(indices_dict.keys())
265 # Check if the index is valid
266 if not all_indices:
267 index_config = IndexConfigModel()
268 indices_list = [index for index in indices_list if index_config.is_valid_index(index)]
269 return Indices(
270 indices_list=indices_list,
271 success=True,
272 message="Indices collected successfully",
273 )
274 except Exception as e:
275 logger.error(f"Failed to collect indices: {e}")
276 raise HTTPException(status_code=500, detail=f"Failed to collect indices: {e}")
277
278
279 class AlertsQueryBuilder:
280 @staticmethod
281 def _get_time_range_start(timerange: str) -> str:
282 """
283 Determines the start time of the time range based on the current time and the provided timerange.
284
285 Args:
286 timerange (str): The time range to collect alerts from. This is a string like "24h", "1w", etc.
287
288 Returns:
289 str: A string representing the start time of the time range in ISO format.
290 """
291 if timerange.endswith("h"):
292 delta = timedelta(hours=int(timerange[:-1]))
293 elif timerange.endswith("d"):
294 delta = timedelta(days=int(timerange[:-1]))
295 elif timerange.endswith("w"):
296 delta = timedelta(weeks=int(timerange[:-1]))
297 else:
298 raise ValueError(
299 "Invalid timerange format. Expected a string like '24h', '1d', '1w', etc.",
300 )
301
302 start = datetime.utcnow() - delta
303 return start.isoformat() + "Z" # Elasticsearch expects the time in ISO format with a Z at the end
304
305 def __init__(self):
306 self.query = {
307 "query": {
308 "bool": {
309 "must": [],
310 },
311 },
312 "sort": [],
313 }
314
315 def add_time_range(self, timerange: str, timestamp_field: str):
316 """
317 Adds a time range filter to the query.
318
319 Args:
320 timerange (str): The time range to filter by.
321 timestamp_field (str): The name of the timestamp field in the index.
322
323 Returns:
324 self: The updated instance of the class.
325 """
326 start = self._get_time_range_start(timerange)
327 range_query = {
328 "range": {
329 timestamp_field: {
330 "gte": start,
331 "lte": "now",
332 },
333 },
334 }
335 if timestamp_field == "timestamp":
336 range_query["range"][timestamp_field]["format"] = "strict_date_optional_time"
337 self.query["query"]["bool"]["must"].append(range_query)
338 return self
339
340 def add_absolute_time_range(self, time_from: str, time_to: str, timestamp_field: str):
341 """
342 Adds an absolute time range filter using ISO-formatted start/end times.
343
344 Args:
345 time_from (str): Start time in ISO format (e.g. '2025-01-01T00:00:00Z').
346 time_to (str): End time in ISO format (e.g. '2025-01-31T23:59:59Z').
347 timestamp_field (str): The name of the timestamp field in the index.
348
349 Returns:
350 self: The updated instance of the class.
351 """
352 range_query = {
353 "range": {
354 timestamp_field: {
355 "gte": time_from,
356 "lte": time_to,
357 },
358 },
359 }
360 if timestamp_field == "timestamp":
361 range_query["range"][timestamp_field]["format"] = "strict_date_optional_time"
362 self.query["query"]["bool"]["must"].append(range_query)
363 return self
364
365 def add_matches(self, matches: Iterable[Tuple[str, str]]):
366 """
367 Adds matches to the query.
368
369 Args:
370 matches (Iterable[Tuple[str, str]]): A collection of field-value pairs to match.
371
372 Returns:
373 self: The current instance of the class.
374 """
375 for field, value in matches:
376 self.query["query"]["bool"]["must"].append({"match": {field: value}})
377 return self
378
379 def add_match_phrase(self, matches: Iterable[Tuple[str, str]]):
380 """
381 Adds match phrases to the query.
382
383 Args:
384 matches (Iterable[Tuple[str, str]]): A collection of field-value pairs to match.
385
386 Returns:
387 self: The instance of the class.
388
389 """
390 for field, value in matches:
391 self.query["query"]["bool"]["must"].append({"match_phrase": {field: value}})
392 return self
393
394 def add_range(self, field: str, value: str):
395 """
396 Adds a range query to the Elasticsearch query.
397
398 Args:
399 field (str): The field to apply the range query on.
400 value (str): The value to compare against in the range query.
401
402 Returns:
403 self: The current instance of the class.
404 """
405 self.query["query"]["bool"]["must"].append({"range": {field: {"gte": value}}})
406 return self
407
408 def add_sort(self, field: str, order: str = "desc"):
409 """
410 Add a sort field to the query.
411
412 Args:
413 field (str): The field to sort by.
414 order (str, optional): The sort order. Defaults to "desc".
415
416 Returns:
417 self: The updated instance of the class.
418 """
419 self.query["sort"].append({field: {"order": order}})
420 return self
421
422 def build(self):
423 """
424 Builds and returns the query.
425
426 Returns:
427 str: The built query.
428 """
429 return self.query
430
431
432 class LogsQueryBuilder:
433 @staticmethod
434 def _get_time_range_start(timerange: str) -> str:
435 """
436 Determines the start time of the time range based on the current time and the provided timerange.
437
438 Args:
439 timerange (str): The time range to collect alerts from. This is a string like "24h", "1w", etc.
440
441 Returns:
442 str: A string representing the start time of the time range in ISO format.
443 """
444 if timerange.endswith("m"):
445 delta = timedelta(minutes=int(timerange[:-1]))
446 elif timerange.endswith("h"):
447 delta = timedelta(hours=int(timerange[:-1]))
448 elif timerange.endswith("d"):
449 delta = timedelta(days=int(timerange[:-1]))
450 elif timerange.endswith("w"):
451 delta = timedelta(weeks=int(timerange[:-1]))
452 else:
453 raise ValueError(
454 "Invalid timerange format. Expected a string like '24h', '1d', '1w', '1m', etc.",
455 )
456
457 start = datetime.utcnow() - delta
458 return start.isoformat() + "Z" # Elasticsearch expects the time in ISO format with a Z at the end
459
460 def __init__(self):
461 self.query = {
462 "query": {
463 "bool": {
464 "must": [],
465 },
466 },
467 "sort": [],
468 }
469
470 def add_time_range(self, timerange: str, timestamp_field: str):
471 """
472 Adds a time range filter to the query.
473
474 Args:
475 timerange (str): The time range to filter by.
476 timestamp_field (str): The name of the timestamp field in the query.
477
478 Returns:
479 self: The updated instance of the class.
480 """
481 start = self._get_time_range_start(timerange)
482 self.query["query"]["bool"]["must"].append(
483 {"range": {timestamp_field: {"gte": start, "lte": "now"}}},
484 )
485 return self
486
487 def add_matches(self, matches: Iterable[Tuple[str, str]]):
488 """
489 Adds matches to the query.
490
491 Args:
492 matches (Iterable[Tuple[str, str]]): A collection of field-value pairs to match.
493
494 Returns:
495 self: The current instance of the class.
496 """
497 for field, value in matches:
498 self.query["query"]["bool"]["must"].append({"match": {field: value}})
499 return self
500
501 def add_match_phrase(self, matches: Iterable[Tuple[str, str]]):
502 """
503 Adds match phrases to the query.
504
505 Args:
506 matches (Iterable[Tuple[str, str]]): A collection of field-value pairs to match.
507
508 Returns:
509 self: The current instance of the class.
510 """
511 for field, value in matches:
512 self.query["query"]["bool"]["must"].append({"match_phrase": {field: value}})
513 return self
514
515 def add_range(self, field: str, value: str):
516 """
517 Adds a range query to the Elasticsearch query.
518
519 Args:
520 field (str): The field to apply the range query on.
521 value (str): The value to compare against in the range query.
522
523 Returns:
524 self: The instance of the class with the range query added.
525 """
526 self.query["query"]["bool"]["must"].append({"range": {field: {"gte": value}}})
527 return self
528
529 def add_sort(self, field: str, order: str = "desc"):
530 """
531 Adds a sort field to the query.
532
533 Args:
534 field (str): The field to sort by.
535 order (str, optional): The sort order. Defaults to "desc".
536
537 Returns:
538 self: The updated instance of the class.
539 """
540 self.query["sort"].append({field: {"order": order}})
541 return self
542
543 def build(self):
544 return self.query
545
546
547 async def get_index_mappings_key_names(index_name: str):
548 """
549 Get the mappings of an index.
550
551 Args:
552 index_name (str): The Name of the index.
553
554 Returns:
555 list: The field names of the index.
556 """
557 es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
558 mappings = es_client.indices.get_mapping(index=index_name)
559 # return only the field names
560 return list(mappings[index_name]["mappings"]["properties"].keys())
561
562
563 async def return_graylog_events_index_names():
564 """
565 Return the index names of the Graylog events.
566
567 Returns:
568 list: The index names of the Graylog events.
569 """
570 es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
571 indices = es_client.indices.get_alias("gl-events*")
572 return list(indices.keys())
573
574
575 # async def get_index_source(index_name: str):
576 # """
577 # Get the 10 latest results from the index and search for where the source contains a field name of `syslog_type` or `integration`
578 # """
579 # es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
580 # query = {"size": 10, "query": {"bool": {"must": [{"exists": {"field": "syslog_type"}}]}}}
581 # response = es_client.search(index=index_name, body=query)
582 # for hit in response["hits"]["hits"]: # Loop through each hit in the response
583 # if "syslog_type" in hit["_source"]: # Check if 'syslog_type' exists in the source of the hit
584 # if hit["_source"]["syslog_type"] == "integration" and "integration" in hit["_source"]:
585 # return hit["_source"]["integration"] # Return the value of 'integration' if 'syslog_type' equals 'integration'
586 # return hit["_source"]["syslog_type"] # Return the value of 'syslog_type' for other cases
587 # raise HTTPException(status_code=404, detail=f"Source not found in index {index_name}")
588
589
590 async def get_index_source(index_name: str):
591 """
592 Get the 10 latest results from the index and search for where the source contains a field name of `syslog_type` or `integration`
593 """
594 es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
595
596 # First search for 'syslog_type'
597 query_syslog_type = {"size": 10, "query": {"bool": {"must": [{"exists": {"field": "syslog_type"}}]}}}
598 response = es_client.search(index=index_name, body=query_syslog_type)
599 for hit in response["hits"]["hits"]: # Loop through each hit in the response
600 if "syslog_type" in hit["_source"]: # Check if 'syslog_type' exists in the source of the hit
601 if hit["_source"]["syslog_type"] == "integration" and "integration" in hit["_source"]:
602 return hit["_source"]["integration"] # Return the value of 'integration' if 'syslog_type' equals 'integration'
603 return hit["_source"]["syslog_type"] # Return the value of 'syslog_type' for other cases
604
605 # If no 'syslog_type' found, search for 'integration'
606 query_integration = {"size": 10, "query": {"bool": {"must": [{"exists": {"field": "integration"}}]}}}
607 response = es_client.search(index=index_name, body=query_integration)
608 for hit in response["hits"]["hits"]: # Loop through each hit in the response
609 if "integration" in hit["_source"]: # Check if 'integration' exists in the source of the hit
610 return hit["_source"]["integration"] # Return the value of 'integration'
611
612 raise HTTPException(status_code=404, detail=f"Source not found in index {index_name}")
613
614
615 async def get_available_indices_via_source(source: str):
616 """
617 Get the available indices based on the source using regex matching.
618
619 Args:
620 source (str): The regex pattern for the source of the index.
621
622 Returns:
623 list: The available indices based on the source regex match.
624 """
625 es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
626 try:
627 indices = es_client.indices.get_alias("*")
628 logger.info(f"Indices: {indices.keys()}")
629 available_indices = []
630 source_pattern = re.compile(source) # Compile the regex pattern for the source
631 exclude_patterns = [re.compile("wazuh-monitoring"), re.compile("wazuh-statistics")]
632 for index in indices.keys():
633 if source_pattern.search(index) and not any(exclude.search(index) for exclude in exclude_patterns):
634 available_indices.append(index)
635 except Exception as e:
636 logger.error(f"An error occurred: {e}")
637 available_indices = []
638
639 logger.info(f"Available indices: {available_indices}")
640 return available_indices
641
642
643 async def resize_wazuh_index_fields():
644 """
645 Resize the Wazuh index fields to the correct size.
646 """
647 es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
648
649 settings = {"index.mapping.total_fields.limit": 2000}
650
651 try:
652 es_client.indices.put_settings(index="wazuh*", body=settings)
653 logger.info("Successfully resized the Wazuh index fields")
654 except Exception as e:
655 logger.error(f"An error occurred: {e}")
656 return None