forked from EmpireProject/Empire
-
-
Notifications
You must be signed in to change notification settings - Fork 576
/
agent_file_api.py
74 lines (56 loc) · 2.25 KB
/
agent_file_api.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
from fastapi import Depends, HTTPException
from empire.server.api.api_router import APIRouter
from empire.server.api.jwt_auth import get_current_active_user
from empire.server.api.v2.agent.agent_file_dto import AgentFile, domain_to_dto_file
from empire.server.api.v2.shared_dependencies import CurrentSession
from empire.server.api.v2.shared_dto import BadRequestResponse, NotFoundResponse
from empire.server.core.agent_file_service import AgentFileService
from empire.server.core.agent_service import AgentService
from empire.server.core.db import models
from empire.server.server import main
agent_file_service: AgentFileService = main.agentfilesv2
agent_service: AgentService = main.agentsv2
router = APIRouter(
prefix="/api/v2/agents/{agent_id}/files",
tags=["agents"],
responses={
404: {"description": "Not found", "model": NotFoundResponse},
400: {"description": "Bad request", "model": BadRequestResponse},
},
dependencies=[Depends(get_current_active_user)],
)
async def get_agent(agent_id: str, db: CurrentSession):
agent = agent_service.get_by_id(db, agent_id)
if agent:
return agent
raise HTTPException(404, f"Agent not found for id {agent_id}")
async def get_file(
uid: int, db: CurrentSession, db_agent: models.Agent = Depends(get_agent)
):
file = agent_file_service.get_file(db, db_agent.session_id, uid)
if file:
return file
raise HTTPException(
404, f"File not found for agent {db_agent.session_id} and file id {uid}"
)
@router.get("/root")
async def read_file_root(
db: CurrentSession, db_agent: models.Agent = Depends(get_agent)
):
file = agent_file_service.get_file_by_path(db, db_agent.session_id, "/")
if file:
return domain_to_dto_file(*file)
raise HTTPException(
404, f'File not found for agent {db_agent.session_id} and file path "/"'
)
@router.get("/{uid}", response_model=AgentFile)
async def read_file(
uid: int,
db_agent: models.Agent = Depends(get_agent),
db_file: tuple[models.AgentFile, list[models.AgentFile]] | None = Depends(get_file),
):
if db_file:
return domain_to_dto_file(*db_file)
raise HTTPException(
404, f'File not found for agent {db_agent.session_id} and file path "/"'
)