From 5214aa4806279907041794c41690a9645930b568 Mon Sep 17 00:00:00 2001 From: Tal Borenstein Date: Sun, 1 Sep 2024 17:13:30 +0300 Subject: [PATCH] feat: basic working --- keep/api/bl/blackouts.py | 20 ------ keep/api/bl/blackouts_bl.py | 54 ++++++++++++++++ .../bl/{enrichments.py => enrichments_bl.py} | 0 keep/api/routes/alerts.py | 2 +- keep/api/tasks/process_event_task.py | 24 ++++++- keep/functions/__init__.py | 2 +- keep/providers/base/base_provider.py | 2 +- tests/test_blackouts_bl.py | 62 +++++++++++++++++++ tests/test_enrichments.py | 2 +- tests/test_functions.py | 2 +- tests/test_search_alerts.py | 2 +- 11 files changed, 144 insertions(+), 28 deletions(-) delete mode 100644 keep/api/bl/blackouts.py create mode 100644 keep/api/bl/blackouts_bl.py rename keep/api/bl/{enrichments.py => enrichments_bl.py} (100%) create mode 100644 tests/test_blackouts_bl.py diff --git a/keep/api/bl/blackouts.py b/keep/api/bl/blackouts.py deleted file mode 100644 index f864fae69..000000000 --- a/keep/api/bl/blackouts.py +++ /dev/null @@ -1,20 +0,0 @@ -from sqlmodel import Session - -from keep.api.core.db import get_session_sync -from keep.api.models.alert import AlertDto -from keep.api.models.db.blackout import BlackoutRule - - -class BlackoutsBl: - def __init__(self, tenant_id: str, session: Session | None) -> None: - self.tenant_id = tenant_id - session = session if session else get_session_sync() - self.blackouts = ( - session.query(BlackoutRule) - .filter(BlackoutRule.tenant_id == tenant_id) - .filter(BlackoutRule.enabled == True) - .all() - ) - - def check_if_alert_in_blackout(alert: AlertDto) -> bool: - return False diff --git a/keep/api/bl/blackouts_bl.py b/keep/api/bl/blackouts_bl.py new file mode 100644 index 000000000..f15c70a53 --- /dev/null +++ b/keep/api/bl/blackouts_bl.py @@ -0,0 +1,54 @@ +import json +import logging + +import celpy +from sqlmodel import Session + +from keep.api.core.db import get_session_sync +from keep.api.models.alert import AlertDto +from keep.api.models.db.blackout import BlackoutRule +from keep.api.utils.cel_utils import preprocess_cel_expression + + +class BlackoutsBl: + def __init__(self, tenant_id: str, session: Session | None) -> None: + self.logger = logging.getLogger(__name__) + self.tenant_id = tenant_id + session = session if session else get_session_sync() + self.blackouts: list[BlackoutRule] = ( + session.query(BlackoutRule) + .filter(BlackoutRule.tenant_id == tenant_id) + .filter(BlackoutRule.enabled == True) + .all() + ) + + def check_if_alert_in_blackout(self, alert: AlertDto) -> bool: + extra = {"tenant_id": self.tenant_id, "fingerprint": alert.fingerprint} + self.logger.info("Checking blackout for alert", extra=extra) + env = celpy.Environment() + + for blackout in self.blackouts: + cel = preprocess_cel_expression(blackout.cel_query) + ast = env.compile(cel) + prgm = env.program(ast) + + payload = alert.dict() + # todo: fix this in the future + payload["source"] = payload["source"][0] + + activation = celpy.json_to_cel(json.loads(json.dumps(payload, default=str))) + + try: + cel_result = prgm.evaluate(activation) + except celpy.evaluation.CELEvalError as e: + if "no such member" in str(e): + continue + # wtf + raise + if cel_result: + self.logger.info( + "Alert is blacked out", extra={**extra, "blackout_id": blackout.id} + ) + return True + self.logger.info("Alert is not blacked out", extra=extra) + return False diff --git a/keep/api/bl/enrichments.py b/keep/api/bl/enrichments_bl.py similarity index 100% rename from keep/api/bl/enrichments.py rename to keep/api/bl/enrichments_bl.py diff --git a/keep/api/routes/alerts.py b/keep/api/routes/alerts.py index 7b0938815..4980d4c65 100644 --- a/keep/api/routes/alerts.py +++ b/keep/api/routes/alerts.py @@ -21,7 +21,7 @@ from pusher import Pusher from keep.api.arq_worker import get_pool -from keep.api.bl.enrichments import EnrichmentsBl +from keep.api.bl.enrichments_bl import EnrichmentsBl from keep.api.core.config import config from keep.api.core.db import get_alert_audit as get_alert_audit_db from keep.api.core.db import get_alerts_by_fingerprint, get_enrichment, get_last_alerts diff --git a/keep/api/tasks/process_event_task.py b/keep/api/tasks/process_event_task.py index 22aa18e0a..1e36ccb41 100644 --- a/keep/api/tasks/process_event_task.py +++ b/keep/api/tasks/process_event_task.py @@ -14,7 +14,8 @@ # internals from keep.api.alert_deduplicator.alert_deduplicator import AlertDeduplicator -from keep.api.bl.enrichments import EnrichmentsBl +from keep.api.bl.blackouts_bl import BlackoutsBl +from keep.api.bl.enrichments_bl import EnrichmentsBl from keep.api.core.db import ( get_alerts_by_fingerprint, get_all_presets, @@ -238,7 +239,26 @@ def __handle_formatted_events( }, ) - # first, filter out any deduplicated events + # first, check for blackouts + blackouts_bl = BlackoutsBl(tenant_id=tenant_id, session=session) + if blackouts_bl.blackouts: + formatted_events = [ + event + for event in formatted_events + if blackouts_bl.check_if_alert_in_blackout(event) is False + ] + else: + logger.debug( + "No blackouts configured for this tenant", extra={"tenant_id": tenant_id} + ) + + if not formatted_events: + logger.info( + "No alerts to process after running blackouts", + extra={"tenant_id": tenant_id}, + ) + + # second, filter out any deduplicated events alert_deduplicator = AlertDeduplicator(tenant_id) for event in formatted_events: diff --git a/keep/functions/__init__.py b/keep/functions/__init__.py index 99a77a5cf..5645ef718 100644 --- a/keep/functions/__init__.py +++ b/keep/functions/__init__.py @@ -10,7 +10,7 @@ from dateutil import parser from dateutil.parser import ParserError -from keep.api.bl.enrichments import EnrichmentsBl +from keep.api.bl.enrichments_bl import EnrichmentsBl from keep.api.core.db import get_alerts_by_fingerprint from keep.api.models.alert import AlertStatus from keep.api.utils.enrichment_helpers import convert_db_alerts_to_dto_alerts diff --git a/keep/providers/base/base_provider.py b/keep/providers/base/base_provider.py index 99576f732..465aab9b8 100644 --- a/keep/providers/base/base_provider.py +++ b/keep/providers/base/base_provider.py @@ -18,7 +18,7 @@ import opentelemetry.trace as trace import requests -from keep.api.bl.enrichments import EnrichmentsBl +from keep.api.bl.enrichments_bl import EnrichmentsBl from keep.api.core.db import get_enrichments from keep.api.models.alert import AlertDto, AlertSeverity, AlertStatus from keep.api.models.db.alert import AlertActionType diff --git a/tests/test_blackouts_bl.py b/tests/test_blackouts_bl.py new file mode 100644 index 000000000..1d3781e49 --- /dev/null +++ b/tests/test_blackouts_bl.py @@ -0,0 +1,62 @@ +from datetime import datetime, timedelta +from unittest.mock import MagicMock + +import pytest + +from keep.api.bl.blackouts_bl import BlackoutsBl +from keep.api.models.alert import AlertDto +from keep.api.models.db.blackout import BlackoutRule + + +@pytest.fixture +def mock_session(): + return MagicMock() + + +@pytest.fixture +def mock_blackout_rule(): + return BlackoutRule( + id=1, + name="Test Blackout", + tenant_id="test-tenant", + cel_query='source == "test-source"', + start_time=datetime.utcnow() - timedelta(hours=1), + end_time=datetime.utcnow() + timedelta(hours=1), + enabled=True, + ) + + +@pytest.fixture +def alert_dto(): + return AlertDto( + id="test-alert", + source=["test-source"], + name="Test Alert", + status="firing", + severity="critical", + lastReceived="2021-08-01T00:00:00Z", + ) + + +def test_alert_in_blackout(mock_session, mock_blackout_rule, alert_dto): + mock_session.query.return_value.filter.return_value.filter.return_value.all.return_value = [ + mock_blackout_rule + ] + + blackout_bl = BlackoutsBl(tenant_id="test-tenant", session=mock_session) + result = blackout_bl.check_if_alert_in_blackout(alert_dto) + + assert result is True + + +def test_alert_not_in_blackout(mock_session, mock_blackout_rule, alert_dto): + # Modify the cel_query so that the alert won't match + mock_blackout_rule.cel_query = 'source == "other-source"' + mock_session.query.return_value.filter.return_value.filter.return_value.all.return_value = [ + mock_blackout_rule + ] + + blackout_bl = BlackoutsBl(tenant_id="test-tenant", session=mock_session) + result = blackout_bl.check_if_alert_in_blackout(alert_dto) + + assert result is False diff --git a/tests/test_enrichments.py b/tests/test_enrichments.py index 790072c74..8ffe79850 100644 --- a/tests/test_enrichments.py +++ b/tests/test_enrichments.py @@ -3,7 +3,7 @@ import pytest -from keep.api.bl.enrichments import EnrichmentsBl +from keep.api.bl.enrichments_bl import EnrichmentsBl from keep.api.core.dependencies import SINGLE_TENANT_UUID from keep.api.models.alert import AlertDto from keep.api.models.db.extraction import ExtractionRule diff --git a/tests/test_functions.py b/tests/test_functions.py index 47a850915..e4f4a385c 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -6,7 +6,7 @@ import pytz import keep.functions as functions -from keep.api.bl.enrichments import EnrichmentsBl +from keep.api.bl.enrichments_bl import EnrichmentsBl from keep.api.core.dependencies import SINGLE_TENANT_UUID from keep.api.models.alert import AlertStatus from keep.api.models.db.alert import AlertActionType diff --git a/tests/test_search_alerts.py b/tests/test_search_alerts.py index 9957fe04e..bb45436ff 100644 --- a/tests/test_search_alerts.py +++ b/tests/test_search_alerts.py @@ -3,7 +3,7 @@ import pytest -from keep.api.bl.enrichments import EnrichmentsBl +from keep.api.bl.enrichments_bl import EnrichmentsBl from keep.api.core.dependencies import SINGLE_TENANT_UUID from keep.api.models.alert import AlertDto from keep.api.models.db.alert import AlertActionType