security-alert-center/backend/tests/test_rdp_flap_auto_disconnect.py

"""Tests for RDP flap auto-disconnect setting and service."""

import uuid
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

from app.config import get_settings
from app.models import Host
from app.models.ui_settings import UI_SETTINGS_ROW_ID, UiSettings
from app.services.ingest import ingest_event
from app.services.problems import maybe_create_problem
from app.services.rdp_flap_auto_disconnect import maybe_auto_disconnect_stuck_rdp_session
from app.services.rdp_flap_settings import get_effective_rdp_flap_settings, upsert_rdp_flap_settings
from app.services.winrm_connect import WinRmCmdResult
from tests.test_ingest import VALID_EVENT


def _payload(**overrides):
    base = {
        **VALID_EVENT,
        "event_id": str(uuid.uuid4()),
        "occurred_at": datetime.now(timezone.utc).isoformat(),
    }
    base.update(overrides)
    return base


def _ingest(db, occurred_at: datetime, **overrides):
    payload = _payload(**overrides)
    payload["occurred_at"] = occurred_at.isoformat()
    event, _ = ingest_event(db, payload)
    db.flush()
    return event


@pytest.fixture
def rdg_settings(monkeypatch):
    monkeypatch.setenv("SAC_RDG_FLAP_WINDOW_MIN_SEC", "1")
    monkeypatch.setenv("SAC_RDG_FLAP_WINDOW_MAX_SEC", "10")
    monkeypatch.setenv("SAC_RDG_FLAP_DEDUP_SEC", "30")
    monkeypatch.setenv("SAC_WIN_ADMIN_USER", r"B26\admin")
    monkeypatch.setenv("SAC_WIN_ADMIN_PASSWORD", "secret")
    get_settings.cache_clear()
    yield
    get_settings.cache_clear()


@pytest.fixture
def rdg_hosts(db_session):
    ws = Host(
        hostname="TSA-PC",
        os_family="windows",
        product="rdp-login-monitor",
        ipv4="192.168.163.100",
    )
    gw = Host(
        hostname="K6A-DC3",
        os_family="windows",
        product="rdp-login-monitor",
        ipv4="192.168.160.40",
    )
    db_session.add_all([ws, gw])
    db_session.commit()
    return ws, gw


def test_rdp_flap_settings_default_disabled(db_session):
    cfg = get_effective_rdp_flap_settings(db_session)
    assert cfg.auto_disconnect is False
    assert cfg.source == "default"


def test_rdp_flap_settings_upsert(db_session):
    upsert_rdp_flap_settings(db_session, auto_disconnect=True)
    cfg = get_effective_rdp_flap_settings(db_session)
    assert cfg.auto_disconnect is True
    row = db_session.get(UiSettings, UI_SETTINGS_ROW_ID)
    assert row is not None
    assert row.auto_rdp_flap_disconnect is True


def test_auto_disconnect_skipped_when_disabled(db_session, rdg_settings, rdg_hosts, monkeypatch):
    ws, gw = rdg_hosts
    t0 = datetime.now(timezone.utc)
    user = r"B26\TSA"
    details = {"user": user, "internal_ip": ws.ipv4}
    gw_payload = {
        "host": {"hostname": gw.hostname, "os_family": "windows", "ipv4": gw.ipv4},
        "source": {"product": "rdp-login-monitor", "product_version": "2.1.8-SAC"},
    }

    _ingest(
        db_session,
        t0,
        **gw_payload,
        type="rdg.connection.success",
        category="auth",
        severity="info",
        title="302",
        summary="302",
        details=details,
    )
    end = _ingest(
        db_session,
        t0 + timedelta(seconds=4),
        **gw_payload,
        type="rdg.connection.disconnected",
        category="auth",
        severity="info",
        title="303",
        summary="303",
        details=details,
    )
    maybe_create_problem(db_session, end)

    mock_logoff = MagicMock()
    monkeypatch.setattr("app.services.rdp_flap_auto_disconnect.execute_logoff_via_winrm", mock_logoff)

    result = maybe_auto_disconnect_stuck_rdp_session(db_session, end)
    assert result is None
    mock_logoff.assert_not_called()


