77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
from typing import Any
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|
from fastapi.responses import Response
|
|
from pydantic import BaseModel, Field
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.auth.api_key import get_api_key_auth
|
|
from app.database import get_db
|
|
from app.services.agent_commands import (
|
|
complete_command,
|
|
resolve_host_by_agent_instance_id,
|
|
)
|
|
from app.services.agent_poll import build_agent_poll_payload
|
|
|
|
router = APIRouter(prefix="/agent", tags=["agent"])
|
|
|
|
|
|
class AgentCommandResultBody(BaseModel):
|
|
status: str = Field(pattern="^(completed|failed)$")
|
|
stdout: str | None = None
|
|
stderr: str | None = None
|
|
|
|
|
|
class AgentCommandsPollResponse(BaseModel):
|
|
config_revision: int = 0
|
|
config: dict[str, Any] = Field(default_factory=dict)
|
|
update: dict[str, Any] = Field(default_factory=dict)
|
|
commands: list[dict[str, Any]] = Field(default_factory=list)
|
|
|
|
|
|
@router.get("/commands", response_model=AgentCommandsPollResponse)
|
|
def poll_agent_commands(
|
|
agent_instance_id: str = Query(..., min_length=1),
|
|
db: Session = Depends(get_db),
|
|
_api_key: str = Depends(get_api_key_auth),
|
|
) -> AgentCommandsPollResponse:
|
|
host = resolve_host_by_agent_instance_id(db, agent_instance_id)
|
|
if host is None:
|
|
raise HTTPException(status_code=404, detail="Unknown agent_instance_id")
|
|
payload = build_agent_poll_payload(db, host)
|
|
return AgentCommandsPollResponse(**payload)
|
|
|
|
|
|
@router.post("/commands/{command_uuid}/result")
|
|
def post_agent_command_result(
|
|
command_uuid: str,
|
|
body: AgentCommandResultBody,
|
|
db: Session = Depends(get_db),
|
|
_api_key: str = Depends(get_api_key_auth),
|
|
) -> dict[str, str]:
|
|
cmd = complete_command(
|
|
db,
|
|
command_uuid,
|
|
status=body.status,
|
|
stdout=body.stdout,
|
|
stderr=body.stderr,
|
|
)
|
|
if cmd is None:
|
|
raise HTTPException(status_code=404, detail="Command not found")
|
|
db.commit()
|
|
return {"status": cmd.status, "command_uuid": cmd.command_uuid}
|
|
|
|
|
|
@router.get("/rdp-bundle/{token}")
|
|
def download_rdp_bundle(token: str) -> Response:
|
|
from app.services.rdp_bundle_delivery import get_rdp_bundle_zip
|
|
|
|
data = get_rdp_bundle_zip(token)
|
|
if not data:
|
|
raise HTTPException(status_code=404, detail="Bundle not found or expired")
|
|
return Response(
|
|
content=data,
|
|
media_type="application/zip",
|
|
headers={"Content-Disposition": 'attachment; filename="sac-rdp-bundle.zip"'},
|
|
)
|