add get_customer_code to get soc alert by id
Taylor committed
Jan 29, 2024 at 15:19 UTC
ad3a89f5089f0a19917ae981a77c4e9a4c5051c7
2 files changed
+37
-3
backend/app/connectors/dfir_iris/routes/alerts.py
+4
-2
@@ -3,8 +3,10 @@ from fastapi import Depends
3
from fastapi import HTTPException
4
from fastapi import Security
5
from loguru import logger
6
+from sqlalchemy.ext.asyncio import AsyncSession
7
8
from app.auth.utils import AuthHandler
9
+from app.db.db_session import get_db
10
from app.connectors.dfir_iris.schema.alerts import AlertResponse
11
from app.connectors.dfir_iris.schema.alerts import AlertsResponse
12
from app.connectors.dfir_iris.schema.alerts import BookmarkedAlertsResponse
@@ -86,7 +88,7 @@ async def get_alerts_filtered(request: FilterAlertsRequest) -> AlertsResponse:
88
description="Get an alert by ID",
89
dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
90
)
89
-async def get_alert_by_id(alert_id: str = Depends(verify_alert_exists)) -> AlertResponse:
91
+async def get_alert_by_id(alert_id: str = Depends(verify_alert_exists), session: AsyncSession = Depends(get_db)) -> AlertResponse:
92
"""
93
Retrieve an alert by its ID.
94
@@ -97,7 +99,7 @@ async def get_alert_by_id(alert_id: str = Depends(verify_alert_exists)) -> Alert
99
AlertResponse: The response containing the alert information.
100
"""
101
logger.info(f"Fetching alert {alert_id}")
100
- return await get_alert(alert_id=alert_id)
102
+ return await get_alert(alert_id=alert_id, session=session)
103
104
105
@dfir_iris_alerts_router.get(
backend/app/connectors/dfir_iris/services/alerts.py
+33
-1
@@ -2,6 +2,11 @@ from fastapi import HTTPException
2
from loguru import logger
3
4
from app.connectors.dfir_iris.schema.alerts import AlertResponse
5
+from sqlalchemy.ext.asyncio import AsyncSession
6
+from sqlalchemy.future import select
7
+from app.integrations.alert_creation_settings.models.alert_creation_settings import (
8
+ AlertCreationSettings,
9
+)
10
from app.connectors.dfir_iris.schema.alerts import AlertsResponse
11
from app.connectors.dfir_iris.schema.alerts import BookmarkedAlertsResponse
12
from app.connectors.dfir_iris.schema.alerts import CaseCreationResponse
@@ -11,6 +16,30 @@ from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
16
from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
17
18
19
+async def get_customer_code(session: AsyncSession, customer_id: int) -> str:
20
+ """
21
+ Retrieves the customer code for a given customer ID.
22
+
23
+ Args:
24
+ session (AsyncSession): The database session.
25
+ customer_id (int): The ID of the customer.
26
+
27
+ Returns:
28
+ The customer code for the given customer ID.
29
+ """
30
+ logger.info(f"Retrieving customer code for customer ID {customer_id}")
31
+ try:
32
+ alert_creation_settings = await session.execute(
33
+ select(AlertCreationSettings).filter(AlertCreationSettings.iris_customer_id == customer_id),
34
+ )
35
+ alert_creation_settings = alert_creation_settings.scalars().first()
36
+ if alert_creation_settings is None:
37
+ return "Customer Not Found"
38
+ return alert_creation_settings.customer_code
39
+ except Exception as e:
40
+ logger.error(f"Error retrieving customer code for customer ID {customer_id}: {e}")
41
+ return "Customer Not Found"
42
+
43
async def get_alerts(request: FilterAlertsRequest) -> AlertsResponse:
44
"""
45
Retrieves alerts from the DFIR-IRIS service.
@@ -54,7 +83,7 @@ def construct_params(request: FilterAlertsRequest) -> dict:
83
return {k: v for k, v in params.items() if v is not None}
84
85
57
-async def get_alert(alert_id: str) -> AlertResponse:
86
+async def get_alert(alert_id: str, session: AsyncSession) -> AlertResponse:
87
"""
88
Retrieves an alert by its ID.
89
@@ -69,6 +98,9 @@ async def get_alert(alert_id: str) -> AlertResponse:
98
"""
99
client, alert = await initialize_client_and_alert("DFIR-IRIS")
100
result = await fetch_and_validate_data(client, alert.get_alert, alert_id)
101
+ # Add the customer code to the alert
102
+ customer_code = await get_customer_code(session, result["data"]["customer"]["customer_id"])
103
+ result["data"]["customer"]["customer_code"] = customer_code
104
return AlertResponse(success=True, message="Successfully fetched alert", alert=result["data"])
105
106