1
-import os
2
-from typing import List
3
-from typing import Optional
4
-
5
-from fastapi import HTTPException
6
-from loguru import logger
7
-from sqlalchemy.ext.asyncio import AsyncSession
8
-from sqlalchemy.future import select
9
-
10
-from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
11
-from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
12
-from app.connectors.utils import get_connector_info_from_db
13
-from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
14
-from app.db.universal_models import Agents
15
-from app.integrations.alert_creation.general.schema.alert import IrisAsset
16
-from app.integrations.alert_creation_settings.models.alert_creation_settings import (
17
- AlertCreationSettings,
18
-)
19
-from app.integrations.alert_escalation.schema.escalate_alert import CreateAlertRequest
20
-from app.integrations.alert_escalation.schema.escalate_alert import CreateAlertResponse
21
-from app.integrations.alert_escalation.schema.escalate_alert import CustomerCodeKeys
22
-from app.integrations.alert_escalation.schema.escalate_alert import GenericAlertModel
23
-from app.integrations.alert_escalation.schema.escalate_alert import GenericSourceModel
24
-from app.integrations.alert_escalation.schema.escalate_alert import IrisAlertContext
25
-from app.integrations.alert_escalation.schema.escalate_alert import IrisAlertPayload
26
-from app.integrations.alert_escalation.schema.escalate_alert import SourceFieldsToRemove
27
-from app.integrations.alert_escalation.schema.escalate_alert import SyslogLevelMapping
28
-from app.integrations.monitoring_alert.services.wazuh import (
29
- handle_customer_notifications,
30
-)
31
-from app.integrations.utils.alerts import get_asset_type_id
32
-
33
-
34
-async def fetch_settings(field: str, value: str, session: AsyncSession):
35
- """
36
- Fetch settings based on the field and value.
37
-
38
- Args:
39
- field (str): The field to check.
40
- value (str): The value to check.
41
- session (AsyncSession): The database session.
42
-
43
- Returns:
44
- AlertCreationSettings: The settings if found, None otherwise.
45
- """
46
- logger.info(f"Checking if {field}: {value} is valid.")
47
- result = await session.execute(
48
- select(AlertCreationSettings).where(
49
- getattr(AlertCreationSettings, field) == value,
50
- ),
51
- )
52
- settings = result.scalars().first()
53
- logger.info(f"Settings: {settings}")
54
- return settings
55
-
56
-
57
-async def is_customer_code_valid(customer_code: str, session: AsyncSession) -> AlertCreationSettings:
58
- """
59
- Check if the customer code is valid.
60
-
61
- Args:
62
- customer_code (str): The customer code to check.
63
- session (AsyncSession): The database session.
64
-
65
- Returns:
66
- bool: True if the customer code is valid, False otherwise.
67
- """
68
- settings = await fetch_settings("customer_code", customer_code, session)
69
-
70
- if settings:
71
- return settings
72
-
73
- # If no settings found with customer_code, try with office365_organization_id
74
- settings = await fetch_settings("office365_organization_id", customer_code, session)
75
-
76
- if settings:
77
- return settings
78
-
79
- raise HTTPException(
80
- status_code=400,
81
- detail=f"Customer code {customer_code} is not valid. Has the customer been provisioned?",
82
- )
83
-
84
-
85
-async def get_single_alert_details(
86
- alert_details: CreateAlertRequest,
87
-) -> GenericAlertModel:
88
- """
89
- Fetches the details of a single alert.
90
-
91
- Args:
92
- alert_details (CreateAlertRequest): The details of the alert to fetch.
93
-
94
- Returns:
95
- GenericAlertModel: The model representing the fetched alert.
96
-
97
- Raises:
98
- HTTPException: If there is an error while fetching the alert details.
99
- """
100
- logger.info(
101
- f"Fetching alert details for alert {alert_details.alert_id} in index {alert_details.index_name}",
102
- )
103
- es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
104
- try:
105
- alert = es_client.get(index=alert_details.index_name, id=alert_details.alert_id)
106
- source_model = GenericSourceModel(**alert["_source"])
107
- return GenericAlertModel(
108
- _source=source_model,
109
- _id=alert["_id"],
110
- _index=alert["_index"],
111
- _version=alert["_version"],
112
- rule_description=source_model.rule_description,
113
- syslog_level=source_model.syslog_level,
114
- )
115
- except Exception as e:
116
- logger.debug(f"Failed to collect alert details: {e}")
117
- raise HTTPException(
118
- status_code=400,
119
- detail=f"Failed to collect alert details: {e}",
120
- )
121
-
122
-
123
-async def set_alert_level(syslog_level: str):
124
- """
125
- Sets the alert level based on the syslog level.
126
-
127
- Args:
128
- syslog_level (str): The syslog level.
129
-
130
- Returns:
131
- int: The alert level.
132
- """
133
- for level in SyslogLevelMapping:
134
- if level.name == syslog_level:
135
- logger.info(f"Setting alert level to {level.value}")
136
- return level.value
137
- return 3
138
-
139
-
140
-def remove_process_name_if_osquery(source_dict: dict) -> None:
141
- """
142
- Remove the process_name field from the source dictionary if rule_group1 is 'osquery'.
143
-
144
- Args:
145
- source_dict (dict): The source dictionary.
146
- """
147
- rule_group1 = source_dict.get("rule_group1")
148
- if rule_group1 == "osquery":
149
- logger.info("Removing process_name field")
150
- source_dict.pop("process_name", None)
151
-
152
-
153
-def get_process_image(source_dict: dict) -> str:
154
- """
155
- Get the process_image field from the source dictionary.
156
-
157
- Args:
158
- source_dict (dict): The source dictionary.
159
-
160
- Returns:
161
- str: The process image.
162
- """
163
- process_image = source_dict.get("process_image")
164
- if not process_image:
165
- process_image = source_dict.get("data_win_eventdata_image")
166
- if not process_image:
167
- process_image = source_dict.get("data_event_Image")
168
- logger.info(f"Process image: {process_image}")
169
- return process_image
170
-
171
-
172
-def get_process_name_from_image(process_image: str) -> str:
173
- """
174
- Get the process name from the process image.
175
-
176
- Args:
177
- process_image (str): The process image.
178
-
179
- Returns:
180
- str: The process name.
181
- """
182
- process_name = os.path.basename(process_image) if process_image else None
183
- logger.info(f"Process name: {process_name}")
184
- return process_name
185
-
186
-
187
-async def get_process_name(source_dict: dict) -> List[str]:
188
- """
189
- Get the process name from the source dictionary.
190
-
191
- Args:
192
- source_dict (dict): The source dictionary.
193
-
194
- Returns:
195
- List[str]: The process name as a list.
196
- """
197
- remove_process_name_if_osquery(source_dict)
198
- process_image = get_process_image(source_dict)
199
- process_name = get_process_name_from_image(process_image)
200
- return [process_name] if process_name else []
201
-
202
-
203
-async def build_alert_context_payload(
204
- alert_details: GenericAlertModel,
205
- customer_alert_creation_settings: AlertCreationSettings,
206
-) -> IrisAlertContext:
207
- """
208
- Builds the payload for the alert context.
209
-
210
- Args:
211
- alert_details (GenericAlertModel): The details of the alert.
212
- agent_data (AgentsResponse): The data of the agent.
213
- session (AsyncSession): The async session.
214
-
215
- Returns:
216
- IrisAlertContext: The built alert context payload.
217
- """
218
- # Convert the _source to a dictionary
219
- source_dict = alert_details._source.to_dict()
220
-
221
- # Remove fields that start with any prefix in SourceFieldsToRemove
222
- for field in SourceFieldsToRemove:
223
- source_dict = {k: v for k, v in source_dict.items() if not k.startswith(field.value)}
224
-
225
- return IrisAlertContext(
226
- customer_iris_id=customer_alert_creation_settings.iris_customer_id,
227
- customer_name=customer_alert_creation_settings.customer_name,
228
- customer_cases_index=customer_alert_creation_settings.iris_index,
229
- alert_id=alert_details._id,
230
- alert_name=alert_details.rule_description,
231
- alert_level=await set_alert_level(alert_details.syslog_level),
232
- process_name=await get_process_name(source_dict),
233
- **source_dict,
234
- )
235
-
236
-
237
-async def construct_soc_alert_url(root_url: str, soc_alert_id: int) -> str:
238
- """Constructs the full URL for the SOC alert.
239
-
240
- Args:
241
- root_url (str): The root URL of the SOC alert system.
242
- soc_alert_id (int): The ID of the SOC alert.
243
-
244
- Returns:
245
- str: The full URL for the SOC alert.
246
-
247
- """
248
- url_path = f"/alerts?cid=1&page=1&per_page=10&sort=desc&alert_ids={soc_alert_id}"
249
- return f"{root_url}{url_path}"
250
-
251
-
252
-async def add_alert_to_document(
253
- es_client,
254
- alert: CreateAlertRequest,
255
- soc_alert_id: int,
256
- session: AsyncSession,
257
-) -> Optional[str]:
258
- """
259
- Update the alert document in Elasticsearch with the provided SOC alert ID URL.
260
-
261
- Parameters:
262
- - es_client: The Elasticsearch client instance to use for the update.
263
- - alert: The alert request object containing alert_id and index_name.
264
- - soc_alert_id: The alert ID as it exists within IRIS.
265
- - session: The database session for retrieving connector information.
266
-
267
- Returns:
268
- - True if the update is successful, False otherwise.
269
- """
270
- try:
271
- connector_info = await get_connector_info_from_db("DFIR-IRIS", session)
272
- full_url = await construct_soc_alert_url(
273
- connector_info["connector_url"],
274
- soc_alert_id,
275
- )
276
- es_client.update(
277
- index=alert.index_name,
278
- id=alert.alert_id,
279
- body={"doc": {"alert_url": full_url}},
280
- )
281
- logger.info(
282
- f"Added alert ID {soc_alert_id} to alert {alert.alert_id} in index {alert.index_name}",
283
- )
284
- return full_url
285
- except Exception as e:
286
- logger.error(
287
- f"Failed to add alert ID {soc_alert_id} to alert {alert.alert_id} in index {alert.index_name}: {e}",
288
- )
289
- # Attempt to remove read-only block
290
- try:
291
- es_client.indices.put_settings(
292
- index=alert.index_name,
293
- body={"index.blocks.write": None},
294
- )
295
- logger.info(
296
- f"Removed read-only block from index {alert.index_name}. Retrying update.",
297
- )
298
-
299
- # Retry the update operation
300
- es_client.update(
301
- index=alert.index_name,
302
- id=alert.alert_id,
303
- body={"doc": {"alert_url": full_url}},
304
- )
305
- logger.info(
306
- f"Added alert ID {soc_alert_id} to alert {alert.alert_id} in index {alert.index_name} after removing read-only block",
307
- )
308
-
309
- # Reenable the write block
310
- es_client.indices.put_settings(
311
- index=alert.index_name,
312
- body={"index.blocks.write": True},
313
- )
314
- return full_url
315
- except Exception as e2:
316
- logger.error(
317
- f"Failed to remove read-only block from index {alert.index_name}: {e2}",
318
- )
319
- return False
320
-
321
-
322
-async def get_customer_code(alert_details: dict):
323
- logger.info(f"Fetching customer code for alert {alert_details}")
324
-
325
- # Iterate over the possible keys and return the value if the key is present
326
- for key in CustomerCodeKeys:
327
- logger.info(f"Checking for key {key.value}")
328
- if key.value in alert_details:
329
- return alert_details[key.value]
330
-
331
- # If none of the keys are present, raise an exception
332
- logger.info(f"Failed to fetch customer code. Valid customer code field names are {', '.join([key.value for key in CustomerCodeKeys])}")
333
- raise HTTPException(
334
- status_code=400,
335
- detail=f"Failed to fetch customer code. Valid customer code field names are {', '.join([key.value for key in CustomerCodeKeys])}",
336
- )
337
-
338
-
339
-async def build_alert_payload(
340
- alert_details: GenericAlertModel,
341
- customer_alert_creation_settings: AlertCreationSettings,
342
-) -> IrisAlertPayload:
343
- """
344
- Builds the alert payload based on the provided alert details, agent data, IoC payload, and session.
345
-
346
- Args:
347
- alert_details (GenericAlertModel): The details of the alert.
348
- agent_data: The data of the agent.
349
- ioc_payload (Optional[IrisIoc]): The IoC payload.
350
- session (AsyncSession): The session object for database operations.
351
-
352
- Returns:
353
- IrisAlertPayload: The built alert payload.
354
-
355
- Raises:
356
- HTTPException: If there is an error while building the alert payload.
357
- """
358
- context_payload = await build_alert_context_payload(
359
- alert_details=alert_details,
360
- customer_alert_creation_settings=customer_alert_creation_settings,
361
- )
362
- logger.info(f"Context payload: {context_payload}")
363
- timefield = customer_alert_creation_settings.timefield
364
- # Get the timefield value from the alert_details
365
- if hasattr(alert_details, timefield):
366
- alert_details.time_field = getattr(alert_details, timefield)
367
- # Check if its part of _source
368
- if hasattr(alert_details._source, timefield):
369
- alert_details.time_field = getattr(alert_details._source, timefield)
370
- logger.info(f"Alert has context: {context_payload}")
371
- try:
372
- return IrisAlertPayload(
373
- alert_title=alert_details._source.rule_description,
374
- alert_description=alert_details._source.rule_description,
375
- alert_source="CoPilot - Manual Escalation",
376
- alert_status_id=3,
377
- alert_severity_id=5,
378
- alert_customer_id=customer_alert_creation_settings.iris_customer_id,
379
- alert_source_content=alert_details._source,
380
- alert_context=context_payload,
381
- alert_source_event_time=alert_details.time_field,
382
- )
383
- except Exception as e:
384
- logger.error(f"Failed to build alert payload: {e}")
385
- raise HTTPException(
386
- status_code=500,
387
- detail=f"Failed to build alert payload: {e}",
388
- )
389
-
390
-
391
-async def retrieve_agent_details_from_db(agent_name: str, session: AsyncSession):
392
- """
393
- Retrieve agent details from the database.
394
-
395
- Args:
396
- agent_name (str): The name of the agent.
397
- session (AsyncSession): The database session.
398
-
399
- Returns:
400
- Agents: The agent details.
401
- """
402
- logger.info(f"Retrieving agent details for {agent_name}")
403
- result = await session.execute(
404
- select(Agents).where(Agents.hostname == agent_name),
405
- )
406
- agent = result.scalars().first()
407
- if agent:
408
- return agent
409
- return None
410
-
411
-
412
-async def add_asset_to_iris_alert(alert_id: int, asset_details: Agents, iris_alert_payload: IrisAlertPayload, session: AsyncSession):
413
- client, alert_client = await initialize_client_and_alert("DFIR-IRIS")
414
- asset_data = {
415
- "assets": [
416
- dict(
417
- IrisAsset(
418
- asset_name=asset_details.hostname,
419
- asset_ip=asset_details.ip_address,
420
- asset_description="",
421
- asset_type_id=await get_asset_type_id(asset_details.os),
422
- asset_tags=f"agent_id:{asset_details.agent_id}",
423
- ).dict(),
424
- ),
425
- ],
426
- }
427
- logger.info(f"Adding asset to alert {alert_id} with asset data: {asset_data}")
428
- await fetch_and_validate_data(
429
- client,
430
- alert_client.update_alert,
431
- alert_id,
432
- asset_data,
433
- )
434
-
435
-
436
-async def add_asset_if_wazuh(alert_details: GenericAlertModel, alert_id: int, iris_alert_payload: IrisAlertPayload, session: AsyncSession):
437
- """
438
- Adds the asset to the alert if the alert is from Wazuh.
439
-
440
- Args:
441
- alert_details (GenericAlertModel): The details of the alert.
442
- alert_id (int): The ID of the alert.
443
- session (AsyncSession): The database session.
444
- """
445
- # Check if `agent_id` is present in the alert details and is not equal to `000`
446
- if hasattr(alert_details._source, "agent_id") and alert_details._source.agent_id != "000":
447
- logger.info(f"Adding asset to alert {alert_id}")
448
- asset_details = await retrieve_agent_details_from_db(alert_details._source.agent_name, session)
449
- if asset_details:
450
- await add_asset_to_iris_alert(alert_id, asset_details, iris_alert_payload, session)
451
- logger.info(f"Asset added to alert {alert_id}")
452
- return None
453
- else:
454
- logger.error(f"Failed to retrieve asset details for {alert_details._source.agent_name}")
455
- return None
456
- logger.info(f"Alert {alert_id} is not from Wazuh. Skipping asset addition.")
457
- return None
458
-
459
-
460
-async def create_alert(
461
- alert: CreateAlertRequest,
462
- session: AsyncSession,
463
-) -> CreateAlertResponse:
464
- """
465
- Creates an alert in IRIS.
466
-
467
- Args:
468
- alert (CreateAlertRequest): The request object containing the alert details.
469
- session (AsyncSession): The database session.
470
-
471
- Returns:
472
- CreateAlertResponse: The response object containing the created alert details.
473
-
474
- Raises:
475
- HTTPException: If there is an error creating the alert.
476
- """
477
- logger.info(f"Creating alert {alert.alert_id} in IRIS")
478
- alert_details = await get_single_alert_details(alert_details=alert)
479
- logger.info(f"Alert details: {alert_details}")
480
-
481
- customer_code = await get_customer_code(dict(alert_details._source))
482
- logger.info(f"Customer code: {customer_code}")
483
- customer_alert_creation_settings = await is_customer_code_valid(customer_code=customer_code, session=session)
484
- logger.info(f"Customer creation settings: {customer_alert_creation_settings}")
485
- iris_alert_payload = await build_alert_payload(
486
- alert_details=alert_details,
487
- customer_alert_creation_settings=customer_alert_creation_settings,
488
- )
489
- logger.info(f"Iris Alert Payload: {iris_alert_payload}")
490
- client, alert_client = await initialize_client_and_alert("DFIR-IRIS")
491
- result = await fetch_and_validate_data(
492
- client,
493
- alert_client.add_alert,
494
- iris_alert_payload.to_dict(),
495
- )
496
- alert_id = result["data"]["alert_id"]
497
- await handle_customer_notifications(
498
- customer_code=customer_code,
499
- alert_payload=iris_alert_payload,
500
- session=session,
501
- )
502
-
503
- await add_asset_if_wazuh(alert_details, alert_id, iris_alert_payload, session)
504
-
505
- es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
506
- iris_url = await add_alert_to_document(
507
- es_client,
508
- alert,
509
- result["data"]["alert_id"],
510
- session=session,
511
- )
512
- try:
513
- alert_id = result["data"]["alert_id"]
514
- return CreateAlertResponse(
515
- alert_id=alert_id,
516
- success=True,
517
- message=f"Alert {alert_id} created successfully",
518
- alert_url=iris_url,
519
- )
520
- except Exception as e:
521
- logger.error(f"Failed to create alert {alert.alert_id}: {e}")
522
- raise HTTPException(
523
- status_code=500,
524
- detail=f"Failed to create alert for ID {alert.alert_id}: {e}",
525
- )