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

"""Auto logoff stuck RDP sessions when RDG flap (or direct login failure) is detected."""

from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import datetime, timezone

from sqlalchemy import select
from sqlalchemy.orm import Session
from sqlalchemy.orm.attributes import flag_modified

from app.models import Event, Host
from app.services.host_sessions import (
    list_windows_sessions,
    mark_event_session_terminated,
    parse_qwinsta_sessions,
    terminate_windows_session,
)
from app.services.rdg_client_host import ClientWorkstationNotFoundError, resolve_client_workstation
from app.services.rdg_session_flap import (
    RDG_END_TYPES,
    RDG_SUCCESS_TYPE,
    _event_user,
    _stored_flap_pair_id,
    event_has_rdg_flap,
    find_rdg_success_before_end,
)
from app.services.rdg_winrm_actions import execute_logoff_via_winrm
from app.services.rdg_workstation_session import (
    WORKSTATION_LOGIN_TYPE,
    _event_login_user,
    _login_already_closed,
    find_workstation_login_for_rdg_end,
    users_match_rdg,
)
from app.services.rdp_flap_settings import get_effective_rdp_flap_settings
from app.services.win_admin_settings import get_effective_win_admin_for_host

logger = logging.getLogger("sac.rdp_flap_auto_disconnect")

AUTO_DISCONNECT_BY = "auto:rdp_flap"
RDP_LOGIN_FAILED = "rdp.login.failed"
AUTO_DISCONNECT_DETAILS_KEY = "rdp_flap_auto_disconnect"


@dataclass(frozen=True)
class AutoDisconnectResult:
    ok: bool
    message: str
    trigger_event_id: int
    workstation_host_id: int | None = None
    login_event_id: int | None = None
    session_ids: tuple[int, ...] = ()


def _details_dict(event: Event) -> dict:
    raw = event.details
    return raw if isinstance(raw, dict) else {}


def _already_auto_disconnected(event: Event) -> bool:
    details = _details_dict(event)
    block = details.get(AUTO_DISCONNECT_DETAILS_KEY)
    if not isinstance(block, dict):
        return False
    return block.get("ok") is True


def _mark_auto_disconnect(event: Event, *, result: AutoDisconnectResult) -> None:
    details = dict(_details_dict(event))
    details[AUTO_DISCONNECT_DETAILS_KEY] = {
        "ok": result.ok,
        "message": result.message,
        "workstation_host_id": result.workstation_host_id,
        "login_event_id": result.login_event_id,
        "session_ids": list(result.session_ids),
        "at": datetime.now(timezone.utc).isoformat(),
    }
    event.details = details
    flag_modified(event, "details")


def _norm_user_filter(user: str) -> str:
    text = (user or "").strip()
    if "\\" in text:
        return text.split("\\")[-1].strip().lower()
    if "@" in text:
        return text.split("@")[0].strip().lower()
    return text.lower()


def _sessions_for_user(sessions, user: str):
    needle = _norm_user_filter(user)
    if not needle:
        return []
    matched = []
    for row in sessions:
        if needle in _norm_user_filter(row.user):
            matched.append(row)
    return matched


def find_open_workstation_login(db: Session, *, host_id: int, user: str) -> Event | None:
    if not user.strip():
        return None
    candidates = db.scalars(
        select(Event)
        .where(
            Event.host_id == host_id,
            Event.type == WORKSTATION_LOGIN_TYPE,
        )
        .order_by(Event.occurred_at.desc())
    ).all()
    for login in candidates:
        if _login_already_closed(login):
            continue
        if users_match_rdg(_event_login_user(login), user):
            return login
    return None


def _rdg_pair_success_event(db: Session, event: Event) -> Event | None:
    if event.type == RDG_SUCCESS_TYPE and event_has_rdg_flap(event):
        pair_id = _stored_flap_pair_id(event)
        if pair_id is not None:
            return db.get(Event, pair_id)
        return event
    if event.type in RDG_END_TYPES:
        pair_id = _stored_flap_pair_id(event)
        if pair_id is not None:
            return db.get(Event, pair_id)
        return find_rdg_success_before_end(db, event)
    return None


