@cryptotaxi247 / CoPilot / commits / b3f39683

Refactor (#136)

* Split Dockerfile into two containers * split project (be/fe) * updated docker * Update logo path in README.md * Move image file to frontend directory * updated nginx config * updated docker file * updated nginx * Add APIRouter for modularizing routes * updated nginx * updated ignore file * Update Dockerfile and docker-compose.yml * added env file * new tmp dockerfile for fe * docker changes * Remove port mapping for copilot-backend container and update nginx.conf for API endpoint but still not working as expected * updated fe package.json * fix httpclient headers * updated dependencies * Fix authentication bug and improve error handling * updated docker * Update Office365 integration documentation and Docker configuration * updated images * Delete unused files and update Docker configuration * Update Dockerfile to expose port 2000 * added ssl * updated ssl generator * precommit fixes * ignore ssl files * remove nginx certs * updated ssl script * add default ss keys * Update agent SOC cases and bookmarked alerts * precommit fixes * Add Docker installation check and update frontend URL * Update image reference for copilot-frontend container * Add SERVER_IP configuration for remote machine connection * Add build-dockers.sh script to build copilot-frontend image * Remove unnecessary SSL configuration * Update API URL in .env.example file * Refactor nginx.conf for improved logging and HTTPS redirection * Update Copilot port and protocol --------- Co-authored-by: Davide Di Modica <webmaster.ddm@gmail.com>

taylor_socfortress committed Feb 9, 2024 at 09:53 UTC b3f39683269d57d9d24b84cf9e7a577e49623ac9

583 files changed +7923 -16749

.dockerignore deleted

-2

@@ -1,2 +0,0 @@

1 -backend/data/copilot.db

2 -backend/copilot.db

.env.example

+2 -12

@@ -1,15 +1,5 @@

1 -# base url

2 -SERVER_IP=YOUR_SERVER_IP

3 -VITE_API_URL=http://YOUR_SERVER_IP:5000

4 -

5 -# value in seconds

6 -VITE_TOKEN_DEBOUNCE_TIME=10

7 -

8 -# alert if value is over

9 -VITE_UNCOMMITTED_JOURNAL_ENTRIES_THRESHOLD=50000

10 -

11 -# value in seconds

12 -VITE_HEALTHCHECKS_INTERVAL=120

1 +# Leave this as is if connecting from a remote machine

2 +SERVER_IP=0.0.0.0

4 # Connector Credentials

5 # ! SETTING UP YOUR CONNECTORS DEMOs https://www.youtube.com/@taylorwalton_socfortress/videos! #

.gitignore

+7 -3

@@ -1,6 +1,10 @@

1 +#server.key

2 +server.csr

3 +#server.crt

4 +

5 # Logs

6 logs

3 -!src/components/logs

7 +!frontend/src/components/logs

8 *.log

9 npm-debug.log*

10 yarn-debug.log*

@@ -15,8 +19,8 @@ dist-ssr

19 coverage

20 *.local

18 -/cypress/videos/

19 -/cypress/screenshots/

22 +cypress/videos/

23 +cypress/screenshots/

25 # Editor directories and files

26 .vscode/*

.nvmrc deleted

-1

@@ -1 +0,0 @@

1 -18.17.1

README.md

+6 -3

@@ -1,6 +1,6 @@

1 <h1 align="center">

3 -<a href="https://www.socfortress.co"><img src="src/assets/images/socfortress_logo.svg" width="300" height="200"></a>

3 +<a href="https://www.socfortress.co"><img src="frontend/src/assets/images/socfortress_logo.svg" width="300" height="200"></a>

5 SOCFortress CoPilot

@@ -13,7 +13,7 @@ SOCFortress CoPilot

14 [SOCFortress CoPilot](https://www.socfortress.co) focuses on providing a single pane of glass for all your security operations needs. Simplify your open source security stack with a single platform focused on making open source security tools easier to use and more accessible.

16 -![demo_timeline](src/assets/images/copilot_gif.gif)

16 +![demo_timeline](frontend/src/assets/images/copilot_gif.gif)

18 ## Table of contents

@@ -71,11 +71,14 @@ cp .env.example .env

72 # Make your changes to the .env file

74 +# Build the copilot-frontend image

75 +bash build-dockers.sh

76 +

77 # Run Copilot

78 docker compose up -d

79 ```

78 -Copilot shall be available on the host interface, port 5173, protocol HTTP - `http://<your_instance_ip>:5173`.

81 +Copilot shall be available on the host interface, port 443, protocol HTTPS - `https://<your_instance_ip>`.

82 By default, an `admin` account is created. The password is printed in stdout the very first time Copilot is started. It won't be printed anymore after that.

83 `Admin user password` can be searched in the logs of the `copilot` docker to find the password. You will use the `plain` password to login to the web interface.

backend/.dockerignore new

@@ -0,0 +1,2 @@

1 +data/copilot.db

2 +copilot.db

backend/Dockerfile renamed

+5 -16

@@ -1,6 +1,7 @@

1 # build with `docker build -t python-backend -f Dockerfile.deb .`

2 # run with `docker run -p 5000:5000 -d python-backend`

3 # Start with the base Debian 11 image

4 +# looking to split into 2 containers, one for the backend and one for the frontend

5 FROM debian:11

7 # Set environment variables

@@ -37,27 +38,15 @@ RUN /opt/venv/bin/pip install setuptools

39 # Install your application's dependencies

40 WORKDIR /opt/copilot/backend

40 -COPY backend/requirements.txt ./

41 +COPY requirements.txt ./

42 RUN /opt/venv/bin/pip install --no-cache-dir -r requirements.txt

44 # Copy your application into the Docker image

44 -WORKDIR /opt/copilot

45 +WORKDIR /opt/copilot/backend

46 COPY . .

47 -# Install Node.js and npm

48 -RUN curl -sL https://deb.nodesource.com/setup_18.x | bash -

49 -RUN apt-get install -y nodejs

50 -

51 -# Install concurrently

52 -RUN npm install -g concurrently

53 -

54 -# Install your Vue.js application's dependencies

55 -WORKDIR /opt/copilot

56 -RUN npm install

57 -

48 # Expose ports

59 -EXPOSE 5000 5173

49 +EXPOSE 5000

51 # Run your application

62 -#CMD ["sh", "-c", "cd backend && python copilot.py & cd /opt/copilot && npm run dev"]

63 -CMD ["sh", "-c", "cd /opt/copilot && npm run start"]

52 +CMD ["sh", "-c", "ls -la && /opt/venv/bin/python copilot.py"]

backend/app/agents/dfir_iris/services/cases.py

+12 -5

@@ -1,13 +1,16 @@

1 from typing import List

3 -from loguru import logger

4 -

3 from app.agents.dfir_iris.schema.cases import AssetCaseIDResponse

4 from app.connectors.dfir_iris.services.assets import get_case_assets

5 from app.connectors.dfir_iris.services.cases import get_all_cases

6 +from loguru import logger

7 +from sqlalchemy.ext.asyncio import AsyncSession

10 -async def collect_agent_soc_cases(agent_id: int) -> AssetCaseIDResponse:

10 +async def collect_agent_soc_cases(

11 + agent_id: int,

12 + session: AsyncSession,

13 +) -> AssetCaseIDResponse:

14 """

15 Get all cases for the given agent ID.

@@ -18,11 +21,15 @@ async def collect_agent_soc_cases(agent_id: int) -> AssetCaseIDResponse:

21 AssetCaseIDResponse: An instance of AssetCaseIDResponse containing the cases for the given agent ID.

22 """

23 logger.info(f"Getting cases for agent: {agent_id}")

21 - all_cases = await get_all_cases()

24 + all_cases = await get_all_cases(session=session)

25 case_ids = await filter_cases_by_agent_id(all_cases, agent_id)

27 logger.info(f"Found cases: {case_ids}")

25 - return AssetCaseIDResponse(case_ids=case_ids, success=True, message="Successfully retrieved cases for agent")

28 + return AssetCaseIDResponse(

29 + case_ids=case_ids,

30 + success=True,

31 + message="Successfully retrieved cases for agent",

32 + )

35 async def filter_cases_by_agent_id(cases, agent_id: int) -> List[int]:

backend/app/agents/routes/agents.py

+127 -47

@@ -1,21 +1,15 @@

1 -from fastapi import APIRouter

2 -from fastapi import BackgroundTasks

3 -from fastapi import Depends

4 -from fastapi import HTTPException

5 -from fastapi import Security

6 -from loguru import logger

7 -from sqlalchemy import delete

8 -from sqlalchemy.ext.asyncio import AsyncSession

9 -from sqlalchemy.future import select

10 -

1 from app.agents.dfir_iris.services.cases import collect_agent_soc_cases

12 -from app.agents.schema.agents import AgentModifyResponse

13 -from app.agents.schema.agents import AgentsResponse

14 -from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse

15 -from app.agents.schema.agents import OutdatedWazuhAgentsResponse

16 -from app.agents.schema.agents import SyncedAgentsResponse

17 -from app.agents.services.status import get_outdated_agents_velociraptor

18 -from app.agents.services.status import get_outdated_agents_wazuh

2 +from app.agents.schema.agents import (

3 + AgentModifyResponse,

4 + AgentsResponse,

5 + OutdatedVelociraptorAgentsResponse,

6 + OutdatedWazuhAgentsResponse,

7 + SyncedAgentsResponse,

8 +)

9 +from app.agents.services.status import (

10 + get_outdated_agents_velociraptor,

11 + get_outdated_agents_wazuh,

12 +)

13 from app.agents.services.sync import sync_agents

14 from app.agents.velociraptor.services.agents import delete_agent_velociraptor

15 from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse

@@ -29,6 +23,11 @@ from app.db.db_session import get_db

23 # App specific imports

24 # from app.db.db_session import session

25 from app.db.universal_models import Agents

26 +from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Security

27 +from loguru import logger

28 +from sqlalchemy import delete

29 +from sqlalchemy.ext.asyncio import AsyncSession

30 +from sqlalchemy.future import select

32 agents_router = APIRouter()

@@ -57,7 +56,10 @@ async def fetch_velociraptor_id(db: AsyncSession, agent_id: str) -> str:

56 raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found")

57 except Exception as e:

58 logger.error(f"Failed to fetch agent {agent_id} from database: {e}")

60 - raise HTTPException(status_code=500, detail=f"Failed to fetch agent {agent_id} from database: {e}")

59 + raise HTTPException(

60 + status_code=500,

61 + detail=f"Failed to fetch agent {agent_id} from database: {e}",

62 + )

65 async def delete_agent_from_database(db: AsyncSession, agent_id: str):

@@ -78,7 +80,10 @@ async def delete_agent_from_database(db: AsyncSession, agent_id: str):

80 except Exception as e:

81 logger.error(f"Failed to delete agent {agent_id} from database: {e}")

82 await db.rollback()

81 - raise HTTPException(status_code=500, detail=f"Failed to delete agent {agent_id} from database: {e}")

83 + raise HTTPException(

84 + status_code=500,

85 + detail=f"Failed to delete agent {agent_id} from database: {e}",

86 + )

89 @agents_router.get(

@@ -101,7 +106,11 @@ async def get_agents(db: AsyncSession = Depends(get_db)) -> AgentsResponse:

106 try:

107 result = await db.execute(select(Agents))

108 agents = result.scalars().all()

104 - return AgentsResponse(agents=agents, success=True, message="Agents fetched successfully")

109 + return AgentsResponse(

110 + agents=agents,

111 + success=True,

112 + message="Agents fetched successfully",

113 + )

114 except Exception as e:

115 logger.error(f"Failed to fetch agents: {e}")

116 raise HTTPException(status_code=500, detail=f"Failed to fetch agents: {e}")

@@ -113,7 +122,10 @@ async def get_agents(db: AsyncSession = Depends(get_db)) -> AgentsResponse:

122 description="Get agent by agent_id",

123 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

124 )

116 -async def get_agent(agent_id: str, db: AsyncSession = Depends(get_db)) -> AgentsResponse:

125 +async def get_agent(

126 + agent_id: str,

127 + db: AsyncSession = Depends(get_db),

128 +) -> AgentsResponse:

129 """

130 Retrieve an agent by agent_id.

131

@@ -132,12 +144,24 @@ async def get_agent(agent_id: str, db: AsyncSession = Depends(get_db)) -> Agents

144 result = await db.execute(select(Agents).filter(Agents.agent_id == agent_id))

145 agent = result.scalars().first()

146 if agent:

135 - return AgentsResponse(agents=[agent], success=True, message="Agent fetched successfully")

147 + return AgentsResponse(

148 + agents=[agent],

149 + success=True,

150 + message="Agent fetched successfully",

151 + )

152 else:

137 - raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")

153 + raise HTTPException(

154 + status_code=404,

155 + detail=f"Agent with agent_id {agent_id} not found",

156 + )

157 except Exception as e:

139 - logger.error(f"Failed to fetch agent: {agent_id} with error {e}. Does it exist?")

140 - raise HTTPException(status_code=500, detail=f"Failed to fetch agent: {agent_id}. Does it exist?")

158 + logger.error(

159 + f"Failed to fetch agent: {agent_id} with error {e}. Does it exist?",

160 + )

161 + raise HTTPException(

162 + status_code=500,

163 + detail=f"Failed to fetch agent: {agent_id}. Does it exist?",

164 + )

165

166

167 @agents_router.get(

@@ -146,7 +170,10 @@ async def get_agent(agent_id: str, db: AsyncSession = Depends(get_db)) -> Agents

170 description="Get agent by hostname",

171 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

172 )

149 -async def get_agent_by_hostname(hostname: str, db: AsyncSession = Depends(get_db)) -> AgentsResponse:

173 +async def get_agent_by_hostname(

174 + hostname: str,

175 + db: AsyncSession = Depends(get_db),

176 +) -> AgentsResponse:

177 """

178 Retrieve an agent by its hostname.

179

@@ -165,9 +192,16 @@ async def get_agent_by_hostname(hostname: str, db: AsyncSession = Depends(get_db

192 result = await db.execute(select(Agents).filter(Agents.hostname == hostname))

193 agent = result.scalars().first()

194 if agent:

168 - return AgentsResponse(agents=[agent], success=True, message="Agent fetched successfully")

195 + return AgentsResponse(

196 + agents=[agent],

197 + success=True,

198 + message="Agent fetched successfully",

199 + )

200 else:

170 - raise HTTPException(status_code=404, detail=f"Agent with hostname {hostname} not found")

201 + raise HTTPException(

202 + status_code=404,

203 + detail=f"Agent with hostname {hostname} not found",

204 + )

205 except Exception as e:

206 logger.error(f"Failed to fetch agent: {e}")

207 # The exception message should not be exposed directly, especially in production

@@ -178,9 +212,14 @@ async def get_agent_by_hostname(hostname: str, db: AsyncSession = Depends(get_db

212 "/sync",

213 response_model=SyncedAgentsResponse,

214 description="Sync agents from Wazuh Manager",

181 - dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "scheduler"))],

215 + dependencies=[

216 + Security(AuthHandler().require_any_scope("admin", "analyst", "scheduler")),

217 + ],

218 )

183 -async def sync_all_agents(backgroud_tasks: BackgroundTasks, session: AsyncSession = Depends(get_db)) -> SyncedAgentsResponse:

219 +async def sync_all_agents(

220 + backgroud_tasks: BackgroundTasks,

221 + session: AsyncSession = Depends(get_db),

222 +) -> SyncedAgentsResponse:

223 """

224 Sync all agents from Wazuh Manager.

225

@@ -197,7 +236,10 @@ async def sync_all_agents(backgroud_tasks: BackgroundTasks, session: AsyncSessio

236 """

237 logger.info("Syncing agents from Wazuh Manager")

238 backgroud_tasks.add_task(sync_agents, session)

200 - return SyncedAgentsResponse(success=True, message="Agents synced started successfully")

239 + return SyncedAgentsResponse(

240 + success=True,

241 + message="Agents synced started successfully",

242 + )

243

244

245 @agents_router.post(

@@ -206,7 +248,10 @@ async def sync_all_agents(backgroud_tasks: BackgroundTasks, session: AsyncSessio

248 description="Mark agent as critical",

249 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

250 )

209 -async def mark_agent_as_critical(agent_id: str, session: AsyncSession = Depends(get_db)) -> AgentModifyResponse:

251 +async def mark_agent_as_critical(

252 + agent_id: str,

253 + session: AsyncSession = Depends(get_db),

254 +) -> AgentModifyResponse:

255 """

256 Marks the specified agent as critical.

257

@@ -220,19 +265,30 @@ async def mark_agent_as_critical(agent_id: str, session: AsyncSession = Depends(

265 logger.info(f"Marking agent {agent_id} as critical")

266 try:

267 # Asynchronously fetch the agent by id

223 - result = await session.execute(select(Agents).filter(Agents.agent_id == agent_id))

268 + result = await session.execute(

269 + select(Agents).filter(Agents.agent_id == agent_id),

270 + )

271 agent = result.scalars().first()

272

273 if not agent:

227 - raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")

274 + raise HTTPException(

275 + status_code=404,

276 + detail=f"Agent with agent_id {agent_id} not found",

277 + )

278

279 agent.critical_asset = True

280 await session.commit()

281

232 - return AgentModifyResponse(success=True, message=f"Agent {agent_id} marked as critical: {True}")

282 + return AgentModifyResponse(

283 + success=True,

284 + message=f"Agent {agent_id} marked as critical: {True}",

285 + )

286 except Exception as e:

287 session.rollback() # Roll back the session in case of error

235 - raise HTTPException(status_code=500, detail=f"Failed to mark agent as critical: {str(e)}")

288 + raise HTTPException(

289 + status_code=500,

290 + detail=f"Failed to mark agent as critical: {str(e)}",

291 + )

292

293

294 @agents_router.post(

@@ -241,7 +297,10 @@ async def mark_agent_as_critical(agent_id: str, session: AsyncSession = Depends(

297 description="Mark agent as not critical",

298 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

299 )

244 -async def mark_agent_as_not_critical(agent_id: str, session: AsyncSession = Depends(get_db)) -> AgentModifyResponse:

300 +async def mark_agent_as_not_critical(

301 + agent_id: str,

302 + session: AsyncSession = Depends(get_db),

303 +) -> AgentModifyResponse:

304 """

305 Marks the specified agent as not critical.

306

@@ -257,19 +316,30 @@ async def mark_agent_as_not_critical(agent_id: str, session: AsyncSession = Depe

316 """

317 logger.info(f"Marking agent {agent_id} as not critical")

318 try:

260 - result = await session.execute(select(Agents).filter(Agents.agent_id == agent_id))

319 + result = await session.execute(

320 + select(Agents).filter(Agents.agent_id == agent_id),

321 + )

322 agent = result.scalars().first()

323

324 if not agent:

264 - raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")

325 + raise HTTPException(

326 + status_code=404,

327 + detail=f"Agent with agent_id {agent_id} not found",

328 + )

329

330 agent.critical_asset = False

331 await session.commit()

332

269 - return AgentModifyResponse(success=True, message=f"Agent {agent_id} marked as not critical")

333 + return AgentModifyResponse(

334 + success=True,

335 + message=f"Agent {agent_id} marked as not critical",

336 + )

337 except Exception as e:

338 await session.rollback() # Roll back the session in case of error

272 - raise HTTPException(status_code=500, detail=f"Failed to mark agent as not critical: {str(e)}")

339 + raise HTTPException(

340 + status_code=500,

341 + detail=f"Failed to mark agent as not critical: {str(e)}",

342 + )

343

344

345 @agents_router.get(

@@ -298,7 +368,7 @@ async def get_agent_vulnerabilities(agent_id: str) -> WazuhAgentVulnerabilitiesR

368 description="Get SOC cases for agent",

369 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

370 )

301 -async def get_agent_soc_cases(agent_id: str):

371 +async def get_agent_soc_cases(agent_id: str, session: AsyncSession = Depends(get_db)):

372 """

373 Fetches the SOC cases of a specific agent.

374

@@ -309,7 +379,7 @@ async def get_agent_soc_cases(agent_id: str):

379 SocCasesResponse: The response containing the agent SOC cases.

380 """

381 logger.info(f"Fetching agent {agent_id} SOC cases")

312 - return await collect_agent_soc_cases(agent_id)

382 + return await collect_agent_soc_cases(agent_id, session)

383

384

385 @agents_router.get(

@@ -318,7 +388,9 @@ async def get_agent_soc_cases(agent_id: str):

388 description="Get all outdated Wazuh agents",

389 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

390 )

321 -async def get_outdated_wazuh_agents(session: AsyncSession = Depends(get_db)) -> OutdatedWazuhAgentsResponse:

391 +async def get_outdated_wazuh_agents(

392 + session: AsyncSession = Depends(get_db),

393 +) -> OutdatedWazuhAgentsResponse:

394 """

395 Retrieve all outdated Wazuh agents.

396

@@ -337,7 +409,9 @@ async def get_outdated_wazuh_agents(session: AsyncSession = Depends(get_db)) ->

409 description="Get all outdated Velociraptor agents",

410 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

411 )

340 -async def get_outdated_velociraptor_agents(session: AsyncSession = Depends(get_db)) -> OutdatedVelociraptorAgentsResponse:

412 +async def get_outdated_velociraptor_agents(

413 + session: AsyncSession = Depends(get_db),

414 +) -> OutdatedVelociraptorAgentsResponse:

415 """

416 Fetches all outdated Velociraptor agents.

417

@@ -358,7 +432,10 @@ async def get_outdated_velociraptor_agents(session: AsyncSession = Depends(get_d

432 description="Delete agent",

433 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

434 )

361 -async def delete_agent(agent_id: str, session: AsyncSession = Depends(get_db)) -> AgentModifyResponse:

435 +async def delete_agent(

436 + agent_id: str,

437 + session: AsyncSession = Depends(get_db),

438 +) -> AgentModifyResponse:

439 """

440 Delete an agent.

441

@@ -376,7 +453,10 @@ async def delete_agent(agent_id: str, session: AsyncSession = Depends(get_db)) -

453 if client_id != "n/a":

454 await delete_agent_velociraptor(client_id)

455 await delete_agent_from_database(db=session, agent_id=agent_id)

379 - return AgentModifyResponse(success=True, message=f"Agent {agent_id} deleted successfully")

456 + return AgentModifyResponse(

457 + success=True,

458 + message=f"Agent {agent_id} deleted successfully",

459 + )

460

461

462 # ! TODO: CURRENTLY UPDATES IN THE DB BUT NEED TO UPDATE IN WAZUH # !

backend/app/agents/schema/agents.py

+1 -3

@@ -1,11 +1,9 @@

1 from typing import List

3 -from pydantic import BaseModel

4 -from pydantic import Field

5 -

3 from app.agents.velociraptor.schema.agents import VelociraptorAgent

4 from app.agents.wazuh.schema.agents import WazuhAgent

5 from app.db.universal_models import Agents

6 +from pydantic import BaseModel, Field

9 class AgentsResponse(BaseModel):

backend/app/agents/services/modify.py

+9 -4

@@ -1,8 +1,7 @@

1 -from fastapi import HTTPException

2 -

1 import app.agents.wazuh.services.agents as wazuh_services

2 from app.db.db_session import session

3 from app.db.universal_models import Agents

4 +from fastapi import HTTPException

7 def delete_agent_db(agent_id: str):

@@ -17,7 +16,10 @@ def delete_agent_db(agent_id: str):

16 """

17 agent = session.query(Agents).filter(Agents.agent_id == agent_id).first()

18 if not agent:

20 - raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found")

19 + raise HTTPException(

20 + status_code=404,

21 + detail=f"Agent with agent_id {agent_id} not found",

22 + )

23 session.delete(agent)

24 session.commit()

25 return {"success": True, "message": f"Agent {agent_id} deleted from database"}

@@ -40,4 +42,7 @@ def delete_agent_wazuh(agent_id: str):

42 wazuh_services.delete_agent(agent_id)

43 return {"success": True, "message": f"Agent {agent_id} deleted from Wazuh"}

44 except Exception as e:

43 - raise HTTPException(status_code=500, detail=f"Failed to delete agent {agent_id} from Wazuh: {e}")

45 + raise HTTPException(

46 + status_code=500,

47 + detail=f"Failed to delete agent {agent_id} from Wazuh: {e}",

48 + )

backend/app/agents/services/status.py

+44 -16

@@ -1,16 +1,17 @@

1 from typing import List

3 +from app.agents.schema.agents import (

4 + OutdatedVelociraptorAgentsResponse,

5 + OutdatedWazuhAgentsResponse,

6 +)

7 +from app.connectors.velociraptor.utils.universal import UniversalService

8 +from app.db.db_session import session

9 +from app.db.universal_models import Agents

10 from fastapi import HTTPException

11 from loguru import logger

12 from sqlalchemy.ext.asyncio import AsyncSession

13 from sqlalchemy.future import select

8 -from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse

9 -from app.agents.schema.agents import OutdatedWazuhAgentsResponse

10 -from app.connectors.velociraptor.utils.universal import UniversalService

11 -from app.db.db_session import session

12 -from app.db.universal_models import Agents

13 -

16 def get_agent(agent_id: str) -> List[Agents]:

17 """

@@ -26,10 +27,15 @@ def get_agent(agent_id: str) -> List[Agents]:

27 return session.query(Agents).filter(Agents.agent_id == agent_id).first()

28 except Exception as e:

29 logger.error(f"Failed to fetch agent with agent_id {agent_id}: {e}")

29 - raise HTTPException(status_code=500, detail=f"Failed to fetch agent with agent_id {agent_id}: {e}")

30 + raise HTTPException(

31 + status_code=500,

32 + detail=f"Failed to fetch agent with agent_id {agent_id}: {e}",

33 + )

32 -async def get_outdated_agents_wazuh(session: AsyncSession) -> OutdatedWazuhAgentsResponse:

36 +async def get_outdated_agents_wazuh(

37 + session: AsyncSession,

38 +) -> OutdatedWazuhAgentsResponse:

39 """

40 Retrieves all agents with outdated Wazuh agent versions from the database asynchronously.

@@ -40,15 +46,23 @@ async def get_outdated_agents_wazuh(session: AsyncSession) -> OutdatedWazuhAgent

46 OutdatedWazuhAgentsResponse: Response object containing the outdated agents.

47 """

48 try:

43 - wazuh_manager_result = await session.execute(select(Agents).filter(Agents.agent_id == "000"))

49 + wazuh_manager_result = await session.execute(

50 + select(Agents).filter(Agents.agent_id == "000"),

51 + )

52 wazuh_manager = wazuh_manager_result.scalars().first()

54 if wazuh_manager is None:

55 logger.error("Wazuh Manager with agent_id '000' not found.")

48 - raise HTTPException(status_code=404, detail="Wazuh Manager with agent_id '000' not found.")

56 + raise HTTPException(

57 + status_code=404,

58 + detail="Wazuh Manager with agent_id '000' not found.",

59 + )

61 outdated_agents_result = await session.execute(

51 - select(Agents).filter(Agents.agent_id != "000", Agents.wazuh_agent_version != wazuh_manager.wazuh_agent_version),

62 + select(Agents).filter(

63 + Agents.agent_id != "000",

64 + Agents.wazuh_agent_version != wazuh_manager.wazuh_agent_version,

65 + ),

66 )

67 outdated_wazuh_agents = outdated_agents_result.scalars().all()

@@ -58,10 +72,15 @@ async def get_outdated_agents_wazuh(session: AsyncSession) -> OutdatedWazuhAgent

72 outdated_wazuh_agents=outdated_wazuh_agents,

73 )

74 except Exception as e:

61 - raise HTTPException(status_code=500, detail=f"Failed to fetch outdated Wazuh agents: {e}")

75 + raise HTTPException(

76 + status_code=500,

77 + detail=f"Failed to fetch outdated Wazuh agents: {e}",

78 + )

64 -async def get_outdated_agents_velociraptor(session: AsyncSession) -> OutdatedVelociraptorAgentsResponse:

81 +async def get_outdated_agents_velociraptor(

82 + session: AsyncSession,

83 +) -> OutdatedVelociraptorAgentsResponse:

84 """

85 Retrieves all agents with outdated Velociraptor client versions from the database asynchronously.

@@ -76,10 +95,16 @@ async def get_outdated_agents_velociraptor(session: AsyncSession) -> OutdatedVel

96 try:

97 # Assuming _get_server_version is an async function

79 - server_version = await velociraptor_service._get_server_version(vql_server_version)

98 + server_version = await velociraptor_service._get_server_version(

99 + vql_server_version,

100 + )

101 agents_result = await session.execute(select(Agents))

102 agents = agents_result.scalars().all()

82 - outdated_velociraptor_agents = [agent for agent in agents if agent.velociraptor_agent_version != server_version]

103 + outdated_velociraptor_agents = [

104 + agent

105 + for agent in agents

106 + if agent.velociraptor_agent_version != server_version

107 + ]

108

109 return OutdatedVelociraptorAgentsResponse(

110 message="Outdated Velociraptor agents fetched successfully.",

@@ -87,4 +112,7 @@ async def get_outdated_agents_velociraptor(session: AsyncSession) -> OutdatedVel

112 outdated_velociraptor_agents=outdated_velociraptor_agents,

113 )

114 except Exception as e:

90 - raise HTTPException(status_code=500, detail=f"Failed to fetch outdated Velociraptor agents: {e}")

115 + raise HTTPException(

116 + status_code=500,

117 + detail=f"Failed to fetch outdated Velociraptor agents: {e}",

118 + )

backend/app/agents/services/sync.py

+38 -15

@@ -1,18 +1,15 @@

1 from typing import List

3 -from loguru import logger

4 -from sqlalchemy.ext.asyncio import AsyncSession

5 -from sqlalchemy.future import select

6 -

3 import app.agents.velociraptor.services.agents as velociraptor_services

4 import app.agents.wazuh.services.agents as wazuh_services

9 -from app.agents.schema.agents import SyncedAgent

10 -from app.agents.schema.agents import SyncedAgentsResponse

5 +from app.agents.schema.agents import SyncedAgent, SyncedAgentsResponse

6 from app.agents.velociraptor.schema.agents import VelociraptorAgent

12 -from app.agents.wazuh.schema.agents import WazuhAgent

13 -from app.agents.wazuh.schema.agents import WazuhAgentsList

7 +from app.agents.wazuh.schema.agents import WazuhAgent, WazuhAgentsList

8 from app.connectors.models import Connectors

9 from app.db.universal_models import Agents

10 +from loguru import logger

11 +from sqlalchemy.ext.asyncio import AsyncSession

12 +from sqlalchemy.future import select

15 async def fetch_wazuh_agents() -> WazuhAgentsList:

@@ -47,7 +44,12 @@ async def fetch_velociraptor_agent(agent_name: str) -> VelociraptorAgent:

44 return await velociraptor_services.collect_velociraptor_agent(agent_name)

50 -async def add_agent_to_db(session: AsyncSession, agent: WazuhAgent, client: VelociraptorAgent, customer_code: str):

47 +async def add_agent_to_db(

48 + session: AsyncSession,

49 + agent: WazuhAgent,

50 + client: VelociraptorAgent,

51 + customer_code: str,

52 +):

53 """Add new agent to database.

55 Args:

@@ -114,7 +116,9 @@ async def get_velociraptor_connector(session):

116 Returns:

117 The first result of the query as a scalar value.

118 """

117 - connector_query = select(Connectors).filter(Connectors.connector_name == "Velociraptor")

119 + connector_query = select(Connectors).filter(

120 + Connectors.connector_name == "Velociraptor",

121 + )

122 result = await session.execute(connector_query)

123 return result.scalars().first()

124

@@ -187,24 +191,43 @@ async def sync_agents(session: AsyncSession) -> SyncedAgentsResponse:

191 try:

192 velociraptor_agent = await process_velociraptor_agent(session, wazuh_agent)

193 except Exception as e:

190 - logger.error(f"Failed to collect Velociraptor Agent for {wazuh_agent.agent_name}: {e}")

194 + logger.error(

195 + f"Failed to collect Velociraptor Agent for {wazuh_agent.agent_name}: {e}",

196 + )

197 continue

198

199 customer_code = extract_customer_code(wazuh_agent.agent_label)

200

201 # Asynchronously fetch the existing agent

196 - existing_agent_query = select(Agents).filter(Agents.hostname == wazuh_agent.agent_name)

202 + existing_agent_query = select(Agents).filter(

203 + Agents.hostname == wazuh_agent.agent_name,

204 + )

205 result = await session.execute(existing_agent_query)

206 existing_agent = result.scalars().first()

207

208 if existing_agent:

201 - await update_agent_in_db(session, existing_agent, wazuh_agent, velociraptor_agent, customer_code)

209 + await update_agent_in_db(

210 + session,

211 + existing_agent,

212 + wazuh_agent,

213 + velociraptor_agent,

214 + customer_code,

215 + )

216 else:

203 - await add_agent_to_db(session, wazuh_agent, velociraptor_agent, customer_code)

217 + await add_agent_to_db(

218 + session,

219 + wazuh_agent,

220 + velociraptor_agent,

221 + customer_code,

222 + )

223

224 # Combine the wazuh agent and velociraptor agent into one object

225 synced_agent = SyncedAgent(**wazuh_agent.dict(), **velociraptor_agent.dict())

226 agents_added_list.append(synced_agent)

227

228 logger.info(f"Agents Added List: {agents_added_list}")

210 - return SyncedAgentsResponse(success=True, message="Agents synced successfully", agents_added=agents_added_list)

229 + return SyncedAgentsResponse(

230 + success=True,

231 + message="Agents synced successfully",

232 + agents_added=agents_added_list,

233 + )

backend/app/agents/velociraptor/schema/agents.py

+1 -2

@@ -1,8 +1,7 @@

1 from datetime import datetime

2 from typing import Optional

4 -from pydantic import BaseModel

5 -from pydantic import Field

4 +from pydantic import BaseModel, Field

7 class VelociraptorAgent(BaseModel):

backend/app/agents/velociraptor/services/agents.py

+51 -16

@@ -1,11 +1,10 @@

1 from datetime import datetime

3 -from fastapi import HTTPException

4 -from loguru import logger

5 -

3 from app.agents.schema.agents import AgentModifyResponse

4 from app.agents.velociraptor.schema.agents import VelociraptorAgent

5 from app.connectors.velociraptor.utils.universal import UniversalService

6 +from fastapi import HTTPException

7 +from loguru import logger

10 def create_query(query: str) -> str:

@@ -39,29 +38,45 @@ async def collect_velociraptor_agent(agent_name: str) -> VelociraptorAgent:

38 client_id = client_id["results"][0]["client_id"]

39 except (KeyError, IndexError, TypeError) as e:

40 logger.error(f"Failed to get client ID for {agent_name}. Error: {e}")

42 - return VelociraptorAgent(client_id="Unknown", client_last_seen="Unknown", client_version="Unknown")

41 + return VelociraptorAgent(

42 + client_id="Unknown",

43 + client_last_seen="Unknown",

44 + client_version="Unknown",

45 + )

47 try:

45 - vql_last_seen_at = f"select last_seen_at from clients(search='host:{agent_name}')"

46 - last_seen_at = await velociraptor_service._get_last_seen_timestamp(vql_last_seen_at)

48 + vql_last_seen_at = (

49 + f"select last_seen_at from clients(search='host:{agent_name}')"

50 + )

51 + last_seen_at = await velociraptor_service._get_last_seen_timestamp(

52 + vql_last_seen_at,

53 + )

54 client_last_seen = datetime.fromtimestamp(

55 int(last_seen_at) / 1000000,

56 ).strftime(

57 "%Y-%m-%dT%H:%M:%S+00:00",

58 ) # Converting to string format

59 except Exception as e:

53 - logger.error(f"Failed to get or convert last seen at for {agent_name}. Error: {e}")

60 + logger.error(

61 + f"Failed to get or convert last seen at for {agent_name}. Error: {e}",

62 + )

63 client_last_seen = "1970-01-01T00:00:00+00:00"

65 try:

66 vql_client_version = f"select * from clients(search='host:{agent_name}')"

67 # client_version = UniversalService()._get_client_version(vql_client_version)

59 - client_version = await velociraptor_service._get_client_version(vql_client_version)

68 + client_version = await velociraptor_service._get_client_version(

69 + vql_client_version,

70 + )

71 except Exception as e:

72 logger.error(f"Failed to get client version for {agent_name}. Error: {e}")

73 client_version = "Unknown"

64 - return VelociraptorAgent(client_id=client_id, client_last_seen=client_last_seen, client_version=client_version)

75 + return VelociraptorAgent(

76 + client_id=client_id,

77 + client_last_seen=client_last_seen,

78 + client_version=client_version,

79 + )

82 def execute_query(universal_service, query: str) -> dict:

@@ -96,10 +111,16 @@ def check_flow_success(flow: dict, client_id: str) -> dict:

111 """

112 if flow["success"]:

113 logger.info(f"Successfully deleted velociraptor client {client_id}")

99 - return {"message": f"Successfully deleted velociraptor client {client_id}", "success": True}

114 + return {

115 + "message": f"Successfully deleted velociraptor client {client_id}",

116 + "success": True,

117 + }

118 else:

119 logger.error(f"Failed to delete velociraptor client {client_id}")

102 - return handle_exception(e="Failed to delete velociraptor client", client_id=client_id)

120 + return handle_exception(

121 + e="Failed to delete velociraptor client",

122 + client_id=client_id,

123 + )

124

125

126 def check_client_in_results(results: dict, client_id: str) -> dict:

@@ -115,12 +136,18 @@ def check_client_in_results(results: dict, client_id: str) -> dict:

136 """

137 if results["results"] == []:

138 logger.info(f"Successfully deleted velociraptor client {client_id}")

118 - return {"message": f"Successfully deleted velociraptor client {client_id}", "success": True}

139 + return {

140 + "message": f"Successfully deleted velociraptor client {client_id}",

141 + "success": True,

142 + }

143

144 for result in results["results"]:

145 if result["client_id"] == client_id:

146 logger.error(f"Failed to delete velociraptor client {client_id}")

123 - return handle_exception(e="Failed to delete velociraptor client", client_id=client_id)

147 + return handle_exception(

148 + e="Failed to delete velociraptor client",

149 + client_id=client_id,

150 + )

151

152

153 def handle_exception(e: Exception, client_id: str) -> dict:

@@ -195,15 +222,23 @@ async def ensure_client_deleted(client_id: str) -> dict:

222 """

223 universal_service = UniversalService()

224 try:

198 - query = create_query("SELECT collect_client(client_id='server', artifacts=['Server.Information.Clients'], env=dict()) FROM scope()")

225 + query = create_query(

226 + "SELECT collect_client(client_id='server', artifacts=['Server.Information.Clients'], env=dict()) FROM scope()",

227 + )

228 flow = execute_query(universal_service, query)

229 flow_id = (

230 flow.get("results")[0]

202 - .get("collect_client(client_id='server', artifacts=['Server.Information.Clients'], env=dict())")

231 + .get(

232 + "collect_client(client_id='server', artifacts=['Server.Information.Clients'], env=dict())",

233 + )

234 .get("flow_id")

235 )

236

206 - results = universal_service.read_collection_results(client_id=client_id, flow_id=flow_id, artifact="Server.Information.Clients")

237 + results = universal_service.read_collection_results(

238 + client_id=client_id,

239 + flow_id=flow_id,

240 + artifact="Server.Information.Clients",

241 + )

242 return check_client_in_results(results, client_id)

243 except Exception as e:

244 return handle_exception(e, client_id)

backend/app/agents/velociraptor/utils/universal.py

+3 -1

@@ -16,5 +16,7 @@ def parse_date(date_string: str) -> datetime:

16 try:

17 return datetime.strptime(date_string, "%Y-%m-%dT%H:%M:%S+00:00")

18 except ValueError:

19 - logger.info(f"Invalid format for date: {date_string}. Using the epoch time as default.")

19 + logger.info(

20 + f"Invalid format for date: {date_string}. Using the epoch time as default.",

21 + )

22 return datetime.strptime("1970-01-01T00:00:00+00:00", "%Y-%m-%dT%H:%M:%S+00:00")

backend/app/agents/wazuh/schema/agents.py

+2 -4

@@ -1,9 +1,7 @@

1 from datetime import datetime

2 -from typing import List

3 -from typing import Optional

2 +from typing import List, Optional

5 -from pydantic import BaseModel

6 -from pydantic import Field

4 +from pydantic import BaseModel, Field

7 class WazuhAgent(BaseModel):

backend/app/agents/wazuh/services/agents.py

+24 -10

@@ -1,12 +1,12 @@

1 +from app.agents.schema.agents import AgentModifyResponse

2 +from app.agents.wazuh.schema.agents import WazuhAgent, WazuhAgentsList

3 +from app.connectors.wazuh_manager.utils.universal import (

4 + send_delete_request,

5 + send_get_request,

6 +)

7 from fastapi import HTTPException

8 from loguru import logger

4 -from app.agents.schema.agents import AgentModifyResponse

5 -from app.agents.wazuh.schema.agents import WazuhAgent

6 -from app.agents.wazuh.schema.agents import WazuhAgentsList

7 -from app.connectors.wazuh_manager.utils.universal import send_delete_request

8 -from app.connectors.wazuh_manager.utils.universal import send_get_request

9 -

11 async def collect_wazuh_agents() -> WazuhAgentsList:

12 """

@@ -16,7 +16,10 @@ async def collect_wazuh_agents() -> WazuhAgentsList:

16 WazuhAgentsList: A list of WazuhAgent objects representing the collected agents.

17 """

18 logger.info("Collecting all agents from Wazuh Manager")

19 - agents_collected = await send_get_request(endpoint="/agents", params={"limit": 1000})

19 + agents_collected = await send_get_request(

20 + endpoint="/agents",

21 + params={"limit": 1000},

22 + )

24 if agents_collected.get("success") is False:

25 raise HTTPException(

@@ -26,7 +29,11 @@ async def collect_wazuh_agents() -> WazuhAgentsList:

29 try:

30 if agents_collected.get("success"):

31 wazuh_agents_list = []

29 - for agent in agents_collected.get("data", {}).get("data", {}).get("affected_items", []):

32 + for agent in (

33 + agents_collected.get("data", {})

34 + .get("data", {})

35 + .get("affected_items", [])

36 + ):

37 os_name = agent.get("os", {}).get("name", "Unknown")

38 last_keep_alive = agent.get("lastKeepAlive", "Unknown")

39 agent_group_list = agent.get("group", [])

@@ -43,7 +50,11 @@ async def collect_wazuh_agents() -> WazuhAgentsList:

50 )

51 wazuh_agents_list.append(wazuh_agent)

46 - return WazuhAgentsList(agents=wazuh_agents_list, success=True, message="Agents collected successfully")

53 + return WazuhAgentsList(

54 + agents=wazuh_agents_list,

55 + success=True,

56 + message="Agents collected successfully",

57 + )

59 except (KeyError, IndexError, HTTPException) as e:

60 # Handle or log the error as needed

@@ -116,4 +127,7 @@ async def delete_agent_wazuh(agent_id: str) -> AgentModifyResponse:

127

128 except Exception as e:

129 # * Catch-all for other exceptions

119 - raise HTTPException(status_code=500, detail=f"Failed to delete agent {agent_id} from Wazuh Manager: {e}")

130 + raise HTTPException(

131 + status_code=500,

132 + detail=f"Failed to delete agent {agent_id} from Wazuh Manager: {e}",

133 + )

backend/app/agents/wazuh/services/vulnerabilities.py

+22 -9

@@ -1,12 +1,13 @@

1 from typing import List

3 +from app.agents.wazuh.schema.agents import (

4 + WazuhAgentVulnerabilities,

5 + WazuhAgentVulnerabilitiesResponse,

6 +)

7 +from app.connectors.wazuh_manager.utils.universal import send_get_request

8 from fastapi import HTTPException

9 from loguru import logger

6 -from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilities

7 -from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse

8 -from app.connectors.wazuh_manager.utils.universal import send_get_request

9 -

12 async def collect_agent_vulnerabilities(agent_id: str):

13 """

@@ -22,11 +23,15 @@ async def collect_agent_vulnerabilities(agent_id: str):

23 HTTPException: If there is an error collecting the vulnerabilities.

24 """

25 logger.info(f"Collecting agent {agent_id} vulnerabilities from Wazuh Manager")

25 - agent_vulnerabilities = await send_get_request(endpoint=f"/vulnerability/{agent_id}")

26 + agent_vulnerabilities = await send_get_request(

27 + endpoint=f"/vulnerability/{agent_id}",

28 + )

29 if agent_vulnerabilities["success"] is False:

30 raise HTTPException(status_code=500, detail=agent_vulnerabilities["message"])

29 - processed_vulnerabilities = process_agent_vulnerabilities(agent_vulnerabilities["data"])

32 + processed_vulnerabilities = process_agent_vulnerabilities(

33 + agent_vulnerabilities["data"],

34 + )

35 return WazuhAgentVulnerabilitiesResponse(

36 vulnerabilities=processed_vulnerabilities,

37 success=True,

@@ -34,7 +39,9 @@ async def collect_agent_vulnerabilities(agent_id: str):

39 )

37 -def process_agent_vulnerabilities(agent_vulnerabilities: dict) -> List[WazuhAgentVulnerabilities]:

42 +def process_agent_vulnerabilities(

43 + agent_vulnerabilities: dict,

44 +) -> List[WazuhAgentVulnerabilities]:

45 """

46 Process agent vulnerabilities and return a list of WazuhAgentVulnerabilities objects.

@@ -48,7 +55,13 @@ def process_agent_vulnerabilities(agent_vulnerabilities: dict) -> List[WazuhAgen

55 HTTPException: If there is an error processing the agent vulnerabilities.

56 """

57 try:

51 - vulnerabilities = agent_vulnerabilities.get("data", {}).get("affected_items", [])

58 + vulnerabilities = agent_vulnerabilities.get("data", {}).get(

59 + "affected_items",

60 + [],

61 + )

62 return [WazuhAgentVulnerabilities(**vuln) for vuln in vulnerabilities]

63 except Exception as e:

54 - raise HTTPException(status_code=500, detail=f"Failed to process agent vulnerabilities: {e}")

64 + raise HTTPException(

65 + status_code=500,

66 + detail=f"Failed to process agent vulnerabilities: {e}",

67 + )

backend/app/auth/models/users.py

+28 -11

@@ -6,12 +6,8 @@ from enum import Enum

6 from typing import Optional

8 import bcrypt

9 -from pydantic import BaseModel

10 -from pydantic import EmailStr

11 -from pydantic import validator

12 -from sqlmodel import Field

13 -from sqlmodel import Relationship

14 -from sqlmodel import SQLModel

9 +from pydantic import BaseModel, EmailStr, validator

10 +from sqlmodel import Field, Relationship, SQLModel

13 class Role(SQLModel, table=True):

@@ -50,7 +46,11 @@ class UserInput(SQLModel):

46 description="Password must be at least 8 characters long and contain at least one uppercase letter, one lowercase letter, one number, and one special character",

47 )

48 email: EmailStr

53 - role_id: RoleEnum = Field(RoleEnum.analyst, description="Role ID 1: admin, 2: analyst", foreign_key="role.id")

49 + role_id: RoleEnum = Field(

50 + RoleEnum.analyst,

51 + description="Role ID 1: admin, 2: analyst",

52 + foreign_key="role.id",

53 + )

56 class UserLogin(SQLModel):

@@ -84,7 +84,12 @@ class SMTPInput(SQLModel):

86 class Password(BaseModel):

87 - length: int = Field(default=12, ge=8, le=128, description="The length of the password")

87 + length: int = Field(

88 + default=12,

89 + ge=8,

90 + le=128,

91 + description="The length of the password",

92 + )

93 hashed: str # Holds the hashed password

94 plain: str # Holds the plain password

@@ -106,11 +111,19 @@ class Password(BaseModel):

111 punctuation = string.punctuation

112

113 # Ensure the password has at least one lowercase, one uppercase, one digit, and one symbol

109 - password_chars = [random.choice(lowercase), random.choice(uppercase), random.choice(digits), random.choice(punctuation)]

114 + password_chars = [

115 + random.choice(lowercase),

116 + random.choice(uppercase),

117 + random.choice(digits),

118 + random.choice(punctuation),

119 + ]

120

121 # Fill the rest of the password length with a random mix of characters

122 if length > 4:

113 - password_chars += random.choices(lowercase + uppercase + digits + punctuation, k=length - 4)

123 + password_chars += random.choices(

124 + lowercase + uppercase + digits + punctuation,

125 + k=length - 4,

126 + )

127

128 # Shuffle the resulting password list to avoid predictable patterns

129 random.shuffle(password_chars)

@@ -122,7 +135,11 @@ class Password(BaseModel):

135 hashed_password = bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt())

136

137 # Return the Password object with both the plain and hashed password

125 - return cls(length=length, hashed=hashed_password.decode("utf-8"), plain=password)

138 + return cls(

139 + length=length,

140 + hashed=hashed_password.decode("utf-8"),

141 + plain=password,

142 + )

143

144

145 # ! PASSWORD RESET TOKEN GENERATION NOT USING FOR NOW! #

backend/app/auth/routes/auth.py

+61 -28

@@ -1,27 +1,21 @@

1 from datetime import timedelta

3 -from fastapi import APIRouter

4 -from fastapi import Depends

5 -from fastapi import HTTPException

6 -from fastapi import Security

7 -from fastapi import status

8 -from fastapi.security import OAuth2PasswordRequestForm

9 -from loguru import logger

10 -from sqlalchemy.ext.asyncio import AsyncSession

11 -

12 -from app.auth.models.users import PasswordReset

13 -from app.auth.models.users import PasswordResetToken

14 -from app.auth.models.users import User

15 -from app.auth.models.users import UserInput

16 -from app.auth.models.users import UserLogin

17 -from app.auth.schema.auth import Token

18 -from app.auth.schema.auth import UserLoginResponse

19 -from app.auth.schema.auth import UserResponse

3 +from app.auth.models.users import (

4 + PasswordReset,

5 + PasswordResetToken,

6 + User,

7 + UserInput,

8 + UserLogin,

9 +)

10 +from app.auth.schema.auth import Token, UserLoginResponse, UserResponse

11 from app.auth.schema.user import UserBaseResponse

21 -from app.auth.services.universal import find_user

22 -from app.auth.services.universal import select_all_users

12 +from app.auth.services.universal import find_user, select_all_users

13 from app.auth.utils import AuthHandler

14 from app.db.db_session import get_db

15 +from fastapi import APIRouter, Depends, HTTPException, Security, status

16 +from fastapi.security import OAuth2PasswordRequestForm

17 +from loguru import logger

18 +from sqlalchemy.ext.asyncio import AsyncSession

20 ACCESS_TOKEN_EXPIRE_MINUTES = 1440

@@ -41,6 +35,7 @@ async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(

35 dict: A dictionary containing the access token and token type.

36 """

37 # user = auth_handler.authenticate_user(form_data.username, form_data.password)

38 +

39 user = await auth_handler.authenticate_user(form_data.username, form_data.password)

40 if not user:

41 raise HTTPException(

@@ -50,6 +45,7 @@ async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(

45 )

46 access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)

47 access_token = await auth_handler.encode_token(user.username, access_token_expires)

48 + logger.info(f"Access token: {access_token}")

49 return {"access_token": access_token, "token_type": "bearer"}

@@ -65,11 +61,19 @@ async def refresh_token(current_user: User = Depends(auth_handler.get_current_us

61 - dict: A dictionary containing the refreshed access token and token type.

62 """

63 access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)

68 - access_token = await auth_handler.encode_token(current_user.username, access_token_expires)

64 + access_token = await auth_handler.encode_token(

65 + current_user.username,

66 + access_token_expires,

67 + )

68 return {"access_token": access_token, "token_type": "bearer"}

72 -@auth_router.post("/register", response_model=UserResponse, status_code=201, description="Register new user")

71 +@auth_router.post(

72 + "/register",

73 + response_model=UserResponse,

74 + status_code=201,

75 + description="Register new user",

76 +)

77 async def register(user: UserInput, session: AsyncSession = Depends(get_db)):

78 """

79 Register a new user.

@@ -85,14 +89,24 @@ async def register(user: UserInput, session: AsyncSession = Depends(get_db)):

89 if any(x.username == user.username for x in users):

90 raise HTTPException(status_code=400, detail="Username is taken")

91 hashed_pwd = auth_handler.get_password_hash(user.password)

88 - u = User(username=user.username, password=hashed_pwd, email=user.email, role_id=user.role_id)

92 + u = User(

93 + username=user.username,

94 + password=hashed_pwd,

95 + email=user.email,

96 + role_id=user.role_id,

97 + )

98 logger.info(f"User: {u}")

99 session.add(u)

100 await session.commit()

101 return {"message": "User created successfully", "success": True}

102

103

95 -@auth_router.post("/login", response_model=UserLoginResponse, description="Login user", deprecated=True)

104 +@auth_router.post(

105 + "/login",

106 + response_model=UserLoginResponse,

107 + description="Login user",

108 + deprecated=True,

109 +)

110 async def login(user: UserLogin):

111 """

112 Logs in a user.

@@ -136,12 +150,24 @@ async def get_users(session: AsyncSession = Depends(get_db)):

150

151 """

152 users = await select_all_users()

139 - return UserBaseResponse(users=users, message="Users retrieved successfully", success=True)

153 + return UserBaseResponse(

154 + users=users,

155 + message="Users retrieved successfully",

156 + success=True,

157 + )

158

159

160 # ! TODO: HAVE LOGIC TO HANDLE PASSWORD RESET VIA A TOKEN BUT NOT IMPLEMENTED YET ! #

143 -@auth_router.post("/reset-token", status_code=200, description="Request password reset", include_in_schema=False)

144 -async def request_password_reset(password_reset_request: PasswordResetToken, session: AsyncSession = Depends(get_db)):

161 +@auth_router.post(

162 + "/reset-token",

163 + status_code=200,

164 + description="Request password reset",

165 + include_in_schema=False,

166 +)

167 +async def request_password_reset(

168 + password_reset_request: PasswordResetToken,

169 + session: AsyncSession = Depends(get_db),

170 +):

171 """

172 Request a password reset.

173

@@ -189,7 +215,10 @@ async def request_password_reset(password_reset_request: PasswordResetToken, ses

215 description="Reset user's password via username",

216 dependencies=[Security(AuthHandler().require_any_scope("admin"))],

217 )

192 -async def reset_password_via_username(request: PasswordReset, session: AsyncSession = Depends(get_db)):

218 +async def reset_password_via_username(

219 + request: PasswordReset,

220 + session: AsyncSession = Depends(get_db),

221 +):

222 """

223 Reset a user's password via the username. Must be an admin.

224

@@ -217,7 +246,11 @@ async def reset_password_via_username(request: PasswordReset, session: AsyncSess

246 description="Reset user's password",

247 dependencies=[Security(AuthHandler().require_any_scope("analyst", "admin"))],

248 )

220 -async def reset_password_me(request: PasswordReset, token: str = Depends(AuthHandler().security), session: AsyncSession = Depends(get_db)):

249 +async def reset_password_me(

250 + request: PasswordReset,

251 + token: str = Depends(AuthHandler().security),

252 + session: AsyncSession = Depends(get_db),

253 +):

254 """

255 Reset a user's password.

256

backend/app/auth/schema/user.py

+1 -2

@@ -1,7 +1,6 @@

1 from typing import List

3 -from pydantic import BaseModel

4 -from pydantic import EmailStr

3 +from pydantic import BaseModel, EmailStr

6 class UserBase(BaseModel):

backend/app/auth/services/universal.py

+16 -11

@@ -1,14 +1,11 @@

1 +from app.auth.models.users import Password, Role, User

2 +from app.db.db_session import async_engine

3 from loguru import logger

5 # ! New with Async

6 from sqlalchemy.ext.asyncio import AsyncSession

7 from sqlmodel import select

7 -from app.auth.models.users import Password

8 -from app.auth.models.users import Role

9 -from app.auth.models.users import User

10 -from app.db.db_session import async_engine

11 -

9 passwords_in_memory = {}

@@ -35,10 +32,14 @@ async def find_user(name: str):

32 Returns:

33 User: The user object if found, None otherwise.

34 """

38 - async with AsyncSession(async_engine) as session:

39 - statement = select(User).where(User.username == name)

40 - result = await session.execute(statement)

41 - return result.scalars().first()

35 + try:

36 + async with AsyncSession(async_engine) as session:

37 + statement = select(User).where(User.username == name)

38 + result = await session.execute(statement)

39 + return result.scalars().first()

40 + except Exception as e:

41 + logger.error(f"Error: {e}")

42 + return None

45 async def get_role(name: str):

@@ -103,7 +104,9 @@ async def create_admin_user(session: AsyncSession):

104 Returns:

105 - None

106 """

106 - if not await check_admin_user_exists(session): # The check function needs to be passed the session as well

107 + if not await check_admin_user_exists(

108 + session,

109 + ): # The check function needs to be passed the session as well

110 # Create the admin user

111 password_model = Password.generate(length=12)

112 admin_user = User(

@@ -133,7 +136,9 @@ async def create_scheduler_user(session: AsyncSession):

136 Returns:

137 - None

138 """

136 - if not await check_scheduler_user_exists(session): # The check function needs to be passed the session as well

139 + if not await check_scheduler_user_exists(

140 + session,

141 + ): # The check function needs to be passed the session as well

142 # Create the scheduler user

143 password_model = Password.generate(length=12)

144 scheduler_user = User(

backend/app/auth/utils.py

+36 -18

@@ -1,22 +1,21 @@

1 -from datetime import datetime

2 -from datetime import timedelta

1 +from datetime import datetime, timedelta

3 import jwt

5 -from fastapi import Depends

6 -from fastapi import HTTPException

7 -from fastapi.security import OAuth2PasswordBearer

8 -from fastapi.security import SecurityScopes

4 +from app.auth.services.universal import find_user, get_role

5 +from fastapi import Depends, HTTPException

6 +from fastapi.security import OAuth2PasswordBearer, SecurityScopes

7 from loguru import logger

8 from passlib.context import CryptContext

12 -from app.auth.services.universal import find_user

13 -from app.auth.services.universal import get_role

14 -

11 class AuthHandler:

12 security = OAuth2PasswordBearer(

18 - tokenUrl="auth/token",

19 - scopes={"admin": "Admin users", "analyst": "SOC Analysts", "scheduler": "Scheduler for automated tasks"},

13 + tokenUrl="api/auth/token",

14 + scopes={

15 + "admin": "Admin users",

16 + "analyst": "SOC Analysts",

17 + "scheduler": "Scheduler for automated tasks",

18 + },

19 )

20 pwd_context = CryptContext(schemes=["bcrypt"])

21 secret = "bL4unrkoxtFs1MT6A7Ns2yMLkduyuqrkTxDV9CjlbNc="

@@ -28,7 +27,11 @@ class AuthHandler:

27 return self.pwd_context.verify(plain_password, hashed_password)

29 # ! TODO: HAVE LOGIC TO HANDLE PASSWORD RESET VIA A TOKEN BUT NOT IMPLEMENTED YET ! #

31 - def generate_reset_token(self, username: str, expires_delta: timedelta = timedelta(minutes=30)):

30 + def generate_reset_token(

31 + self,

32 + username: str,

33 + expires_delta: timedelta = timedelta(minutes=30),

34 + ):

35 """

36 Generates a password reset token.

@@ -78,7 +81,10 @@ class AuthHandler:

81 if payload["sub"] == user.username:

82 return payload["sub"]

83 else:

81 - raise HTTPException(status_code=401, detail="Invalid token. Username does not match.")

84 + raise HTTPException(

85 + status_code=401,

86 + detail="Invalid token. Username does not match.",

87 + )

88 except jwt.ExpiredSignatureError:

89 raise HTTPException(status_code=401, detail="Token expired")

90 except jwt.InvalidTokenError:

@@ -98,13 +104,21 @@ class AuthHandler:

104 otherwise False.

105 """

106 user = await find_user(username)

101 - if not user or not self.verify_password(password, user.password):

102 - logger.info("Password is not verified")

107 + try:

108 + if not user or not self.verify_password(password, user.password):

109 + logger.info("Password is not verified")

110 + return False

111 + return user

112 + except Exception as e:

113 + logger.error(f"Error: {e}")

114 return False

104 - return user

115

116 # ! New with Async

107 - async def encode_token(self, username: str, access_token_expires: timedelta = timedelta(hours=24)):

117 + async def encode_token(

118 + self,

119 + username: str,

120 + access_token_expires: timedelta = timedelta(hours=24),

121 + ):

122 """

123 Encodes a JWT token with the given username and expiration time.

124

@@ -147,7 +161,11 @@ class AuthHandler:

161 except jwt.InvalidTokenError:

162 return "Invalid token", []

163

150 - async def get_current_user(self, security_scopes: SecurityScopes, token: str = Depends(security)):

164 + async def get_current_user(

165 + self,

166 + security_scopes: SecurityScopes,

167 + token: str = Depends(security),

168 + ):

169 """

170 Retrieves the current user based on the provided security scopes and token.

171

backend/app/connectors/cortex/routes/analyzers.py

+21 -15

@@ -1,17 +1,14 @@

1 from typing import List

3 -from fastapi import APIRouter

4 -from fastapi import Depends

5 -from fastapi import HTTPException

6 -from fastapi import Security

7 -from loguru import logger

8 -

3 from app.auth.utils import AuthHandler

10 -from app.connectors.cortex.schema.analyzers import AnalyzersResponse

11 -from app.connectors.cortex.schema.analyzers import RunAnalyzerBody

12 -from app.connectors.cortex.schema.analyzers import RunAnalyzerResponse

13 -from app.connectors.cortex.services.analyzers import get_analyzers

14 -from app.connectors.cortex.services.analyzers import run_analyzer

4 +from app.connectors.cortex.schema.analyzers import (

5 + AnalyzersResponse,

6 + RunAnalyzerBody,

7 + RunAnalyzerResponse,

8 +)

9 +from app.connectors.cortex.services.analyzers import get_analyzers, run_analyzer

10 +from fastapi import APIRouter, Depends, HTTPException, Security

11 +from loguru import logger

13 # App specific imports

@@ -45,7 +42,10 @@ async def verify_analyzer_exists(run_analyzer_body: RunAnalyzerBody) -> RunAnaly

42 """

43 available_analyzers = await get_available_analyzers()

44 if run_analyzer_body.analyzer_name not in available_analyzers:

48 - raise HTTPException(status_code=400, detail=f"Analyzer {run_analyzer_body.analyzer_name} does not exist.")

45 + raise HTTPException(

46 + status_code=400,

47 + detail=f"Analyzer {run_analyzer_body.analyzer_name} does not exist.",

48 + )

49 return run_analyzer_body

@@ -72,7 +72,9 @@ async def get_all_analyzers() -> AnalyzersResponse:

72 description="Run an analyzer",

73 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

74 )

75 -async def run_analyzer_route(run_analyzer_body: RunAnalyzerBody = Depends(verify_analyzer_exists)) -> RunAnalyzerResponse:

75 +async def run_analyzer_route(

76 + run_analyzer_body: RunAnalyzerBody = Depends(verify_analyzer_exists),

77 +) -> RunAnalyzerResponse:

78 """

79 Run an analyzer.

@@ -82,9 +84,13 @@ async def run_analyzer_route(run_analyzer_body: RunAnalyzerBody = Depends(verify

84 Returns:

85 RunAnalyzerResponse: The response containing the result of running the analyzer.

86 """

85 - is_valid, data_type = RunAnalyzerBody.is_valid_datatype(run_analyzer_body.analyzer_data)

87 + is_valid, data_type = RunAnalyzerBody.is_valid_datatype(

88 + run_analyzer_body.analyzer_data,

89 + )

90 if not is_valid:

91 raise HTTPException(status_code=400, detail=f"Invalid data type: {data_type}")

89 - logger.info(f"Running analyzer {run_analyzer_body.analyzer_name} with data {run_analyzer_body.analyzer_data} of type {data_type}")

93 + logger.info(

94 + f"Running analyzer {run_analyzer_body.analyzer_name} with data {run_analyzer_body.analyzer_data} of type {data_type}",

95 + )

96 return await run_analyzer(run_analyzer_body, data_type)

backend/app/connectors/cortex/schema/analyzers.py

+25 -14

@@ -1,16 +1,12 @@

1 import ipaddress

2 import re

3 -from typing import Any

4 -from typing import Dict

5 -from typing import List

6 -from typing import Optional

7 -from typing import Tuple

3 +from typing import Any, Dict, List, Optional, Tuple

9 -from pydantic import BaseModel

10 -from pydantic import Field

11 -from pydantic import validator

5 +from pydantic import BaseModel, Field, validator

13 -HASH_REGEX = re.compile(r"[a-fA-F\d]{32}|[a-fA-F\d]{64}") # Update this regex to match your specific hash format

7 +HASH_REGEX = re.compile(

8 + r"[a-fA-F\d]{32}|[a-fA-F\d]{64}",

9 +) # Update this regex to match your specific hash format

10 DOMAIN_REGEX = re.compile(

11 r"^(?:[a-z0-9](?:[a-z0-9\-]{0,61}[a-z0-9])?\.)+[a-z]{2,6}$",

12 ) # Update this regex to match your specific domain format

@@ -24,8 +20,14 @@ class AnalyzersResponse(BaseModel):

21 class RunAnalyzerBody(BaseModel):

22 analyzer_name: str = Field(..., description="Name of the analyzer to be run.")

27 - analyzer_data: str = Field(..., description="The Indicator of Compromise (IoC) to be analyzed.")

28 - data_type: Optional[str] = Field(default=None, description="Data type determined after validation")

23 + analyzer_data: str = Field(

24 + ...,

25 + description="The Indicator of Compromise (IoC) to be analyzed.",

26 + )

27 + data_type: Optional[str] = Field(

28 + default=None,

29 + description="Data type determined after validation",

30 + )

32 @validator("analyzer_data", pre=True, always=True)

33 def validate_and_set_data_type(cls, value: str, values: dict) -> str:

@@ -70,7 +72,16 @@ class RunAnalyzerResponse(BaseModel):

74 class AnalyzerJobData(BaseModel):

73 - data: str = Field(..., description="The Indicator of Compromise (IoC) to be analyzed.")

74 - dataType: str = Field(..., description="The type of the IoC (e.g., 'IP', 'hash', 'domain').")

75 + data: str = Field(

76 + ...,

77 + description="The Indicator of Compromise (IoC) to be analyzed.",

78 + )

79 + dataType: str = Field(

80 + ...,

81 + description="The type of the IoC (e.g., 'IP', 'hash', 'domain').",

82 + )

83 tlp: int = Field(1, description="Traffic Light Protocol (TLP) level.")

76 - message: str = Field("custom message sent to analyzer", description="Custom message.")

84 + message: str = Field(

85 + "custom message sent to analyzer",

86 + description="Custom message.",

87 + )

backend/app/connectors/cortex/services/analyzers.py

+37 -21

@@ -1,22 +1,19 @@

1 -from typing import Dict

2 -from typing import List

3 -from typing import Union

1 +from typing import Dict, List, Union

3 +from app.connectors.cortex.schema.analyzers import (

4 + AnalyzerJobData,

5 + AnalyzersResponse,

6 + RunAnalyzerBody,

7 + RunAnalyzerResponse,

8 +)

9 +from app.connectors.cortex.utils.universal import ( # Importing create_cortex_client; Importing from universal.py

10 + create_cortex_client,

11 + run_and_wait_for_analyzer,

12 +)

13 from cortex4py.api import Api

14 from fastapi import HTTPException

15 from loguru import logger

9 -from app.connectors.cortex.schema.analyzers import AnalyzerJobData

10 -from app.connectors.cortex.schema.analyzers import AnalyzersResponse

11 -from app.connectors.cortex.schema.analyzers import RunAnalyzerBody

12 -from app.connectors.cortex.schema.analyzers import RunAnalyzerResponse

13 -from app.connectors.cortex.utils.universal import (

14 - create_cortex_client, # Importing create_cortex_client

15 -)

16 -from app.connectors.cortex.utils.universal import (

17 - run_and_wait_for_analyzer, # Importing from universal.py

18 -)

19 -

18 async def fetch_analyzers(api: Api) -> List[Dict]:

19 """

@@ -97,10 +94,17 @@ async def get_analyzers() -> AnalyzersResponse:

94 analyzers = await fetch_analyzers(api)

95 analyzer_names = extract_analyzer_names(analyzers)

100 - return AnalyzersResponse(success=True, message="Successfully fetched analyzers", analyzers=analyzer_names)

97 + return AnalyzersResponse(

98 + success=True,

99 + message="Successfully fetched analyzers",

100 + analyzers=analyzer_names,

101 + )

102

103

103 -async def run_analyzer(run_analyzer_body: RunAnalyzerBody, data_type: str) -> RunAnalyzerResponse:

104 +async def run_analyzer(

105 + run_analyzer_body: RunAnalyzerBody,

106 + data_type: str,

107 +) -> RunAnalyzerResponse:

108 """

109 Runs an analyzer with the given analyzer name, analyzer data, and data type.

110

@@ -116,13 +120,25 @@ async def run_analyzer(run_analyzer_body: RunAnalyzerBody, data_type: str) -> Ru

120

121 analyzer_name = run_analyzer_body.analyzer_name

122 analyzer_data = run_analyzer_body.analyzer_data

119 - logger.info(f"Running analyzer {analyzer_name} with data {analyzer_data} of type {data_type}")

123 + logger.info(

124 + f"Running analyzer {analyzer_name} with data {analyzer_data} of type {data_type}",

125 + )

126 job_data = AnalyzerJobData(data=analyzer_data, dataType=data_type)

127

122 - result = await run_and_wait_for_analyzer(analyzer_name=analyzer_name, job_data=job_data)

128 + result = await run_and_wait_for_analyzer(

129 + analyzer_name=analyzer_name,

130 + job_data=job_data,

131 + )

132

133 if result is None:

134 logger.error(f"Failed to run analyzer {analyzer_name}")

126 - raise HTTPException(status_code=500, detail=f"Failed to run analyzer {analyzer_name}")

127 -

128 - return RunAnalyzerResponse(success=True, message="Successfully ran analyzer", report=result)

135 + raise HTTPException(

136 + status_code=500,

137 + detail=f"Failed to run analyzer {analyzer_name}",

138 + )

139 +

140 + return RunAnalyzerResponse(

141 + success=True,

142 + message="Successfully ran analyzer",

143 + report=result,

144 + )

backend/app/connectors/cortex/utils/universal.py

+45 -16

@@ -1,14 +1,12 @@

1 import time

2 -from typing import Any

3 -from typing import Dict

4 -

5 -from cortex4py.api import Api

6 -from fastapi import HTTPException

7 -from loguru import logger

2 +from typing import Any, Dict

4 from app.connectors.cortex.schema.analyzers import AnalyzerJobData

5 from app.connectors.utils import get_connector_info_from_db

6 from app.db.db_session import get_db_session

7 +from cortex4py.api import Api

8 +from fastapi import HTTPException

9 +from loguru import logger

12 async def verify_cortex_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]:

@@ -21,18 +19,35 @@ async def verify_cortex_credentials(attributes: Dict[str, Any]) -> Dict[str, Any

19 logger.info(f"Verifying the Cortex connection to {attributes['connector_url']}")

21 try:

24 - api = Api(attributes["connector_url"], attributes["connector_api_key"], verify_cert=False)

22 + api = Api(

23 + attributes["connector_url"],

24 + attributes["connector_api_key"],

25 + verify_cert=False,

26 + )

27 # Get Cortex Status

28 status = api.status

29 if status:

30 logger.debug("Cortex connection successful")

29 - return {"connectionSuccessful": True, "message": "Cortex connection successful"}

31 + return {

32 + "connectionSuccessful": True,

33 + "message": "Cortex connection successful",

34 + }

35 else:

31 - logger.error(f"Connection to {attributes['connector_url']} failed with error.")

32 - return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error."}

36 + logger.error(

37 + f"Connection to {attributes['connector_url']} failed with error.",

38 + )

39 + return {

40 + "connectionSuccessful": False,

41 + "message": f"Connection to {attributes['connector_url']} failed with error.",

42 + }

43 except Exception as e:

34 - logger.error(f"Connection to {attributes['connector_url']} failed with error: {e}")

35 - return {"connectionSuccessful": False, "message": f"Connection to {attributes['connector_url']} failed with error: {e}"}

44 + logger.error(

45 + f"Connection to {attributes['connector_url']} failed with error: {e}",

46 + )

47 + return {

48 + "connectionSuccessful": False,

49 + "message": f"Connection to {attributes['connector_url']} failed with error: {e}",

50 + }

53 async def verify_cortex_connection(connector_name: str) -> str:

@@ -62,10 +77,17 @@ async def create_cortex_client(connector_name: str) -> Api:

77 if attributes is None:

78 logger.error("No Wazuh Indexer connector found in the database")

79 return None

65 - return Api(attributes["connector_url"], attributes["connector_api_key"], verify_cert=False)

80 + return Api(

81 + attributes["connector_url"],

82 + attributes["connector_api_key"],

83 + verify_cert=False,

84 + )

68 -async def run_and_wait_for_analyzer(analyzer_name: str, job_data: AnalyzerJobData) -> Dict[str, Any]:

87 +async def run_and_wait_for_analyzer(

88 + analyzer_name: str,

89 + job_data: AnalyzerJobData,

90 +) -> Dict[str, Any]:

91 """

92 Runs an analyzer by name and waits for the job to complete.

@@ -83,7 +105,10 @@ async def run_and_wait_for_analyzer(analyzer_name: str, job_data: AnalyzerJobDat

105 job = api.analyzers.run_by_name(analyzer_name, job_data.dict(), force=1)

106 return await monitor_analyzer_job(api, job)

107 except Exception as e:

86 - raise HTTPException(status_code=500, detail=f"Error running analyzer {analyzer_name}: {e}")

108 + raise HTTPException(

109 + status_code=500,

110 + detail=f"Error running analyzer {analyzer_name}: {e}",

111 + )

112

113

114 async def monitor_analyzer_job(api: Api, job: Any) -> Dict[str, Any]:

@@ -138,4 +163,8 @@ async def retrieve_final_report(api: Api, job_id: str) -> Dict[str, Any]:

163 """

164 report = api.jobs.get_report(job_id).report

165 final_report = report["full"]

141 - return {"success": True, "message": "Analyzer ran successfully", "report": final_report}

166 + return {

167 + "success": True,

168 + "message": "Analyzer ran successfully",

169 + "report": final_report,

170 + }

backend/app/connectors/dfir_iris/routes/alerts.py

+57 -31

@@ -1,26 +1,26 @@

1 -from fastapi import APIRouter

2 -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 -

1 from app.auth.utils import AuthHandler

9 -from app.connectors.dfir_iris.schema.alerts import AlertResponse

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

13 -from app.connectors.dfir_iris.schema.alerts import DeleteAlertResponse

14 -from app.connectors.dfir_iris.schema.alerts import DeleteMultipleAlertsRequest

15 -from app.connectors.dfir_iris.schema.alerts import FilterAlertsRequest

16 -from app.connectors.dfir_iris.services.alerts import bookmark_alert

17 -from app.connectors.dfir_iris.services.alerts import create_case

18 -from app.connectors.dfir_iris.services.alerts import delete_alert

19 -from app.connectors.dfir_iris.services.alerts import get_alert

20 -from app.connectors.dfir_iris.services.alerts import get_alerts

21 -from app.connectors.dfir_iris.services.alerts import get_bookmarked_alerts

2 +from app.connectors.dfir_iris.schema.alerts import (

3 + AlertResponse,

4 + AlertsResponse,

5 + BookmarkedAlertsResponse,

6 + CaseCreationResponse,

7 + DeleteAlertResponse,

8 + DeleteMultipleAlertsRequest,

9 + FilterAlertsRequest,

10 +)

11 +from app.connectors.dfir_iris.services.alerts import (

12 + bookmark_alert,

13 + create_case,

14 + delete_alert,

15 + get_alert,

16 + get_alerts,

17 + get_bookmarked_alerts,

18 +)

19 from app.connectors.dfir_iris.utils.universal import check_alert_exists

20 from app.db.db_session import get_db

21 +from fastapi import APIRouter, Depends, HTTPException, Security

22 +from loguru import logger

23 +from sqlalchemy.ext.asyncio import AsyncSession

25 # App specific imports

@@ -53,7 +53,9 @@ dfir_iris_alerts_router = APIRouter()

53 description="Get all bookmarked alerts",

54 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

55 )

56 -async def get_all_bookmarked_alerts() -> BookmarkedAlertsResponse:

56 +async def get_all_bookmarked_alerts(

57 + session: AsyncSession = Depends(get_db),

58 +) -> BookmarkedAlertsResponse:

59 """

60 Fetches all bookmarked alerts.

@@ -61,7 +63,7 @@ async def get_all_bookmarked_alerts() -> BookmarkedAlertsResponse:

63 BookmarkedAlertsResponse: The response containing the bookmarked alerts.

64 """

65 logger.info("Fetching all bookmarked alerts")

64 - return await get_bookmarked_alerts()

66 + return await get_bookmarked_alerts(session=session)

69 @dfir_iris_alerts_router.post(

@@ -70,7 +72,10 @@ async def get_all_bookmarked_alerts() -> BookmarkedAlertsResponse:

72 description="Get alerts from IRIS based on the provided filters",

73 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

74 )

73 -async def get_alerts_filtered(request: FilterAlertsRequest, session: AsyncSession = Depends(get_db)) -> AlertsResponse:

75 +async def get_alerts_filtered(

76 + request: FilterAlertsRequest,

77 + session: AsyncSession = Depends(get_db),

78 +) -> AlertsResponse:

79 """

80 Retrieve alerts from DFIR-IRIS based on the provided filters.

@@ -88,7 +93,10 @@ async def get_alerts_filtered(request: FilterAlertsRequest, session: AsyncSessio

93 description="Get an alert by ID",

94 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

95 )

91 -async def get_alert_by_id(alert_id: str = Depends(verify_alert_exists), session: AsyncSession = Depends(get_db)) -> AlertResponse:

96 +async def get_alert_by_id(

97 + alert_id: str = Depends(verify_alert_exists),

98 + session: AsyncSession = Depends(get_db),

99 +) -> AlertResponse:

100 """

101 Retrieve an alert by its ID.

102

@@ -119,13 +127,21 @@ async def get_all_alerts_assigned_to_user(user_id: int) -> AlertsResponse:

127 AlertsResponse: The response containing the fetched alerts assigned to the user.

128 """

129 logger.info(f"Fetching all alerts assigned to user {user_id}")

122 - alerts = (await get_alerts(request=FilterAlertsRequest(alert_owner_id=user_id, per_page=1000))).alerts

130 + alerts = (

131 + await get_alerts(

132 + request=FilterAlertsRequest(alert_owner_id=user_id, per_page=1000),

133 + )

134 + ).alerts

135 alerts_assigned_to_user = []

136 for alert in alerts:

137 if alert["alert_owner_id"] == user_id:

138 alerts_assigned_to_user.append(alert)

139

128 - return AlertsResponse(success=True, message="Successfully fetched alerts assigned to user", alerts=alerts_assigned_to_user)

140 + return AlertsResponse(

141 + success=True,

142 + message="Successfully fetched alerts assigned to user",

143 + alerts=alerts_assigned_to_user,

144 + )

145

146

147 @dfir_iris_alerts_router.post(

@@ -134,7 +150,9 @@ async def get_all_alerts_assigned_to_user(user_id: int) -> AlertsResponse:

150 description="Assign an alert to a user",

151 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

152 )

137 -async def create_case_from_alert(alert_id: str = Depends(verify_alert_exists)) -> CaseCreationResponse:

153 +async def create_case_from_alert(

154 + alert_id: str = Depends(verify_alert_exists),

155 +) -> CaseCreationResponse:

156 """

157 Create a case from an alert.

158

@@ -154,7 +172,9 @@ async def create_case_from_alert(alert_id: str = Depends(verify_alert_exists)) -

172 description="Bookmark an alert",

173 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

174 )

157 -async def bookmark_alert_route(alert_id: str = Depends(verify_alert_exists)) -> AlertResponse:

175 +async def bookmark_alert_route(

176 + alert_id: str = Depends(verify_alert_exists),

177 +) -> AlertResponse:

178 """

179 Bookmark an alert.

180

@@ -174,7 +194,9 @@ async def bookmark_alert_route(alert_id: str = Depends(verify_alert_exists)) ->

194 description="Unbookmark an alert",

195 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

196 )

177 -async def unbookmark_alert_route(alert_id: str = Depends(verify_alert_exists)) -> AlertResponse:

197 +async def unbookmark_alert_route(

198 + alert_id: str = Depends(verify_alert_exists),

199 +) -> AlertResponse:

200 """

201 Unbookmark an alert.

202

@@ -194,7 +216,9 @@ async def unbookmark_alert_route(alert_id: str = Depends(verify_alert_exists)) -

216 description="Delete multiple alerts",

217 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

218 )

197 -async def delete_multiple_alerts_route(request: DeleteMultipleAlertsRequest) -> DeleteAlertResponse:

219 +async def delete_multiple_alerts_route(

220 + request: DeleteMultipleAlertsRequest,

221 +) -> DeleteAlertResponse:

222 """

223 Delete multiple alerts.

224

@@ -237,7 +261,9 @@ async def purge_alerts_route() -> DeleteAlertResponse:

261 description="Delete an alert",

262 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

263 )

240 -async def delete_alert_route(alert_id: str = Depends(verify_alert_exists)) -> DeleteAlertResponse:

264 +async def delete_alert_route(

265 + alert_id: str = Depends(verify_alert_exists),

266 +) -> DeleteAlertResponse:

267 """

268 Delete an alert.

269

backend/app/connectors/dfir_iris/routes/assets.py

+5 -7

@@ -1,13 +1,9 @@

1 -from fastapi import APIRouter

2 -from fastapi import Depends

3 -from fastapi import HTTPException

4 -from fastapi import Security

5 -from loguru import logger

6 -

1 from app.auth.utils import AuthHandler

2 from app.connectors.dfir_iris.schema.assets import AssetResponse

3 from app.connectors.dfir_iris.services.assets import get_case_assets

4 from app.connectors.dfir_iris.utils.universal import check_case_exists

5 +from fastapi import APIRouter, Depends, HTTPException, Security

6 +from loguru import logger

8 # App specific imports

@@ -39,7 +35,9 @@ dfir_iris_assets_router = APIRouter()

35 description="Get all assets for a case",

36 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

37 )

42 -async def get_case_assets_route(case_id: int = Depends(verify_case_exists)) -> AssetResponse:

38 +async def get_case_assets_route(

39 + case_id: int = Depends(verify_case_exists),

40 +) -> AssetResponse:

41 """

42 Retrieve all assets for a given case.

backend/app/connectors/dfir_iris/routes/cases.py

+42 -29

@@ -1,31 +1,31 @@

1 from datetime import timedelta

3 -from fastapi import APIRouter

4 -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 -

3 from app.auth.utils import AuthHandler

11 -from app.connectors.dfir_iris.schema.cases import CaseOlderThanBody

12 -from app.connectors.dfir_iris.schema.cases import CaseResponse

13 -from app.connectors.dfir_iris.schema.cases import CasesBreachedResponse

14 -from app.connectors.dfir_iris.schema.cases import ClosedCaseResponse

15 -from app.connectors.dfir_iris.schema.cases import PurgeCaseResponse

16 -from app.connectors.dfir_iris.schema.cases import ReopenedCaseResponse

17 -from app.connectors.dfir_iris.schema.cases import SingleCaseBody

18 -from app.connectors.dfir_iris.schema.cases import SingleCaseResponse

19 -from app.connectors.dfir_iris.schema.cases import TimeUnit

20 -from app.connectors.dfir_iris.services.cases import close_case

21 -from app.connectors.dfir_iris.services.cases import delete_single_case

22 -from app.connectors.dfir_iris.services.cases import get_all_cases

23 -from app.connectors.dfir_iris.services.cases import get_cases_older_than

24 -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

4 +from app.connectors.dfir_iris.schema.cases import (

5 + CaseOlderThanBody,

6 + CaseResponse,

7 + CasesBreachedResponse,

8 + ClosedCaseResponse,

9 + PurgeCaseResponse,

10 + ReopenedCaseResponse,

11 + SingleCaseBody,

12 + SingleCaseResponse,

13 + TimeUnit,

14 +)

15 +from app.connectors.dfir_iris.services.cases import (

16 + close_case,

17 + delete_single_case,

18 + get_all_cases,

19 + get_cases_older_than,

20 + get_single_case,

21 + purge_cases,

22 + reopen_case,

23 +)

24 from app.connectors.dfir_iris.utils.universal import check_case_exists

25 from app.db.db_session import get_db

26 +from fastapi import APIRouter, Depends, HTTPException, Security

27 +from loguru import logger

28 +from sqlalchemy.ext.asyncio import AsyncSession

31 async def verify_case_exists(case_id: int) -> int:

@@ -94,7 +94,9 @@ async def get_cases_route(session: AsyncSession = Depends(get_db)) -> CaseRespon

94 description="Get all cases older than a specified date",

95 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

96 )

97 -async def get_cases_older_than_route(case_older_than_body: CaseOlderThanBody = Depends(get_timedelta)) -> CaseResponse:

97 +async def get_cases_older_than_route(

98 + case_older_than_body: CaseOlderThanBody = Depends(get_timedelta),

99 +) -> CaseResponse:

100 """

101 Fetches all cases older than a specified date.

102

@@ -104,7 +106,9 @@ async def get_cases_older_than_route(case_older_than_body: CaseOlderThanBody = D

106 Returns:

107 CaseResponse: The response containing the cases older than the specified date.

108 """

107 - logger.info(f"Fetching all cases older than {case_older_than_body.older_than} ({case_older_than_body.time_unit.value})")

109 + logger.info(

110 + f"Fetching all cases older than {case_older_than_body.older_than} ({case_older_than_body.time_unit.value})",

111 + )

112 return await get_cases_older_than(case_older_than_body)

113

114

@@ -131,7 +135,9 @@ async def purge_cases_route() -> PurgeCaseResponse:

135 description="Purge a single case",

136 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

137 )

134 -async def purge_single_case_route(case_id: int = Depends(verify_case_exists)) -> PurgeCaseResponse:

138 +async def purge_single_case_route(

139 + case_id: int = Depends(verify_case_exists),

140 +) -> PurgeCaseResponse:

141 """

142 Purge a single case by its ID.

143

@@ -152,7 +158,10 @@ async def purge_single_case_route(case_id: int = Depends(verify_case_exists)) ->

158 description="Get a single case",

159 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

160 )

155 -async def get_single_case_route(case_id: int = Depends(verify_case_exists), session: AsyncSession = Depends(get_db)) -> SingleCaseResponse:

161 +async def get_single_case_route(

162 + case_id: int = Depends(verify_case_exists),

163 + session: AsyncSession = Depends(get_db),

164 +) -> SingleCaseResponse:

165 """

166 Retrieve a single case by its ID.

167

@@ -173,7 +182,9 @@ async def get_single_case_route(case_id: int = Depends(verify_case_exists), sess

182 description="Close a single case",

183 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

184 )

176 -async def close_single_case_route(case_id: int = Depends(verify_case_exists)) -> ClosedCaseResponse:

185 +async def close_single_case_route(

186 + case_id: int = Depends(verify_case_exists),

187 +) -> ClosedCaseResponse:

188 """

189 Close a single case by its ID.

190

@@ -194,7 +205,9 @@ async def close_single_case_route(case_id: int = Depends(verify_case_exists)) ->

205 description="Open a single case",

206 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

207 )

197 -async def reopen_single_case_route(case_id: int = Depends(verify_case_exists)) -> ReopenedCaseResponse:

208 +async def reopen_single_case_route(

209 + case_id: int = Depends(verify_case_exists),

210 +) -> ReopenedCaseResponse:

211 """

212 Open a single case by its ID.

213

backend/app/connectors/dfir_iris/routes/notes.py

+16 -13

@@ -1,18 +1,15 @@

1 from typing import Optional

3 -from fastapi import APIRouter

4 -from fastapi import Depends

5 -from fastapi import HTTPException

6 -from fastapi import Security

7 -from loguru import logger

8 -

3 from app.auth.utils import AuthHandler

10 -from app.connectors.dfir_iris.schema.notes import NoteCreationBody

11 -from app.connectors.dfir_iris.schema.notes import NoteCreationResponse

12 -from app.connectors.dfir_iris.schema.notes import NotesResponse

13 -from app.connectors.dfir_iris.services.notes import create_case_note

14 -from app.connectors.dfir_iris.services.notes import get_case_notes

4 +from app.connectors.dfir_iris.schema.notes import (

5 + NoteCreationBody,

6 + NoteCreationResponse,

7 + NotesResponse,

8 +)

9 +from app.connectors.dfir_iris.services.notes import create_case_note, get_case_notes

10 from app.connectors.dfir_iris.utils.universal import check_case_exists

11 +from fastapi import APIRouter, Depends, HTTPException, Security

12 +from loguru import logger

15 async def verify_case_exists(case_id: int) -> int:

@@ -42,7 +39,10 @@ dfir_iris_notes_router = APIRouter()

39 description="Get all notes for a case",

40 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

41 )

45 -async def get_case_notes_route(case_id: int = Depends(verify_case_exists), search_term: Optional[str] = "%") -> NotesResponse:

42 +async def get_case_notes_route(

43 + case_id: int = Depends(verify_case_exists),

44 + search_term: Optional[str] = "%",

45 +) -> NotesResponse:

46 """

47 Retrieve all notes for a specific case.

@@ -63,7 +63,10 @@ async def get_case_notes_route(case_id: int = Depends(verify_case_exists), searc

63 description="Create a note for a case",

64 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

65 )

66 -async def create_case_note_route(case_id: int, note_creation_body: NoteCreationBody) -> NoteCreationResponse:

66 +async def create_case_note_route(

67 + case_id: int,

68 + note_creation_body: NoteCreationBody,

69 +) -> NoteCreationResponse:

70 """

71 Create a note for a case.

backend/app/connectors/dfir_iris/routes/users.py

+20 -15

@@ -1,18 +1,17 @@

1 -from fastapi import APIRouter

2 -from fastapi import Depends

3 -from fastapi import HTTPException

4 -from fastapi import Security

5 -from loguru import logger

6 -

1 from app.auth.utils import AuthHandler

2 from app.connectors.dfir_iris.schema.alerts import AlertResponse

9 -from app.connectors.dfir_iris.schema.users import User

10 -from app.connectors.dfir_iris.schema.users import UsersResponse

11 -from app.connectors.dfir_iris.services.users import assign_user_to_alert

12 -from app.connectors.dfir_iris.services.users import delete_user_from_alert

13 -from app.connectors.dfir_iris.services.users import get_users

14 -from app.connectors.dfir_iris.utils.universal import check_alert_exists

15 -from app.connectors.dfir_iris.utils.universal import check_user_exists

3 +from app.connectors.dfir_iris.schema.users import User, UsersResponse

4 +from app.connectors.dfir_iris.services.users import (

5 + assign_user_to_alert,

6 + delete_user_from_alert,

7 + get_users,

8 +)

9 +from app.connectors.dfir_iris.utils.universal import (

10 + check_alert_exists,

11 + check_user_exists,

12 +)

13 +from fastapi import APIRouter, Depends, HTTPException, Security

14 +from loguru import logger

17 def verify_user_exists(user_id: int) -> int:

@@ -77,7 +76,10 @@ async def get_all_users() -> UsersResponse:

76 description="Assign a user to an alert",

77 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

78 )

80 -async def assign_user_to_alert_route(alert_id: str = Depends(verify_alert_exists), user_id: int = Depends(verify_user_exists)) -> User:

79 +async def assign_user_to_alert_route(

80 + alert_id: str = Depends(verify_alert_exists),

81 + user_id: int = Depends(verify_user_exists),

82 +) -> User:

83 """

84 Assigns a user to an alert.

@@ -101,7 +103,10 @@ async def assign_user_to_alert_route(alert_id: str = Depends(verify_alert_exists

103 description="Delete a user from an alert",

104 dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],

105 )

104 -async def delete_user_from_alert_route(alert_id: str = Depends(verify_alert_exists), user_id: int = Depends(verify_user_exists)) -> User:

106 +async def delete_user_from_alert_route(

107 + alert_id: str = Depends(verify_alert_exists),

108 + user_id: int = Depends(verify_user_exists),

109 +) -> User:

110 """

111 Delete a user from an alert.

112

backend/app/connectors/dfir_iris/schema/admin.py

+2 -5

@@ -1,11 +1,8 @@

1 import uuid

2 from datetime import datetime

3 -from typing import Dict

4 -from typing import List

5 -from typing import Optional

3 +from typing import Dict, List, Optional

7 -from pydantic import UUID4

8 -from pydantic import BaseModel

5 +from pydantic import UUID4, BaseModel

8 class CreateCustomerData(BaseModel):

backend/app/connectors/dfir_iris/schema/alerts.py

+22 -11

@@ -1,27 +1,32 @@

1 from enum import Enum

2 -from typing import Any

3 -from typing import Dict

4 -from typing import List

5 -from typing import Optional

2 +from typing import Any, Dict, List, Optional

7 -from pydantic import BaseModel

8 -from pydantic import Field

4 +from pydantic import BaseModel, Field

7 class AlertsResponse(BaseModel):

12 - alerts: Optional[List[Dict[str, Any]]] = Field([], description="The alerts returned from the search.")

8 + alerts: Optional[List[Dict[str, Any]]] = Field(

9 + [],

10 + description="The alerts returned from the search.",

11 + )

12 message: str

13 success: bool

16 class AlertResponse(BaseModel):

18 - alert: Optional[Dict[str, Any]] = Field({}, description="The alert returned from the search.")

17 + alert: Optional[Dict[str, Any]] = Field(

18 + {},

19 + description="The alert returned from the search.",

20 + )

21 message: str

22 success: bool

25 class BookmarkedAlertsResponse(BaseModel):

24 - bookmarked_alerts: Optional[List[Dict[str, Any]]] = Field([], description="The alerts returned from the search.")

26 + bookmarked_alerts: Optional[List[Dict[str, Any]]] = Field(

27 + [],

28 + description="The alerts returned from the search.",

29 + )

30 message: str

31 success: bool

@@ -43,9 +48,15 @@ class SortOrder(Enum):

48 class FilterAlertsRequest(BaseModel):

49 per_page: int = Field(1000, description="The number of alerts to return per page.")

50 page: int = Field(1, description="The page number to return.")

46 - sort: SortOrder = Field(SortOrder.desc, description="The sort order for the alerts.")

51 + sort: SortOrder = Field(

52 + SortOrder.desc,

53 + description="The sort order for the alerts.",

54 + )

55 alert_title: Optional[str] = Field(None, description="The title of the alert.")

48 - alert_owner_id: Optional[int] = Field(None, description="The ID of the alert owner.")

56 + alert_owner_id: Optional[int] = Field(

57 + None,

58 + description="The ID of the alert owner.",

59 + )

62 class CaseModificationHistory(BaseModel):

backend/app/connectors/dfir_iris/schema/assets.py

+6 -5

@@ -1,8 +1,6 @@

1 -from typing import List

2 -from typing import Optional

1 +from typing import List, Optional

4 -from pydantic import BaseModel

5 -from pydantic import Field

3 +from pydantic import BaseModel, Field

6 class AssetState(BaseModel):

@@ -20,7 +18,10 @@ class Asset(BaseModel):

18 analysis_status: str

19 analysis_status_id: int

20 asset_compromise_status_id: Optional[int]

23 - asset_description: Optional[str] = Field(None, description="The description of the asset.")

21 + asset_description: Optional[str] = Field(

22 + None,

23 + description="The description of the asset.",

24 + )

25 asset_domain: Optional[str]

26 asset_icon_compromised: str

27 asset_icon_not_compromised: str

backend/app/connectors/dfir_iris/schema/cases.py

+3 -8

@@ -1,13 +1,8 @@

1 -from datetime import date

2 -from datetime import timedelta

1 +from datetime import date, timedelta

2 from enum import Enum

4 -from typing import Dict

5 -from typing import List

6 -from typing import Optional

7 -from typing import Union

3 +from typing import Dict, List, Optional, Union

9 -from pydantic import BaseModel

10 -from pydantic import Field

5 +from pydantic import BaseModel, Field

8 class CaseModel(BaseModel):

backend/app/connectors/dfir_iris/schema/notes.py

+6 -6

@@ -1,9 +1,6 @@

1 -from typing import Dict

2 -from typing import List

3 -from typing import Optional

1 +from typing import Dict, List, Optional

5 -from pydantic import BaseModel

6 -from pydantic import Field

3 +from pydantic import BaseModel, Field

6 class CustomAttributes(BaseModel):

@@ -44,7 +41,10 @@ class NotesResponse(BaseModel):

42 class NotesQueryParams(BaseModel):

43 case_id: int

47 - search_term: Optional[str] = Field("%", description="Search term to filter notes by. Defaults to wildcard search (%).")

44 + search_term: Optional[str] = Field(

45 + "%",

46 + description="Search term to filter notes by. Defaults to wildcard search (%).",

47 + )

50 class NoteCreationBody(BaseModel):

backend/app/connectors/dfir_iris/services/alerts.py

+101 -34

@@ -1,20 +1,23 @@

1 +from app.connectors.dfir_iris.schema.alerts import (

2 + AlertResponse,

3 + AlertsResponse,

4 + BookmarkedAlertsResponse,

5 + CaseCreationResponse,

6 + DeleteAlertResponse,

7 + FilterAlertsRequest,

8 +)

9 +from app.connectors.dfir_iris.utils.universal import (

10 + fetch_and_validate_data,

11 + initialize_client_and_alert,

12 +)

13 +from app.integrations.alert_creation_settings.models.alert_creation_settings import (

14 + AlertCreationSettings,

15 +)

16 from fastapi import HTTPException

17 from loguru import logger

18 from sqlalchemy.ext.asyncio import AsyncSession

19 from sqlalchemy.future import select

6 -from app.connectors.dfir_iris.schema.alerts import AlertResponse

7 -from app.connectors.dfir_iris.schema.alerts import AlertsResponse

8 -from app.connectors.dfir_iris.schema.alerts import BookmarkedAlertsResponse

9 -from app.connectors.dfir_iris.schema.alerts import CaseCreationResponse

10 -from app.connectors.dfir_iris.schema.alerts import DeleteAlertResponse

11 -from app.connectors.dfir_iris.schema.alerts import FilterAlertsRequest

12 -from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data

13 -from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert

14 -from app.integrations.alert_creation_settings.models.alert_creation_settings import (

15 - AlertCreationSettings,

16 -)

17 -

22 async def get_customer_code(session: AsyncSession, customer_id: int) -> str:

23 """

@@ -30,18 +33,25 @@ async def get_customer_code(session: AsyncSession, customer_id: int) -> str:

33 logger.info(f"Retrieving customer code for customer ID {customer_id}")

34 try:

35 alert_creation_settings = await session.execute(

33 - select(AlertCreationSettings).filter(AlertCreationSettings.iris_customer_id == customer_id),

36 + select(AlertCreationSettings).filter(

37 + AlertCreationSettings.iris_customer_id == customer_id,

38 + ),

39 )

40 alert_creation_settings = alert_creation_settings.scalars().first()

41 if alert_creation_settings is None:

42 return "Customer Not Found"

43 return alert_creation_settings.customer_code

44 except Exception as e:

40 - logger.error(f"Error retrieving customer code for customer ID {customer_id}: {e}")

45 + logger.error(

46 + f"Error retrieving customer code for customer ID {customer_id}: {e}",

47 + )

48 return "Customer Not Found"

44 -async def get_alerts(request: FilterAlertsRequest, session: AsyncSession) -> AlertsResponse:

51 +async def get_alerts(

52 + request: FilterAlertsRequest,

53 + session: AsyncSession,

54 +) -> AlertsResponse:

55 """

56 Retrieves alerts from the DFIR-IRIS service.

@@ -54,13 +64,25 @@ async def get_alerts(request: FilterAlertsRequest, session: AsyncSession) -> Ale

64 try:

65 client, alert = await initialize_client_and_alert("DFIR-IRIS")

66 params = construct_params(request)

57 - result = await fetch_and_validate_data(client, lambda: alert.filter_alerts(**params))

58 - logger.info(f"Successfully fetched length {len(result['data']['alerts'])} alerts")

67 + result = await fetch_and_validate_data(

68 + client,

69 + lambda: alert.filter_alerts(**params),

70 + )

71 + logger.info(

72 + f"Successfully fetched length {len(result['data']['alerts'])} alerts",

73 + )

74 # Add the customer code to each alert

75 for alert in result["data"]["alerts"]:

61 - customer_code = await get_customer_code(session, alert["customer"]["customer_id"])

76 + customer_code = await get_customer_code(

77 + session,

78 + alert["customer"]["customer_id"],

79 + )

80 alert["customer"]["customer_code"] = customer_code

63 - return AlertsResponse(success=True, message="Successfully fetched alerts", alerts=result["data"]["alerts"])

81 + return AlertsResponse(

82 + success=True,

83 + message="Successfully fetched alerts",

84 + alerts=result["data"]["alerts"],

85 + )

86 except Exception as e:

87 logger.error(f"Error fetching alerts: {e}")

88 raise HTTPException(status_code=500, detail=f"Error fetching alerts: {e}")

@@ -104,9 +126,16 @@ async def get_alert(alert_id: str, session: AsyncSession) -> AlertResponse:

126 client, alert = await initialize_client_and_alert("DFIR-IRIS")

127 result = await fetch_and_validate_data(client, alert.get_alert, alert_id)

128 # Add the customer code to the alert

107 - customer_code = await get_customer_code(session, result["data"]["customer"]["customer_id"])

129 + customer_code = await get_customer_code(

130 + session,

131 + result["data"]["customer"]["customer_id"],

132 + )

133 result["data"]["customer"]["customer_code"] = customer_code

109 - return AlertResponse(success=True, message="Successfully fetched alert", alert=result["data"])

134 + return AlertResponse(

135 + success=True,

136 + message="Successfully fetched alert",

137 + alert=result["data"],

138 + )

139

140

141 async def create_case(alert_id: str) -> CaseCreationResponse:

@@ -124,9 +153,16 @@ async def create_case(alert_id: str) -> CaseCreationResponse:

153 alert_details = await fetch_and_validate_data(client, alert.get_alert, alert_id)

154 params = construct_case_creation_params(alert_details["data"])

155 logger.info(f"Creating case with params {params}")

127 - result = await fetch_and_validate_data(client, lambda: alert.escalate_alert(int(alert_id), **params))

156 + result = await fetch_and_validate_data(

157 + client,

158 + lambda: alert.escalate_alert(int(alert_id), **params),

159 + )

160 logger.info(f"Successfully created case for alert: {result}")

129 - return CaseCreationResponse(success=True, message="Successfully created case for alert", case=result["data"])

161 + return CaseCreationResponse(

162 + success=True,

163 + message="Successfully created case for alert",

164 + case=result["data"],

165 + )

166

167

168 def construct_case_creation_params(alert_details: dict) -> dict:

@@ -144,7 +180,9 @@ def construct_case_creation_params(alert_details: dict) -> dict:

180 "case_tags": alert_details["alert_tags"],

181 "escalation_note": "Case created from CoPilot",

182 "iocs_import_list": [ioc["ioc_uuid"] for ioc in alert_details["iocs"]],

147 - "assets_import_list": [asset["asset_uuid"] for asset in alert_details["assets"]],

183 + "assets_import_list": [

184 + asset["asset_uuid"] for asset in alert_details["assets"]

185 + ],

186 }

187

188 # Replace None values with the string "None"

@@ -164,27 +202,52 @@ async def bookmark_alert(alert_id: str, bookmarked: bool) -> AlertResponse:

202 """

203 client, alert = await initialize_client_and_alert("DFIR-IRIS")

204 if bookmarked:

167 - result = await fetch_and_validate_data(client, alert.update_alert, alert_id, {"alert_tags": "bookmarked"})

168 - return AlertResponse(success=True, message="Successfully bookmarked alert", alert=result["data"])

169 - result = await fetch_and_validate_data(client, alert.update_alert, alert_id, {"alert_tags": ""})

170 - return AlertResponse(success=True, message="Successfully removed bookmark from alert", alert=result["data"])

171 -

172 -

173 -async def get_bookmarked_alerts() -> BookmarkedAlertsResponse:

205 + result = await fetch_and_validate_data(

206 + client,

207 + alert.update_alert,

208 + alert_id,

209 + {"alert_tags": "bookmarked"},

210 + )

211 + return AlertResponse(

212 + success=True,

213 + message="Successfully bookmarked alert",

214 + alert=result["data"],

215 + )

216 + result = await fetch_and_validate_data(

217 + client,

218 + alert.update_alert,

219 + alert_id,

220 + {"alert_tags": ""},

221 + )

222 + return AlertResponse(

223 + success=True,

224 + message="Successfully removed bookmark from alert",

225 + alert=result["data"],

226 + )

227 +

228 +

229 +async def get_bookmarked_alerts(session: AsyncSession) -> BookmarkedAlertsResponse:

230 """

231 Retrieves the bookmarked alerts from the system.

232

233 Returns:

234 BookmarkedAlertsResponse: The response object containing the bookmarked alerts.

235 """

180 - alerts = await get_alerts(request=FilterAlertsRequest(per_page=10000))

236 + alerts = await get_alerts(

237 + request=FilterAlertsRequest(per_page=10000),

238 + session=session,

239 + )

240 alerts = alerts.alerts

241 bookmarked_alerts = []

242 for alert in alerts:

243 if alert["alert_tags"] is not None and "bookmarked" in alert["alert_tags"]:

244 bookmarked_alerts.append(alert)

245

187 - return BookmarkedAlertsResponse(success=True, message="Successfully fetched bookmarked alerts", bookmarked_alerts=bookmarked_alerts)

246 + return BookmarkedAlertsResponse(

247 + success=True,

248 + message="Successfully fetched bookmarked alerts",

249 + bookmarked_alerts=bookmarked_alerts,

250 + )

251

252

253 async def delete_alert(alert_id: int) -> DeleteAlertResponse:

@@ -199,4 +262,8 @@ async def delete_alert(alert_id: int) -> DeleteAlertResponse:

262 """

263 client, alert = await initialize_client_and_alert("DFIR-IRIS")

264 result = await fetch_and_validate_data(client, alert.delete_alert, alert_id)

202 - return DeleteAlertResponse(success=True, message="Successfully deleted alert", alert=result["data"])

265 + return DeleteAlertResponse(

266 + success=True,

267 + message="Successfully deleted alert",

268 + alert=result["data"],

269 + )

backend/app/connectors/dfir_iris/services/assets.py

+5 -6

@@ -1,11 +1,10 @@

1 +from app.connectors.dfir_iris.schema.assets import Asset, AssetResponse, AssetState

2 +from app.connectors.dfir_iris.utils.universal import (

3 + fetch_and_validate_data,

4 + initialize_client_and_case,

5 +)

6 from fastapi import HTTPException

3 -from app.connectors.dfir_iris.schema.assets import Asset

4 -from app.connectors.dfir_iris.schema.assets import AssetResponse

5 -from app.connectors.dfir_iris.schema.assets import AssetState

6 -from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data

7 -from app.connectors.dfir_iris.utils.universal import initialize_client_and_case

8 -

9 async def get_case_assets(case_id: int) -> AssetResponse:

10 """

backend/app/connectors/dfir_iris/services/cases.py

+79 -32

@@ -1,27 +1,29 @@

1 from datetime import datetime

2 -from typing import Dict

3 -from typing import List

4 -

2 +from typing import Dict, List

3 +

4 +from app.connectors.dfir_iris.schema.cases import (

5 + CaseOlderThanBody,

6 + CaseResponse,

7 + CasesBreachedResponse,

8 + ClosedCaseResponse,

9 + PurgeCaseResponse,

10 + ReopenedCaseResponse,

11 + SingleCaseBody,

12 + SingleCaseResponse,

13 +)

14 +from app.connectors.dfir_iris.utils.universal import (

15 + create_dfir_iris_client,

16 + fetch_and_parse_data,

17 +)

18 +from app.integrations.alert_creation_settings.models.alert_creation_settings import (

19 + AlertCreationSettings,

20 +)

21 from dfir_iris_client.case import Case

22 from fastapi import HTTPException

23 from loguru import logger

24 from sqlalchemy.ext.asyncio import AsyncSession

25 from sqlalchemy.future import select