security-alert-center/backend/app/services/agent_commands.py

"""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))