def _disconnect_on_workstation(
    db: Session,
    *,
    trigger_event: Event,
    workstation: Host,
    rdg_event: Event | None,
    user: str,
    login_event: Event | None,
) -> AutoDisconnectResult:
    win_cfg = get_effective_win_admin_for_host(db, workstation)
    sessions, qwinsta = list_windows_sessions(workstation, win_cfg)
    if qwinsta is None or not qwinsta.ok:
        message = qwinsta.message if qwinsta else "qwinsta failed"
        return AutoDisconnectResult(
            ok=False,
            message=message,
            trigger_event_id=trigger_event.id,
            workstation_host_id=workstation.id,
            login_event_id=login_event.id if login_event else None,
        )

    parsed = parse_qwinsta_sessions(qwinsta.stdout, filter_user=user)
    matched = _sessions_for_user(parsed, user)
    if not matched and qwinsta.stdout.strip():
        parsed = parse_qwinsta_sessions(qwinsta.stdout)
        matched = _sessions_for_user(parsed, user)
    if not matched:
        snippet = " ".join(qwinsta.stdout.split())[:240]
        parsed_note = f", parsed={len(parsed)} session(s)" if parsed else ""
        message = (
            f"No matching Windows session for user {user} on {workstation.hostname}"
            f"{parsed_note}"
        )
        if snippet:
            message = f"{message}; qwinsta: {snippet}"
        return AutoDisconnectResult(
            ok=False,
            message=message,
            trigger_event_id=trigger_event.id,
            workstation_host_id=workstation.id,
            login_event_id=login_event.id if login_event else None,
        )

    logged_off: list[int] = []
    errors: list[str] = []
    for row in matched:
        if rdg_event is not None:
            cmd = execute_logoff_via_winrm(
                db,
                rdg_event,
                session_id=int(row.session_id),
                requested_by=AUTO_DISCONNECT_BY,
            )
            if cmd.status == "completed":
                logged_off.append(int(row.session_id))
            else:
                errors.append(cmd.result_stderr or cmd.result_stdout or f"logoff {row.session_id} failed")
        else:
            result = terminate_windows_session(workstation, win_cfg, row.session_id)
            if result is not None and result.ok:
                logged_off.append(int(row.session_id))
            else:
                errors.append(result.message if result else f"logoff {row.session_id} failed")

    if login_event is not None and logged_off:
        mark_event_session_terminated(login_event, by_username=AUTO_DISCONNECT_BY)

    if logged_off:
        message = f"Auto logoff session(s) {', '.join(str(s) for s in logged_off)} on {workstation.hostname}"
        ok = True
    else:
        ok = False
        message = "; ".join(errors) if errors else "logoff failed"
    return AutoDisconnectResult(
        ok=ok,
        message=message,
        trigger_event_id=trigger_event.id,
        workstation_host_id=workstation.id,
        login_event_id=login_event.id if login_event else None,
        session_ids=tuple(logged_off),
    )


def _auto_disconnect_rdg_flap(db: Session, event: Event) -> AutoDisconnectResult | None:
    if not event_has_rdg_flap(event):
        return None
    if _already_auto_disconnected(event):
        return None

    rdg_success = _rdg_pair_success_event(db, event)
    if rdg_success is None:
        return None

    user = _event_user(event) or _event_user(rdg_success)
    if not user:
        return None

    try:
        workstation = resolve_client_workstation(db, rdg_success)
    except ClientWorkstationNotFoundError as exc:
        result = AutoDisconnectResult(
            ok=False,
            message=str(exc),
            trigger_event_id=event.id,
        )
        _mark_auto_disconnect(event, result=result)
        return result

    rdg_end = event if event.type in RDG_END_TYPES else db.get(Event, _stored_flap_pair_id(event) or -1)
    login_event = find_workstation_login_for_rdg_end(db, rdg_end) if rdg_end is not None else None

    result = _disconnect_on_workstation(
        db,
        trigger_event=event,
        workstation=workstation,
        rdg_event=rdg_success,
        user=user,
        login_event=login_event,
    )
    _mark_auto_disconnect(event, result=result)
    return result


def _auto_disconnect_direct_rdp_failed(db: Session, event: Event) -> AutoDisconnectResult | None:
    if event.type != RDP_LOGIN_FAILED:
        return None
    if _already_auto_disconnected(event):
        return None
    host = event.host
    if host is None:
        return None

    user = _event_login_user(event)
    if not user:
        return None

    login_event = find_open_workstation_login(db, host_id=host.id, user=user)
    if login_event is None:
        return None

    result = _disconnect_on_workstation(
        db,
        trigger_event=event,
        workstation=host,
        rdg_event=None,
        user=user,
        login_event=login_event,
    )
    _mark_auto_disconnect(event, result=result)
    return result


def maybe_auto_disconnect_stuck_rdp_session(db: Session, event: Event) -> AutoDisconnectResult | None:
    """Log off stuck user session when auto-disconnect is enabled."""
    if not get_effective_rdp_flap_settings(db).auto_disconnect:
        return None

    result = _auto_disconnect_rdg_flap(db, event)
    if result is not None:
        if result.ok:
            logger.info(
                "auto rdp flap disconnect ok event_id=%s sessions=%s",
                event.event_id,
                result.session_ids,
            )
        else:
            logger.warning(
                "auto rdp flap disconnect failed event_id=%s: %s",
                event.event_id,
                result.message,
            )
        return result

    result = _auto_disconnect_direct_rdp_failed(db, event)
    if result is not None:
        if result.ok:
            logger.info(
                "auto direct rdp disconnect ok event_id=%s sessions=%s",
                event.event_id,
                result.session_ids,
            )
        else:
            logger.warning(
                "auto direct rdp disconnect failed event_id=%s: %s",
                event.event_id,
                result.message,
            )
    return result