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

"""WinRM bundle delivery and CLIXML error parsing tests."""

from pathlib import Path
from unittest.mock import patch

from app.services.rdp_bundle_delivery import (
    _UTF8_BOM,
    build_rdp_bundle_zip,
    get_rdp_bundle_zip,
    register_rdp_bundle_zip,
)
from app.services.winrm_connect import (
    RDP_BUNDLE_REQUIRED,
    RDP_LEGACY_STAGING,
    RDP_REMOTE_STAGING_DIRNAME,
    WinRmCmdResult,
    _clixml_to_plain,
    _custom_deploy_body,
    _decode_winrm_bytes,
    _deploy_from_staging_body,
    _download_bundle_body,
    _prepare_staging_body,
    _winrm_failure_detail,
    run_winrm_rdp_monitor_update,
)


def test_decode_winrm_bytes_reads_russian_cp866_console():
    sample = "Деплой: корень шары".encode("cp866")
    decoded = _decode_winrm_bytes(sample)
    assert "Деплой" in decoded
    assert "шары" in decoded


def test_decode_winrm_bytes_handles_utf8_multibyte():
    text = "\u0414\u0435\u043f\u043b\u043e\u0439: \u043a\u043e\u0440\u0435\u043d\u044c \u0448\u0430\u0440\u044b"
    assert _decode_winrm_bytes(text.encode("utf-8")) == text


def test_deploy_from_staging_includes_deploy_log_tail():
    script = _deploy_from_staging_body()
    assert "deploy.log" in script
    assert "Get-Content" in script
    assert "-Encoding UTF8" in script
    assert RDP_REMOTE_STAGING_DIRNAME in script


def test_prepare_staging_recreates_remote_dir():
    script = _prepare_staging_body()
    assert "Remove-Item" in script
    assert RDP_REMOTE_STAGING_DIRNAME in script
    assert RDP_LEGACY_STAGING in script


def test_download_bundle_uses_invoke_webrequest():
    script = _download_bundle_body(
        "https://sac.example/api/v1/agent/rdp-bundle/token",
    )
    assert "Invoke-WebRequest" in script
    assert "$env:TEMP" in script
    assert RDP_REMOTE_STAGING_DIRNAME in script
    assert "ZipFile]::ExtractToDirectory" in script
    assert "_sac_staging" not in script or RDP_LEGACY_STAGING in script
    assert "Expand-Archive" not in script


def test_winrm_failure_detail_prefers_error_line_over_progress_stdout():
    stdout = "Downloading bundle: https://sac.test/bundle\nERROR: destination not empty"
    detail = _winrm_failure_detail(stdout, "", 1)
    assert detail == "destination not empty"


def test_deploy_from_staging_runs_local_bundle():
    script = _deploy_from_staging_body()
    assert "-SourceShareRoot" in script
    assert "Deploy-LoginMonitor.ps1" in script


def test_custom_deploy_script_escapes_single_quotes():
    script = _custom_deploy_body(r"C:\Users\O'Brien\Deploy-LoginMonitor.ps1")
    assert "O''Brien" in script


def test_rdp_bundle_zip_adds_utf8_bom_for_ps1(tmp_path: Path, monkeypatch):
    monkeypatch.setattr("app.services.rdp_bundle_delivery.BUNDLE_DIR", tmp_path)
    repo = tmp_path / "repo"
    repo.mkdir()
    (repo / "Deploy-LoginMonitor.ps1").write_text("# test", encoding="utf-8")
    zip_bytes = build_rdp_bundle_zip(repo, ("Deploy-LoginMonitor.ps1",))
    import zipfile
    import io

    with zipfile.ZipFile(io.BytesIO(zip_bytes)) as archive:
        data = archive.read("Deploy-LoginMonitor.ps1")
    assert data.startswith(_UTF8_BOM)


def test_rdp_bundle_zip_roundtrip(tmp_path: Path, monkeypatch):
    monkeypatch.setattr("app.services.rdp_bundle_delivery.BUNDLE_DIR", tmp_path)
    repo = tmp_path / "repo"
    repo.mkdir()
    for name in RDP_BUNDLE_REQUIRED:
        (repo / name).write_text(name, encoding="utf-8")
    zip_bytes = build_rdp_bundle_zip(repo, tuple(RDP_BUNDLE_REQUIRED))
    token = register_rdp_bundle_zip(zip_bytes)
    assert get_rdp_bundle_zip(token) == zip_bytes


def test_resolve_winrm_cmd_uses_system32_for_builtins():
    from app.services.winrm_connect import _resolve_winrm_cmd

    assert _resolve_winrm_cmd("hostname") == "%SystemRoot%\\System32\\hostname.exe"
    assert _resolve_winrm_cmd("qwinsta") == "%SystemRoot%\\System32\\qwinsta.exe"
    assert _resolve_winrm_cmd("logoff 2 /v") == "%SystemRoot%\\System32\\logoff.exe 2 /v"


def test_run_winrm_rdp_monitor_update_downloads_bundle_from_sac(tmp_path: Path):
    repo = tmp_path / "repo"
    repo.mkdir()
    for name in RDP_BUNDLE_REQUIRED:
        (repo / name).write_text(f"content-{name}", encoding="utf-8")

    calls: list[str] = []

    def fake_run_ps(*, script: str, **kwargs) -> WinRmCmdResult:
        calls.append(script)
        if "Remove-Item" in script:
            return WinRmCmdResult(ok=True, message="staging ok", target="pc", stdout="Staging ready")
        if "Invoke-WebRequest" in script:
            return WinRmCmdResult(ok=True, message="download ok", target="pc", stdout="Bundle extracted")
        if "powershell.exe -NoProfile" in script or "& $deploy" in script:
            return WinRmCmdResult(
                ok=True,
                message="WinRM OK (pc), exit 0",
                target="pc",
                stdout="deployed 1.2.3",
                exit_code=0,
            )
        return WinRmCmdResult(ok=True, message="ok", target="pc", exit_code=0)

    with (
        patch("app.services.winrm_connect._fetch_rdp_bundle_dir", return_value=repo),
        patch("app.services.winrm_connect._bundle_download_url", return_value="https://sac.test/bundle.zip"),
        patch("app.services.winrm_connect.run_winrm_ps", side_effect=fake_run_ps),
    ):
        result = run_winrm_rdp_monitor_update(
            target="pc",
            user="B26\\admin",
            password="pw",
            repo_url="https://git.example.com/RDP-login-monitor.git",
        )

    assert result.ok is True
    assert "SAC served bundle" in result.stdout
    assert any("Invoke-WebRequest" in call for call in calls)
    assert not any("FromBase64String" in call for call in calls)


def test_winrm_failure_detail_never_returns_raw_clixml_progress():
    clixml = (
        "#< CLIXML <Objs Version=\"1.1.0.1\">"
        '<Obj S="progress"><MS><PR N="Record"><AV>Preparing modules for first use.</AV></PR></MS></Obj>'
    )
    detail = _winrm_failure_detail("", clixml, 1)
    assert "#< CLIXML" not in detail


def test_clixml_to_plain_strips_progress_only_blob():
    clixml = "#< CLIXML <Objs><Obj S=\"progress\"><AV>Preparing modules for first use.</AV></Obj></Objs>"
    assert _clixml_to_plain(clixml) == ""