"""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_config 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_config(db) 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, ) matched = _sessions_for_user(parse_qwinsta_sessions(qwinsta.stdout, filter_user=user), user) if not matched and qwinsta.stdout.strip(): matched = _sessions_for_user(parse_qwinsta_sessions(qwinsta.stdout), user) if not matched: return AutoDisconnectResult( ok=False, message="No matching Windows session for user", 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: return AutoDisconnectResult( ok=False, message=str(exc), trigger_event_id=event.id, ) 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 win_cfg = get_effective_win_admin_config(db) if not win_cfg.configured: logger.warning( "auto rdp flap disconnect skipped: win admin not configured (event_id=%s)", event.event_id, ) 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