Alert context fix (#254)
* feat: Add RawGenericSourceModel to alert schema This code change adds the `RawGenericSourceModel` class to the alert schema in the `alert.py` file. The `RawGenericSourceModel` includes fields such as `timestamp`, `timestamp_utc`, `rule_description`, and `syslog_level` that are used when creating the IRIS alert. This enhancement improves the alert creation process by providing additional information for the alert. * precommit fixes
taylor_socfortress committed
Jun 26, 2024 at 11:08 UTC
ba630083520db6a66e79436095bcdc6a15e326b1
3 files changed
+51
backend/app/integrations/alert_creation/general/schema/alert.py
+23
@@ -15,6 +15,28 @@ class ValidIocFields(Enum):
15
THREAT_INTEL_VALUE = "threat_intel_value"
16
17
18
+class RawGenericSourceModel(BaseModel):
19
+ timestamp: str = Field(..., description="The timestamp of the alert.")
20
+ timestamp_utc: Optional[str] = Field(
21
+ ...,
22
+ description="The UTC timestamp of the alert.",
23
+ )
24
+ rule_description: Optional[str] = Field(
25
+ "No autogenerated rule_description found",
26
+ description="The timefield of the alert to be used when creating the IRIS alert.",
27
+ )
28
+ syslog_level: Optional[str] = Field(
29
+ "No autogenerated syslog_level found",
30
+ description="The timefield of the alert to be used when creating the IRIS alert.",
31
+ )
32
+
33
+ class Config:
34
+ extra = Extra.allow
35
+
36
+ def to_dict(self):
37
+ return self.dict(exclude_none=True)
38
+
39
+
40
class CreateAlertRequest(BaseModel):
41
index: str = Field(
42
...,
@@ -56,6 +78,7 @@ class CreateAlertRequest(BaseModel):
78
None,
79
description="The IoC type of the alert which is needed for when we add the IoC to IRIS.",
80
)
81
+ _source: RawGenericSourceModel
82
83
class Config:
84
allow_population_by_field_name = True
backend/app/integrations/monitoring_alert/schema/monitoring_alert.py
+13
@@ -14,6 +14,16 @@ from app.integrations.alert_creation.general.schema.alert import IrisAsset
14
from app.integrations.alert_creation.general.schema.alert import IrisIoc
15
16
17
+class WazuhSourceFieldsToRemove(Enum):
18
+ GL2 = "gl2"
19
+ RULE_MITRE_TACTIC = "rule_mitre_tactic"
20
+ RULE_MITRE_ID = "rule_mitre_id"
21
+ RULE_MITRE_TECHNIQUE = "rule_mitre_technique"
22
+ RULE_ID = "rule_id"
23
+ MESSAGE = "message"
24
+ # Add more fields as needed
25
+
26
+
27
class MonitoringAlertsRequestModel(BaseModel):
28
id: Optional[int] = None
29
alert_id: str
@@ -336,6 +346,9 @@ class WazuhIrisAlertContext(BaseModel):
346
description="Name of the process",
347
)
348
349
+ class Config:
350
+ extra = Extra.allow
351
+
352
353
class WazuhIrisAlertPayload(BaseModel):
354
alert_title: str = Field(
backend/app/integrations/monitoring_alert/services/wazuh.py
+15
@@ -21,6 +21,7 @@ from app.integrations.alert_creation.general.schema.alert import ValidIocFields
21
from app.integrations.alert_creation.general.services.alert_multi_exclude import (
22
AlertDetailsService,
23
)
24
+from app.integrations.alert_escalation.schema.escalate_alert import GenericSourceModel
25
from app.integrations.alert_escalation.schema.general_alert import (
26
CreateAlertRequest as AddAlertRequest,
27
)
@@ -41,6 +42,9 @@ from app.integrations.monitoring_alert.schema.monitoring_alert import (
42
from app.integrations.monitoring_alert.schema.monitoring_alert import (
43
WazuhIrisAlertPayload,
44
)
45
+from app.integrations.monitoring_alert.schema.monitoring_alert import (
46
+ WazuhSourceFieldsToRemove,
47
+)
48
from app.integrations.monitoring_alert.utils.db_operations import remove_alert_id
49
from app.integrations.utils.alerts import get_asset_type_id
50
from app.integrations.utils.alerts import validate_ioc_type
@@ -297,6 +301,13 @@ async def build_alert_context_payload(
301
Returns:
302
WazuhIrisAlertContext: The built alert context payload.
303
"""
304
+ # Convert the _source to a dictionary
305
+ source_dict = alert_details._source.to_dict()
306
+
307
+ # Remove fields that start with any prefix in SourceFieldsToRemove
308
+ for field in WazuhSourceFieldsToRemove:
309
+ source_dict = {k: v for k, v in source_dict.items() if not k.startswith(field.value)}
310
+
311
return WazuhIrisAlertContext(
312
customer_iris_id=(
313
await get_customer_alert_settings(
@@ -331,6 +342,7 @@ async def build_alert_context_payload(
342
"No rule mitre technique found",
343
),
344
process_name=alert_details.process_name,
345
+ **source_dict,
346
)
347
348
@@ -423,6 +435,7 @@ async def create_alert_details(alert_details: WazuhAlertModel) -> CreateAlertReq
435
Returns:
436
CreateAlertRequest: The alert details object.
437
"""
438
+ logger.info(f"Creating alert details for alert: {alert_details}")
439
return CreateAlertRequest(
440
index=alert_details._index,
441
id=alert_details._id,
@@ -437,6 +450,7 @@ async def create_alert_details(alert_details: WazuhAlertModel) -> CreateAlertReq
450
timestamp_utc=alert_details._source.get("timestamp_utc", alert_details._source["timestamp"]),
451
process_id=alert_details._source.get("process_id", "No process ID found"),
452
process_name=await get_process_name(alert_details.to_dict()),
453
+ _source=GenericSourceModel(**alert_details._source),
454
)
455
456
@@ -456,6 +470,7 @@ async def create_and_update_alert_in_iris(
470
"""
471
logger.info("Alert does not exist in IRIS. Creating alert.")
472
alert_details = await create_alert_details(alert_details)
473
+ logger.info(f"Alert details: {alert_details}")
474
agent_details = await get_agent_by_hostname(alert_details.agent_name, session)
475
ioc_payload = await build_ioc_payload(alert_details)
476
iris_alert_payload = await build_alert_payload(