"""Queue and resolve agent commands (qwinsta, logoff).""" from __future__ import annotations import uuid from datetime import datetime, timezone from fastapi import HTTPException from sqlalchemy import select from sqlalchemy.orm import Session from app.config import get_settings from app.models import AgentCommand, Event, Host from app.services.rdg_session_flap import event_has_rdg_flap from app.services.win_admin_settings import get_effective_win_admin_config def _win_admin_configured() -> bool: return get_effective_win_admin_config().configured def win_admin_run_as() -> dict[str, str] | None: cfg = get_effective_win_admin_config() if not cfg.configured: return None return { "user": cfg.user, "password": cfg.password, } def command_to_dict(cmd: AgentCommand, *, include_run_as: bool = False) -> dict: out = { "id": cmd.command_uuid, "type": cmd.command_type, "params": cmd.params if isinstance(cmd.params, dict) else {}, "status": cmd.status, "result_stdout": cmd.result_stdout, "result_stderr": cmd.result_stderr, "created_at": cmd.created_at.isoformat() if cmd.created_at else None, "completed_at": cmd.completed_at.isoformat() if cmd.completed_at else None, } if include_run_as and cmd.status == "pending": run_as = win_admin_run_as() if run_as: out["run_as"] = run_as return out def queue_qwinsta(db: Session, event: Event, *, requested_by: str) -> AgentCommand: if not event_has_rdg_flap(event): raise HTTPException(status_code=400, detail="Event is not flagged as RDG session flap") if not _win_admin_configured(): raise HTTPException( status_code=503, detail="Windows domain admin is not configured (Settings or SAC_WIN_ADMIN_*)", ) details = event.details if isinstance(event.details, dict) else {} user = details.get("user") cmd = AgentCommand( command_uuid=str(uuid.uuid4()), host_id=event.host_id, event_id=event.id, command_type="qwinsta", params={"user": user} if user else {}, status="pending", requested_by=requested_by, ) db.add(cmd) db.flush() return cmd def queue_logoff( db: Session, event: Event, *, session_id: int, requested_by: str, ) -> AgentCommand: if not event_has_rdg_flap(event): raise HTTPException(status_code=400, detail="Event is not flagged as RDG session flap") if not _win_admin_configured(): raise HTTPException( status_code=503, detail="Windows domain admin is not configured (Settings or SAC_WIN_ADMIN_*)", ) cmd = AgentCommand( command_uuid=str(uuid.uuid4()), host_id=event.host_id, event_id=event.id, command_type="logoff", params={"session_id": session_id}, status="pending", requested_by=requested_by, ) db.add(cmd) db.flush() return cmd def get_command_for_host_poll(db: Session, host: Host, *, limit: int = 5) -> list[AgentCommand]: return list( db.scalars( select(AgentCommand) .where( AgentCommand.host_id == host.id, AgentCommand.status == "pending", ) .order_by(AgentCommand.created_at.asc()) .limit(limit) ).all() ) def resolve_host_by_agent_instance_id(db: Session, agent_instance_id: str) -> Host | None: if not agent_instance_id.strip(): return None return db.scalar(select(Host).where(Host.agent_instance_id == agent_instance_id.strip())) def complete_command( db: Session, command_uuid: str, *, status: str, stdout: str | None = None, stderr: str | None = None, ) -> AgentCommand | None: cmd = db.scalar(select(AgentCommand).where(AgentCommand.command_uuid == command_uuid)) if cmd is None: return None if cmd.status != "pending": return cmd cmd.status = status cmd.result_stdout = stdout cmd.result_stderr = stderr cmd.completed_at = datetime.now(timezone.utc) db.flush() return cmd def get_command_by_uuid(db: Session, command_uuid: str) -> AgentCommand | None: return db.scalar(select(AgentCommand).where(AgentCommand.command_uuid == command_uuid))