@cryptotaxi247 / CoPilot / commits / ba630083

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(