| 1 | from typing import List |
| 2 | from typing import Tuple |
| 3 | |
| 4 | from loguru import logger |
| 5 | |
| 6 | from app.connectors.wazuh_indexer.utils.universal import ( |
| 7 | create_wazuh_indexer_client_async, |
| 8 | ) |
| 9 | from app.connectors.wazuh_indexer.utils.universal import ( |
| 10 | return_graylog_events_index_names, |
| 11 | ) |
| 12 | from app.incidents.schema.alert_collection import AlertPayloadItem |
| 13 | from app.incidents.schema.alert_collection import AlertsPayload |
| 14 | from app.incidents.schema.incident_alert import CreateAlertRequest |
| 15 | from app.incidents.schema.incident_alert import IndexNamesResponse |
| 16 | |
| 17 | |
| 18 | async def get_graylog_event_indices() -> IndexNamesResponse: |
| 19 | """ |
| 20 | Get the Graylog event indices. Get the Graylog event indices for the Graylog events. |
| 21 | |
| 22 | Returns: |
| 23 | List[str]: The list of Graylog event indices. |
| 24 | """ |
| 25 | # return await return_graylog_events_index_names() |
| 26 | return IndexNamesResponse(index_names=await return_graylog_events_index_names(), success=True, message="Success") |
| 27 | |
| 28 | |
| 29 | # async def construct_query(): |
| 30 | # """ |
| 31 | # Constructs the query to find alerts where `fields.COPILOT_ALERT_ID` is NONE. |
| 32 | # """ |
| 33 | # return {"query": {"bool": {"must": [{"term": {"fields.COPILOT_ALERT_ID": "NONE"}}]}}} |
| 34 | |
| 35 | |
| 36 | async def construct_query(): |
| 37 | """ |
| 38 | Constructs the query to find alerts where `fields.COPILOT_ALERT_ID` is NONE. |
| 39 | """ |
| 40 | return {"query": {"bool": {"must": [{"term": {"fields.COPILOT_ALERT_ID": "NONE"}}]}}, "sort": [{"timestamp": {"order": "asc"}}]} |
| 41 | |
| 42 | |
| 43 | async def fetch_alerts_for_index(es_client, index, query): |
| 44 | """ |
| 45 | Fetches alerts for a given index that match the query using the Elasticsearch scroll API. |
| 46 | """ |
| 47 | # Start the initial search request |
| 48 | response = await es_client.search( |
| 49 | index=index, |
| 50 | body=query, |
| 51 | scroll="2m", |
| 52 | size=1000, # Keep the search context open for 2 minutes # Number of results per "page" |
| 53 | ) |
| 54 | scroll_id = response["_scroll_id"] |
| 55 | hits = response["hits"]["hits"] |
| 56 | |
| 57 | # Keep fetching results while there are still results to fetch |
| 58 | while len(response["hits"]["hits"]): |
| 59 | response = await es_client.scroll(scroll_id=scroll_id, scroll="2m") # Extend the scroll context for another 2 minutes |
| 60 | # Update the scroll ID in case it changes |
| 61 | scroll_id = response["_scroll_id"] |
| 62 | hits.extend(response["hits"]["hits"]) |
| 63 | |
| 64 | # Close the scroll context |
| 65 | await es_client.clear_scroll(scroll_id=scroll_id) |
| 66 | |
| 67 | return [AlertPayloadItem(**hit) for hit in hits] |
| 68 | |
| 69 | |
| 70 | async def fetch_alerts_batch(es_client, index: str, query: dict, batch_size: int = 100) -> Tuple[List[AlertPayloadItem], int]: |
| 71 | """ |
| 72 | Fetches a single batch of alerts for a given index that match the query. |
| 73 | |
| 74 | Args: |
| 75 | es_client: Wazuh Indexer client |
| 76 | index: Index name to query |
| 77 | query: Query to execute |
| 78 | batch_size: Number of alerts to fetch (default 100) |
| 79 | |
| 80 | Returns: |
| 81 | Tuple of (list of alerts, total count of matching documents) |
| 82 | """ |
| 83 | try: |
| 84 | response = await es_client.search( |
| 85 | index=index, |
| 86 | body=query, |
| 87 | size=batch_size, |
| 88 | # sort=[{"@timestamp": {"order": "asc"}}], # Process oldest first |
| 89 | ) |
| 90 | |
| 91 | hits = response["hits"]["hits"] |
| 92 | total = response["hits"]["total"]["value"] if isinstance(response["hits"]["total"], dict) else response["hits"]["total"] |
| 93 | |
| 94 | logger.info(f"Fetched {len(hits)} alerts from index {index}. Total available: {total}") |
| 95 | |
| 96 | return [AlertPayloadItem(**hit) for hit in hits], total |
| 97 | except Exception as e: |
| 98 | logger.error(f"Error fetching alerts from index {index}: {e}") |
| 99 | return [], 0 |
| 100 | |
| 101 | |
| 102 | # async def get_alerts_not_created_in_copilot() -> AlertsPayload: |
| 103 | # """ |
| 104 | # Get the Graylog event indices. Then get all the results from the list of indices, where `copilot_alert_id` does not exist. |
| 105 | # """ |
| 106 | # indices = await return_graylog_events_index_names() |
| 107 | # logger.info(f"Indices: {indices}") |
| 108 | # es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer") |
| 109 | # query = await construct_query() |
| 110 | |
| 111 | # alerts_not_created = [] |
| 112 | # for index in indices: |
| 113 | # alerts = await fetch_alerts_for_index(es_client, index, query) |
| 114 | # alerts_not_created.extend(alerts) |
| 115 | |
| 116 | # logger.info(f"Alerts not created: {len(alerts_not_created)} alerts found") |
| 117 | # return AlertsPayload(alerts=alerts_not_created) |
| 118 | |
| 119 | |
| 120 | async def get_alerts_not_created_in_copilot(batch_size: int = 100) -> Tuple[AlertsPayload, int]: |
| 121 | """ |
| 122 | Get a batch of alerts that have not been created in CoPilot yet. |
| 123 | |
| 124 | Args: |
| 125 | batch_size: Maximum number of alerts to return (default 100) |
| 126 | |
| 127 | Returns: |
| 128 | Tuple of (AlertsPayload with alerts, total count remaining) |
| 129 | """ |
| 130 | indices = await return_graylog_events_index_names() |
| 131 | logger.info(f"Checking indices: {indices}") |
| 132 | |
| 133 | es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer") |
| 134 | query = await construct_query() |
| 135 | |
| 136 | alerts_to_process = [] |
| 137 | total_remaining = 0 |
| 138 | |
| 139 | # Fetch from each index until we have enough alerts or run out |
| 140 | for index in indices: |
| 141 | if len(alerts_to_process) >= batch_size: |
| 142 | break |
| 143 | |
| 144 | remaining_to_fetch = batch_size - len(alerts_to_process) |
| 145 | alerts, index_total = await fetch_alerts_batch(es_client, index, query, remaining_to_fetch) |
| 146 | |
| 147 | alerts_to_process.extend(alerts) |
| 148 | total_remaining += index_total |
| 149 | |
| 150 | logger.info(f"Returning {len(alerts_to_process)} alerts. Total remaining across all indices: {total_remaining}") |
| 151 | |
| 152 | return AlertsPayload(alerts=alerts_to_process), total_remaining |
| 153 | |
| 154 | |
| 155 | async def get_original_alert_id(origin_context: str): |
| 156 | """ |
| 157 | Get the original alert id from the origin context. |
| 158 | """ |
| 159 | # Assuming the ID is the last part after the last colon and before the last underscore |
| 160 | try: |
| 161 | return origin_context.split(":")[-1].split("_")[-1] |
| 162 | except IndexError: # In case the origin_context does not follow the expected pattern |
| 163 | return None |
| 164 | |
| 165 | |
| 166 | async def get_original_alert_index_name(origin_context: str): |
| 167 | """ |
| 168 | Get the original alert index name from the origin context. |
| 169 | """ |
| 170 | # Assuming the index name is the part after 'es:' and before the next colon |
| 171 | try: |
| 172 | return origin_context.split("es:")[-1].split(":")[0] |
| 173 | except IndexError: # In case the origin_context does not follow the expected pattern |
| 174 | return None |
| 175 | |
| 176 | |
| 177 | async def add_copilot_alert_id(index_data: CreateAlertRequest, alert_id: int): |
| 178 | """ |
| 179 | Add the CoPilot alert ID to the Graylog event. |
| 180 | """ |
| 181 | es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer") |
| 182 | body = {"doc": {"fields": {"COPILOT_ALERT_ID": f"{alert_id}"}}} |
| 183 | try: |
| 184 | await es_client.update(index=index_data.index_name, id=index_data.alert_id, body=body) |
| 185 | logger.info(f"Added CoPilot alert ID {alert_id} to Graylog event {index_data.alert_id} in index {index_data.index_name}") |
| 186 | except Exception as e: |
| 187 | logger.error( |
| 188 | f"Failed to add CoPilot alert ID {alert_id} to Graylog event {index_data.alert_id} in index {index_data.index_name}: {e}", |
| 189 | ) |
| 190 | |
| 191 | # Attempt to remove read-only block |
| 192 | try: |
| 193 | await es_client.indices.put_settings(index=index_data.index_name, body={"index.blocks.write": None}) |
| 194 | logger.info(f"Removed read-only block from index {index_data.index_name}. Retrying update.") |
| 195 | |
| 196 | # Retry the update operation |
| 197 | await es_client.update(index=index_data.index_name, id=index_data.alert_id, body=body) |
| 198 | logger.info( |
| 199 | f"Added CoPilot alert ID {alert_id} to Graylog event {index_data.alert_id} in index {index_data.index_name} after removing read-only block", |
| 200 | ) |
| 201 | |
| 202 | # Re-enable the write block |
| 203 | await es_client.indices.put_settings(index=index_data.index_name, body={"index.blocks.write": True}) |
| 204 | except Exception as e2: |
| 205 | logger.error(f"Failed to remove read-only block from index {index_data.index_name}: {e2}") |
| 206 | |
| 207 | return None |