@cryptotaxi247 / CoPilot / commits / d9203195

Enable timefield (#145)

* fallback to `timefield` if `timestamp_utc` cannot be mapped * precommit fixes

taylor_socfortress committed Feb 11, 2024 at 09:32 UTC d9203195a7b24e3dbe0dabbf718e5009dfea7fa2
2 files changed +43 -26
backend/app/connectors/wazuh_indexer/services/alerts.py
+32 -23
@@ -3,6 +3,7 @@ from typing import List
3 from typing import Optional
4 from typing import Type
5
6 +from elasticsearch7.exceptions import RequestError
7 from fastapi import HTTPException
8 from loguru import logger
9
@@ -92,34 +93,42 @@ async def collect_alerts_generic(
93 """
94 es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
95 query_builder = AlertsQueryBuilder()
95 - query_builder.add_time_range(
96 - timerange=body.timerange,
97 - timestamp_field=body.timestamp_field,
98 - )
99 - query_builder.add_matches(matches=[(body.alert_field, body.alert_value)])
100 - query_builder.add_sort(body.timestamp_field)
96
102 - if is_host_specific:
103 - query_builder.add_match_phrase(matches=[("agent_name", body.agent_name)])
97 + try:
98 + query_builder.add_time_range(
99 + timerange=body.timerange,
100 + timestamp_field=body.timestamp_field,
101 + )
102 + query_builder.add_matches(matches=[(body.alert_field, body.alert_value)])
103 + query_builder.add_sort(body.timestamp_field)
104
105 - query = query_builder.build()
105 + if is_host_specific:
106 + query_builder.add_match_phrase(matches=[("agent_name", body.agent_name)])
107 +
108 + query = query_builder.build()
109
107 - try:
110 alerts = es_client.search(index=index_name, body=query, size=body.size)
109 - logger.info(f"Alerts collected: {alerts}")
110 - alerts_list = [alert for alert in alerts["hits"]["hits"]]
111 - logger.info(f"Alerts collected: {alerts_list}")
112 - return CollectAlertsResponse(
113 - alerts=alerts_list,
114 - success=True,
115 - message="Alerts collected successfully",
116 - )
117 - except Exception as e:
111 + except RequestError as e:
112 logger.warning(f"An error occurred while collecting alerts: {e}")
119 - raise HTTPException(
120 - status_code=500,
121 - detail=f"An error occurred while collecting alerts: {e}",
122 - )
113 + if "No mapping found for [timestamp_utc] in order to sort on" in str(e):
114 + logger.warning("Retrying with timestamp field set to 'timestamp'")
115 + body.timestamp_field = "timestamp"
116 + return await collect_alerts_generic(index_name, body, is_host_specific)
117 + else:
118 + logger.warning(f"An error occurred while collecting alerts: {e}")
119 + raise HTTPException(
120 + status_code=500,
121 + detail=f"An error occurred while collecting alerts: {e}",
122 + )
123 +
124 + logger.info(f"Alerts collected: {alerts}")
125 + alerts_list = [alert for alert in alerts["hits"]["hits"]]
126 + logger.info(f"Alerts collected: {alerts_list}")
127 + return CollectAlertsResponse(
128 + alerts=alerts_list,
129 + success=True,
130 + message="Alerts collected successfully",
131 + )
132
133
134 async def get_alerts_generic(
backend/app/connectors/wazuh_indexer/utils/universal.py
+11 -3
@@ -254,9 +254,17 @@ class AlertsQueryBuilder:
254 self: The updated instance of the class.
255 """
256 start = self._get_time_range_start(timerange)
257 - self.query["query"]["bool"]["must"].append(
258 - {"range": {timestamp_field: {"gte": start, "lte": "now"}}},
259 - )
257 + range_query = {
258 + "range": {
259 + timestamp_field: {
260 + "gte": start,
261 + "lte": "now",
262 + },
263 + },
264 + }
265 + if timestamp_field == "timestamp":
266 + range_query["range"][timestamp_field]["format"] = "strict_date_optional_time"
267 + self.query["query"]["bool"]["must"].append(range_query)
268 return self
269
270 def add_matches(self, matches: Iterable[Tuple[str, str]]):