138 lines
3.7 KiB
Python
138 lines
3.7 KiB
Python
"""
|
|
©AngelaMos | 2026
|
|
Alert.py
|
|
"""
|
|
|
|
from typing import Any
|
|
from datetime import datetime, UTC
|
|
from enum import StrEnum
|
|
|
|
from mongoengine import (
|
|
StringField,
|
|
DateTimeField,
|
|
IntField,
|
|
ListField,
|
|
ObjectIdField,
|
|
)
|
|
|
|
from app.config import settings
|
|
from app.core.streaming import publish_event
|
|
from app.models.Base import BaseDocument
|
|
from app.models.LogEvent import LogEvent, Severity
|
|
|
|
|
|
class AlertStatus(StrEnum):
|
|
"""
|
|
Alert lifecycle states
|
|
"""
|
|
NEW = "new"
|
|
ACKNOWLEDGED = "acknowledged"
|
|
INVESTIGATING = "investigating"
|
|
RESOLVED = "resolved"
|
|
FALSE_POSITIVE = "false_positive"
|
|
|
|
|
|
class Alert(BaseDocument):
|
|
"""
|
|
Alert generated by a correlation rule match
|
|
"""
|
|
meta: dict[str, Any] = { # noqa: RUF012
|
|
"collection": "alerts",
|
|
"ordering": ["-created_at"],
|
|
"indexes": ["status", "severity", "rule_id"],
|
|
}
|
|
|
|
rule_id = ObjectIdField(required=True)
|
|
rule_name = StringField(required=True)
|
|
severity = StringField(
|
|
required=True,
|
|
choices=[s.value for s in Severity],
|
|
)
|
|
title = StringField(required=True)
|
|
matched_event_ids = ListField(ObjectIdField())
|
|
matched_event_count = IntField(default=0)
|
|
group_value = StringField()
|
|
status = StringField(
|
|
required=True,
|
|
default=AlertStatus.NEW,
|
|
choices=[s.value for s in AlertStatus],
|
|
)
|
|
mitre_tactic = StringField()
|
|
mitre_technique = StringField()
|
|
acknowledged_by = StringField()
|
|
acknowledged_at = DateTimeField()
|
|
resolved_at = DateTimeField()
|
|
notes = StringField()
|
|
|
|
@classmethod
|
|
def create_from_rule(
|
|
cls,
|
|
rule: Any,
|
|
matched_event_ids: list[str],
|
|
group_value: str,
|
|
) -> Alert:
|
|
"""
|
|
Create an alert from a fired rule and publish to the alert stream
|
|
"""
|
|
alert = cls(
|
|
rule_id=rule.id,
|
|
rule_name=rule.name,
|
|
severity=rule.severity,
|
|
title=f"{rule.name} [{group_value}]",
|
|
matched_event_ids=matched_event_ids,
|
|
matched_event_count=len(matched_event_ids),
|
|
group_value=group_value,
|
|
mitre_tactic=rule.mitre_tactic,
|
|
mitre_technique=rule.mitre_technique,
|
|
)
|
|
alert.save() # type: ignore[no-untyped-call]
|
|
|
|
publish_event(
|
|
settings.ALERT_STREAM_KEY,
|
|
{
|
|
"id": str(alert.id),
|
|
"rule_name": alert.rule_name,
|
|
"severity": alert.severity,
|
|
"title": alert.title,
|
|
"group_value": alert.group_value,
|
|
"matched_event_count": alert.matched_event_count,
|
|
"status": alert.status,
|
|
},
|
|
)
|
|
|
|
return alert
|
|
|
|
def update_status(
|
|
self,
|
|
status: str,
|
|
username: str | None = None,
|
|
notes: str | None = None,
|
|
) -> None:
|
|
"""
|
|
Transition the alert to a new status with optional metadata
|
|
"""
|
|
self.status = status
|
|
if notes is not None:
|
|
self.notes = notes
|
|
|
|
now = datetime.now(UTC)
|
|
if status == AlertStatus.ACKNOWLEDGED and username:
|
|
self.acknowledged_by = username
|
|
self.acknowledged_at = now
|
|
elif status in (AlertStatus.RESOLVED, AlertStatus.FALSE_POSITIVE):
|
|
self.resolved_at = now
|
|
|
|
self.save() # type: ignore[no-untyped-call]
|
|
|
|
def get_with_events(self) -> dict[str, Any]:
|
|
"""
|
|
Return the alert with its matched log events loaded
|
|
"""
|
|
events = list(
|
|
LogEvent.objects(id__in=self.matched_event_ids) # type: ignore[no-untyped-call]
|
|
)
|
|
return {
|
|
"alert": self,
|
|
"matched_events": events,
|
|
}
|