Add customer code to case response
Taylor committed
Jan 29, 2024 at 12:15 UTC
d5443d21358b258efd002dd02225a9052f13cdae
3 files changed
+44
-4
backend/app/connectors/dfir_iris/routes/cases.py
+4
-2
@@ -5,6 +5,7 @@ from fastapi import Depends
5
from fastapi import HTTPException
6
from fastapi import Security
7
from loguru import logger
8
+from sqlalchemy.ext.asyncio import AsyncSession
9
10
from app.auth.utils import AuthHandler
11
from app.connectors.dfir_iris.schema.cases import CaseOlderThanBody
@@ -24,6 +25,7 @@ from app.connectors.dfir_iris.services.cases import get_single_case
25
from app.connectors.dfir_iris.services.cases import purge_cases
26
from app.connectors.dfir_iris.services.cases import reopen_case
27
from app.connectors.dfir_iris.utils.universal import check_case_exists
28
+from app.db.db_session import get_db
29
30
31
async def verify_case_exists(case_id: int) -> int:
@@ -75,7 +77,7 @@ def get_timedelta(older_than: int, time_unit: TimeUnit) -> CaseOlderThanBody:
77
description="Get all cases",
78
dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
79
)
78
-async def get_cases_route() -> CaseResponse:
80
+async def get_cases_route(session: AsyncSession = Depends(get_db)) -> CaseResponse:
81
"""
82
Get all cases.
83
@@ -83,7 +85,7 @@ async def get_cases_route() -> CaseResponse:
85
CaseResponse: The response containing all cases.
86
"""
87
logger.info("Fetching all cases")
86
- return await get_all_cases()
88
+ return await get_all_cases(session=session)
89
90
91
@dfir_iris_cases_router.post(
backend/app/connectors/dfir_iris/schema/cases.py
+1
@@ -28,6 +28,7 @@ class CaseModel(BaseModel):
28
owner_id: int
29
state_id: int
30
state_name: str
31
+ customer_code: str
32
33
34
class CaseResponse(BaseModel):
backend/app/connectors/dfir_iris/services/cases.py
+39
-2
@@ -4,6 +4,7 @@ from typing import List
4
5
from dfir_iris_client.case import Case
6
from fastapi import HTTPException
7
+from sqlalchemy.ext.asyncio import AsyncSession
8
from loguru import logger
9
10
from app.connectors.dfir_iris.schema.cases import CaseOlderThanBody
@@ -16,7 +17,10 @@ from app.connectors.dfir_iris.schema.cases import SingleCaseBody
17
from app.connectors.dfir_iris.schema.cases import SingleCaseResponse
18
from app.connectors.dfir_iris.utils.universal import create_dfir_iris_client
19
from app.connectors.dfir_iris.utils.universal import fetch_and_parse_data
19
-
20
+from app.integrations.alert_creation_settings.models.alert_creation_settings import (
21
+ AlertCreationSettings,
22
+)
23
+from sqlalchemy.future import select
24
25
async def get_client_and_cases() -> Dict:
26
"""
@@ -31,6 +35,36 @@ async def get_client_and_cases() -> Dict:
35
result = await fetch_and_parse_data(dfir_iris_client, case.list_cases)
36
return result
37
38
+async def get_customer_code(session: AsyncSession, client_name: str) -> str:
39
+ """
40
+ Retrieves the customer code for a given customer ID.
41
+
42
+ Args:
43
+ session (AsyncSession): The database session.
44
+ customer_id (int): The ID of the customer.
45
+
46
+ Returns:
47
+ The customer code for the given customer ID.
48
+ """
49
+ try:
50
+ alert_creation_settings = await session.execute(
51
+ select(AlertCreationSettings).filter(
52
+ AlertCreationSettings.iris_customer_name == client_name
53
+ )
54
+ )
55
+ alert_creation_settings = alert_creation_settings.scalars().first()
56
+ if alert_creation_settings is None:
57
+ return "Customer Not Found"
58
+ logger.info(
59
+ f"Alert creation settings for customer ID {client_name}: {alert_creation_settings}"
60
+ )
61
+ logger.info(
62
+ f"Customer code for customer ID {client_name}: {alert_creation_settings.customer_code}"
63
+ )
64
+ return alert_creation_settings.customer_code
65
+ except Exception as e:
66
+ logger.error(f"Error retrieving customer code for customer ID {client_name}: {e}")
67
+ return "Customer Not Found"
68
69
def filter_open_cases(cases: List[Dict]) -> List[Dict]:
70
"""
@@ -70,7 +104,7 @@ def filter_cases_older_than(cases: List[Dict], older_than: datetime) -> List[Dic
104
return filtered_cases
105
106
73
-async def get_all_cases() -> CaseResponse:
107
+async def get_all_cases(session: AsyncSession) -> CaseResponse:
108
"""
109
Retrieves all cases from DFIR-IRIS.
110
@@ -85,6 +119,9 @@ async def get_all_cases() -> CaseResponse:
119
if not result["success"]:
120
logger.error(f"Failed to get all cases: {result['message']}")
121
raise HTTPException(status_code=500, detail=f"Failed to get all cases: {result['message']}")
122
+ # For the `customer_id` get the customer code from the database and append it to the case
123
+ for case in result["data"]:
124
+ case["customer_code"] = await get_customer_code(session, case["client_name"])
125
return CaseResponse(success=True, message="Successfully fetched all cases", cases=result["data"])
126
except Exception as err:
127
logger.error(f"Failed to get all cases: {err}")