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

"""RDG qwinsta/logoff via WinRM on client workstation (variant B)."""

from __future__ import annotations

import uuid
from datetime import datetime, timezone

from fastapi import HTTPException
from sqlalchemy.orm import Session

from app.models import AgentCommand, Event, Host
from app.services.rdg_client_host import ClientWorkstationNotFoundError, resolve_client_workstation
from app.services.rdg_display import event_supports_rdg_client_qwinsta
from app.services.rdg_session_flap import event_internal_ip
from app.services.win_admin_settings import get_effective_win_admin_for_host
from app.services.winrm_connect import (
    WinRmCmdResult,
    run_winrm_logoff,
    run_winrm_on_host_targets,
    run_winrm_qwinsta,
)


def _require_win_admin(db: Session, host: Host | None = None):
    if host is not None:
        cfg = get_effective_win_admin_for_host(db, host)
    else:
        from app.services.win_admin_settings import get_effective_win_admin_config

        cfg = get_effective_win_admin_config(db)
    if not cfg.configured:
        raise HTTPException(
            status_code=503,
            detail="Windows admin is not configured (host override or Settings → Windows)",
        )
    return cfg


def _require_rdg_client_qwinsta(db: Session, event: Event) -> None:
    if event_supports_rdg_client_qwinsta(event, db):
        return
    raise HTTPException(
        status_code=400,
        detail="Event is not an RD Gateway connection with internal_ip (client workstation)",
    )


def _resolve_client(db: Session, event: Event) -> Host:
    try:
        return resolve_client_workstation(db, event)
    except ClientWorkstationNotFoundError as exc:
        if not exc.internal_ip:
            raise HTTPException(
                status_code=400,
                detail="Event has no internal_ip (client workstation address)",
            ) from exc
        raise HTTPException(status_code=404, detail=str(exc)) from exc


def _format_result_message(result: WinRmCmdResult, *, attempts: list[str]) -> str:
    message = result.message
    if attempts:
        message = f"{message} (пробовали: {', '.join(attempts)})"
    return message


def _persist_command(
    db: Session,
    *,
    client_host: Host,
    event: Event,
    command_type: str,
    params: dict,
    requested_by: str,
    result: WinRmCmdResult,
    attempts: list[str],
) -> AgentCommand:
    now = datetime.now(timezone.utc)
    cmd = AgentCommand(
        command_uuid=str(uuid.uuid4()),
        host_id=client_host.id,
        event_id=event.id,
        command_type=command_type,
        params=params,
        status="completed" if result.ok else "failed",
        result_stdout=result.stdout or None,
        result_stderr=(result.stderr or _format_result_message(result, attempts=attempts) or None)
        if not result.ok
        else None,
        requested_by=requested_by,
        created_at=now,
        completed_at=now,
    )
    db.add(cmd)
    db.flush()
    return cmd


def execute_qwinsta_via_winrm(db: Session, event: Event, *, requested_by: str) -> AgentCommand:
    _require_rdg_client_qwinsta(db, event)
    client_host = _resolve_client(db, event)
    cfg = _require_win_admin(db, client_host)

    details = event.details if isinstance(event.details, dict) else {}
    user = details.get("user")
    internal_ip = event_internal_ip(event)

    result, attempts = run_winrm_on_host_targets(
        client_host,
        user=cfg.user,
        password=cfg.password,
        action=lambda target: run_winrm_qwinsta(target=target, user=cfg.user, password=cfg.password),
    )
    assert result is not None

    params = {
        "execution": "winrm",
        "user": user,
        "internal_ip": internal_ip,
        "client_hostname": client_host.hostname,
        "client_host_id": client_host.id,
        "winrm_target": result.target,
    }
    return _persist_command(
        db,
        client_host=client_host,
        event=event,
        command_type="qwinsta",
        params=params,
        requested_by=requested_by,
        result=result,
        attempts=attempts,
    )


def execute_logoff_via_winrm(
    db: Session,
    event: Event,
    *,
    session_id: int,
    requested_by: str,
) -> AgentCommand:
    _require_rdg_client_qwinsta(db, event)
    client_host = _resolve_client(db, event)
    cfg = _require_win_admin(db, client_host)

    details = event.details if isinstance(event.details, dict) else {}
    user = details.get("user")
    internal_ip = event_internal_ip(event)

    result, attempts = run_winrm_on_host_targets(
        client_host,
        user=cfg.user,
        password=cfg.password,
        action=lambda target: run_winrm_logoff(
            target=target,
            user=cfg.user,
            password=cfg.password,
            session_id=session_id,
        ),
    )
    assert result is not None

    params = {
        "execution": "winrm",
        "user": user,
        "internal_ip": internal_ip,
        "client_hostname": client_host.hostname,
        "client_host_id": client_host.id,
        "session_id": session_id,
        "winrm_target": result.target,
    }
    return _persist_command(
        db,
        client_host=client_host,
        event=event,
        command_type="logoff",
        params=params,
        requested_by=requested_by,
        result=result,
        attempts=attempts,
    )