@cryptotaxi247 / CoPilot / commits / bd421df5

Add SOC cases endpoint and purge cases functionality

Taylor committed Jan 11, 2024 at 10:19 UTC bd421df5bcc2647db558a8182f54e4893142b6d8
5 files changed +134 -3
backend/app/agents/dfir_iris/services/cases.py
+24
@@ -6,3 +6,27 @@ from app.connectors.dfir_iris.services.cases import get_all_cases
6 from app.connectors.dfir_iris.schema.cases import CaseResponse
7
8 from app.connectors.dfir_iris.services.cases import get_all_cases
9 +from app.connectors.dfir_iris.services.assets import get_case_assets
10 +
11 +
12 +async def collect_agent_soc_cases(agent_id: int) -> CaseResponse:
13 + """
14 + Get all cases for the given agent ID.
15 +
16 + Args:
17 + agent_id (int): The ID of the agent to get cases for.
18 +
19 + 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.
24 + """
25 + logger.info(f"Getting cases for agent: {agent_id}")
26 + cases = await get_all_cases()
27 + # for every case, get the assets
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 + logger.info(f"Assets for case: {case.case_id} are: {assets.assets}")
32 + return None
backend/app/agents/routes/agents.py
+20
@@ -20,6 +20,7 @@ from app.agents.velociraptor.services.agents import delete_agent_velociraptor
20 from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse
21 from app.agents.wazuh.services.agents import delete_agent_wazuh
22 from app.agents.wazuh.services.vulnerabilities import collect_agent_vulnerabilities
23 +from app.agents.dfir_iris.services.cases import collect_agent_soc_cases
24
25 # App specific imports
26 from app.auth.routes.auth import AuthHandler
@@ -291,6 +292,25 @@ async def get_agent_vulnerabilities(agent_id: str) -> WazuhAgentVulnerabilitiesR
292 logger.info(f"Fetching agent {agent_id} vulnerabilities")
293 return await collect_agent_vulnerabilities(agent_id)
294
295 +@agents_router.get(
296 + "/{agent_id}/soc_cases",
297 + #response_model=SocCasesResponse,
298 + description="Get SOC cases for agent",
299 + dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
300 +)
301 +async def get_agent_soc_cases(agent_id: str):
302 + """
303 + Fetches the SOC cases of a specific agent.
304 +
305 + Args:
306 + agent_id (str): The ID of the agent.
307 +
308 + Returns:
309 + SocCasesResponse: The response containing the agent SOC cases.
310 + """
311 + logger.info(f"Fetching agent {agent_id} SOC cases")
312 + return await collect_agent_soc_cases(agent_id)
313 +
314
315 @agents_router.get(
316 "/wazuh/outdated",
backend/app/connectors/dfir_iris/routes/cases.py
+17
@@ -17,6 +17,7 @@ from app.connectors.dfir_iris.services.cases import get_all_cases
17 from app.connectors.dfir_iris.services.cases import get_cases_older_than
18 from app.connectors.dfir_iris.services.cases import get_single_case
19 from app.connectors.dfir_iris.utils.universal import check_case_exists
20 +from app.connectors.dfir_iris.services.cases import purge_cases
21
22
23 async def verify_case_exists(case_id: int) -> int:
@@ -98,6 +99,22 @@ async def get_cases_older_than_route(case_older_than_body: CaseOlderThanBody = D
99 logger.info(f"Fetching all cases older than {case_older_than_body.older_than} ({case_older_than_body.time_unit.value})")
100 return await get_cases_older_than(case_older_than_body)
101
102 +@dfir_iris_cases_router.delete(
103 + "/purge",
104 + response_model=CaseResponse,
105 + description="Purge all cases",
106 + dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
107 +)
108 +async def purge_cases_route() -> CaseResponse:
109 + """
110 + Purge all cases.
111 +
112 + Returns:
113 + CaseResponse: The response containing all cases.
114 + """
115 + logger.info("Purging all cases")
116 + return await purge_cases()
117 +
118
119 @dfir_iris_cases_router.get(
120 "/{case_id}",
backend/app/connectors/dfir_iris/schema/assets.py
+7 -3
@@ -9,23 +9,27 @@ class AssetState(BaseModel):
9 object_last_update: str
10 object_state: int
11
12 +class AssetIocLink(BaseModel):
13 + ioc_id: int
14 + ioc_value: str
15 + asset_id: int
16
17 class Asset(BaseModel):
18 analysis_status: str
19 analysis_status_id: int
20 asset_compromise_status_id: Optional[int]
17 - asset_description: str
21 + asset_description: Optional[str] = Field(None, description="The description of the asset.")
22 asset_domain: Optional[str]
23 asset_icon_compromised: str
24 asset_icon_not_compromised: str
25 asset_id: int
22 - asset_ip: str
26 + asset_ip: Optional[str] = Field(None, description="The IP address of the asset.")
27 asset_name: str
28 asset_tags: Optional[str]
29 asset_type: Optional[str]
30 asset_type_id: int
31 asset_uuid: str
28 - ioc_links: Optional[None]
32 + ioc_links: Optional[List[AssetIocLink]] = None
33 link: List
34
35
backend/app/connectors/dfir_iris/services/cases.py
+66
@@ -129,3 +129,69 @@ async def get_single_case(case_id: SingleCaseBody) -> SingleCaseResponse:
129 case = Case(session=dfir_iris_client)
130 result = await fetch_and_parse_data(dfir_iris_client, case.get_case, case_id)
131 return SingleCaseResponse(success=True, message="Successfully fetched single case", case=result["data"])
132 +
133 +async def purge_cases() -> CaseResponse:
134 + """
135 + Purges all cases from DFIR-IRIS.
136 +
137 + Returns:
138 + CaseResponse: The response object containing the success status, message, and cases data.
139 +
140 + Raises:
141 + HTTPException: If there is an error purging the cases.
142 + """
143 + dfir_iris_client = await create_dfir_iris_client("DFIR-IRIS")
144 + case = Case(session=dfir_iris_client)
145 +
146 + case_ids = await get_case_ids_to_purge()
147 +
148 + for case_id in case_ids:
149 + await purge_case(dfir_iris_client, case, case_id)
150 +
151 + return CaseResponse(success=True, message="Successfully purged all cases")
152 +
153 +
154 +async def get_case_ids_to_purge() -> List[int]:
155 + """
156 + Retrieves all case IDs to be purged, skipping over specific cases as needed.
157 +
158 + Returns:
159 + List[int]: List of case IDs to purge.
160 + """
161 + result = await get_client_and_cases()
162 + handle_cases_retrieval_failure(result)
163 +
164 + # Extract case IDs, skipping over specific cases
165 + return [case["case_id"] for case in result["data"] if case["case_id"] != 1]
166 +
167 +
168 +def handle_cases_retrieval_failure(result: Dict) -> None:
169 + """
170 + Handles failure in retrieving cases.
171 +
172 + Args:
173 + result (Dict): The result of the cases retrieval attempt.
174 + """
175 + if not result["success"]:
176 + error_message = f"Failed to get all cases: {result['message']}"
177 + logger.error(error_message)
178 + raise HTTPException(status_code=500, detail=error_message)
179 +
180 +
181 +async def purge_case(client, case, case_id) -> None:
182 + """
183 + Purges a single case.
184 +
185 + Args:
186 + client: The DFIR-IRIS client.
187 + case: The case object.
188 + case_id (int): The ID of the case to purge.
189 + """
190 + try:
191 + logger.info(f"Purging case: {case_id}")
192 + await fetch_and_parse_data(client, case.delete_case, case_id)
193 + except Exception as err:
194 + error_message = f"Failed to purge case {case_id}: {err}"
195 + logger.error(error_message)
196 + raise HTTPException(status_code=500, detail=error_message)
197 +