Refactor collect_agent_soc_cases function to filter cases by agent ID
Taylor committed
Jan 11, 2024 at 10:57 UTC
e3ea78ed1a78bb9e6fe5b99d06155824991d5a3e
2 files changed
+81
-12
backend/app/agents/dfir_iris/schema/cases.py
new
+7
@@ -0,0 +1,7 @@
1
+from pydantic import BaseModel
2
+from typing import List
3
+
4
+class AssetCaseIDResponse(BaseModel):
5
+ message: str
6
+ success: bool
7
+ case_ids: List[int]
backend/app/agents/dfir_iris/services/cases.py
+74
-12
@@ -1,15 +1,43 @@
1
from fastapi import HTTPException
2
from loguru import logger
3
+from typing import List
4
5
6
from app.connectors.dfir_iris.services.cases import get_all_cases
6
-from app.connectors.dfir_iris.schema.cases import CaseResponse
7
+from app.agents.dfir_iris.schema.cases import AssetCaseIDResponse
8
9
from app.connectors.dfir_iris.services.cases import get_all_cases
10
from app.connectors.dfir_iris.services.assets import get_case_assets
11
12
12
-async def collect_agent_soc_cases(agent_id: int) -> CaseResponse:
13
+# async def collect_agent_soc_cases(agent_id: int) -> AssetCaseIDResponse:
14
+# """
15
+# Get all cases for the given agent ID.
16
+
17
+# Args:
18
+# agent_id (int): The ID of the agent to get cases for.
19
+
20
+# Returns:
21
+# CaseResponse: An instance of CaseResponse containing the cases for the given agent ID.
22
+
23
+# Raises:
24
+# HTTPException: If the agent does not exist.
25
+# """
26
+# logger.info(f"Getting cases for agent: {agent_id}")
27
+# cases = await get_all_cases()
28
+# for case in cases.cases:
29
+# logger.info(f"Getting assets for case: {case.case_id}")
30
+# assets = await get_case_assets(case.case_id)
31
+# case_ids = []
32
+# for asset in assets.assets:
33
+# if f"agent_id:{agent_id}" in asset.asset_tags:
34
+# logger.info(f"Found case for agent: {agent_id}")
35
+# case_ids.append(case.case_id)
36
+# logger.info(f"Found cases: {case_ids}")
37
+# cases = AssetCaseID(case_ids=case_ids)
38
+# return AssetCaseIDResponse(case_ids=cases, success=True, message="Successfully retrieved cases for agent")
39
+
40
+async def collect_agent_soc_cases(agent_id: int) -> AssetCaseIDResponse:
41
"""
42
Get all cases for the given agent ID.
43
@@ -17,16 +45,50 @@ async def collect_agent_soc_cases(agent_id: int) -> CaseResponse:
45
agent_id (int): The ID of the agent to get cases for.
46
47
Returns:
20
- CaseResponse: An instance of CaseResponse containing the cases for the given agent ID.
21
-
22
- Raises:
23
- HTTPException: If the agent does not exist.
48
+ AssetCaseIDResponse: An instance of AssetCaseIDResponse containing the cases for the given agent ID.
49
"""
50
logger.info(f"Getting cases for agent: {agent_id}")
26
- cases = await get_all_cases()
27
- # for every case, get the assets
51
+ all_cases = await get_all_cases()
52
+ case_ids = await filter_cases_by_agent_id(all_cases, agent_id)
53
+
54
+ logger.info(f"Found cases: {case_ids}")
55
+ return AssetCaseIDResponse(case_ids=case_ids, success=True, message="Successfully retrieved cases for agent")
56
+
57
+
58
+async def filter_cases_by_agent_id(cases, agent_id: int) -> List[int]:
59
+ """
60
+ Filters cases by the given agent ID and collects all associated case IDs.
61
+
62
+ Args:
63
+ cases: All available cases.
64
+ agent_id (int): Agent ID to filter by.
65
+
66
+ Returns:
67
+ List[int]: List of case IDs associated with the given agent ID.
68
+ """
69
+ case_ids = []
70
for case in cases.cases:
29
- logger.info(f"Getting assets for case: {case.case_id}")
30
- assets = await get_case_assets(case.case_id)
31
- logger.info(f"Assets for case: {case.case_id} are: {assets.assets}")
32
- return None
71
+ if await is_agent_in_case(case.case_id, agent_id):
72
+ case_ids.append(case.case_id)
73
+ return case_ids
74
+
75
+
76
+async def is_agent_in_case(case_id: int, agent_id: int) -> bool:
77
+ """
78
+ Checks if a given agent ID is associated with a case.
79
+
80
+ Args:
81
+ case_id (int): The case ID to check.
82
+ agent_id (int): The agent ID to check.
83
+
84
+ Returns:
85
+ bool: True if the agent is associated with the case, False otherwise.
86
+ """
87
+ logger.info(f"Getting assets for case: {case_id}")
88
+ assets = await get_case_assets(case_id)
89
+ for asset in assets.assets:
90
+ if f"agent_id:{agent_id}" in asset.asset_tags:
91
+ logger.info(f"Found case with agent: {agent_id} in case: {case_id}")
92
+ return True
93
+ return False
94
+