def test_auto_disconnect_rdg_flap_calls_logoff(db_session, rdg_settings, rdg_hosts, monkeypatch):
    ws, gw = rdg_hosts
    upsert_rdp_flap_settings(db_session, auto_disconnect=True)
    t0 = datetime.now(timezone.utc)
    user = r"B26\TSA"
    details = {"user": user, "internal_ip": ws.ipv4}
    gw_payload = {
        "host": {"hostname": gw.hostname, "os_family": "windows", "ipv4": gw.ipv4},
        "source": {"product": "rdp-login-monitor", "product_version": "2.1.8-SAC"},
    }

    _ingest(
        db_session,
        t0,
        **gw_payload,
        type="rdg.connection.success",
        category="auth",
        severity="info",
        title="302",
        summary="302",
        details=details,
    )
    _ingest(
        db_session,
        t0 + timedelta(seconds=1),
        host={"hostname": ws.hostname, "os_family": "windows", "ipv4": ws.ipv4},
        source={"product": "rdp-login-monitor", "product_version": "2.1.8-SAC"},
        type="rdp.login.success",
        category="auth",
        severity="info",
        title="login",
        summary="login",
        details={"user": user},
    )
    end = _ingest(
        db_session,
        t0 + timedelta(seconds=4),
        **gw_payload,
        type="rdg.connection.disconnected",
        category="auth",
        severity="info",
        title="303",
        summary="303",
        details=details,
    )
    maybe_create_problem(db_session, end)

    qwinsta_stdout = "SESSIONNAME       USERNAME                 ID  STATE\n rdp-tcp#0         B26\\TSA                    5  Active\n"
    monkeypatch.setattr(
        "app.services.rdp_flap_auto_disconnect.list_windows_sessions",
        lambda host, cfg: (
            [],
            WinRmCmdResult(ok=True, message="ok", target=host.ipv4 or "", stdout=qwinsta_stdout),
        ),
    )
    mock_logoff = MagicMock(return_value=SimpleNamespace(status="completed", result_stderr=None, result_stdout="ok"))
    monkeypatch.setattr("app.services.rdp_flap_auto_disconnect.execute_logoff_via_winrm", mock_logoff)

    result = maybe_auto_disconnect_stuck_rdp_session(db_session, end)
    assert result is not None
    assert result.ok is True
    assert result.session_ids == (5,)
    mock_logoff.assert_called_once()
    assert end.details["rdp_flap_auto_disconnect"]["ok"] is True


def test_auto_disconnect_direct_rdp_failed(db_session, rdg_settings, rdg_hosts, monkeypatch):
    ws, _gw = rdg_hosts
    upsert_rdp_flap_settings(db_session, auto_disconnect=True)
    t0 = datetime.now(timezone.utc)
    user = r"B26\TSA"
    ws_payload = {
        "host": {"hostname": ws.hostname, "os_family": "windows", "ipv4": ws.ipv4},
        "source": {"product": "rdp-login-monitor", "product_version": "2.1.8-SAC"},
    }

    _ingest(
        db_session,
        t0,
        **ws_payload,
        type="rdp.login.success",
        category="auth",
        severity="info",
        title="login",
        summary="login",
        details={"user": user},
    )
    failed = _ingest(
        db_session,
        t0 + timedelta(seconds=30),
        **ws_payload,
        type="rdp.login.failed",
        category="auth",
        severity="warning",
        title="failed",
        summary="failed",
        details={"user": user},
    )

    qwinsta_stdout = "SESSIONNAME       USERNAME                 ID  STATE\n rdp-tcp#0         B26\\TSA                    3  Active\n"
    monkeypatch.setattr(
        "app.services.rdp_flap_auto_disconnect.list_windows_sessions",
        lambda host, cfg: (
            [],
            WinRmCmdResult(ok=True, message="ok", target=host.ipv4 or "", stdout=qwinsta_stdout),
        ),
    )
    monkeypatch.setattr(
        "app.services.rdp_flap_auto_disconnect.terminate_windows_session",
        lambda host, cfg, sid: WinRmCmdResult(ok=True, message="ok", target=host.ipv4 or ""),
    )

    result = maybe_auto_disconnect_stuck_rdp_session(db_session, failed)
    assert result is not None
    assert result.ok is True
    assert result.session_ids == (3,)