fix: deliver RDP bundle via SAC HTTP zip download, not WinRM chunks (0.20.15)

backend/app/services/rdp_bundle_delivery.py

@@ -0,0 +1,57 @@
+"""Short-lived in-memory RDP agent bundles for WinRM clients to download."""
+
+from __future__ import annotations
+
+import io
+import secrets
+import time
+import zipfile
+from pathlib import Path
+from threading import Lock
+
+TTL_SECONDS = 600
+
+_lock = Lock()
+_entries: dict[str, tuple[bytes, float]] = {}
+
+
+def _prune_expired(now: float) -> None:
+    expired = [token for token, (_, expires_at) in _entries.items() if expires_at <= now]
+    for token in expired:
+        _entries.pop(token, None)
+
+
+def build_rdp_bundle_zip(repo_dir: Path, filenames: tuple[str, ...]) -> bytes:
+    buffer = io.BytesIO()
+    with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive:
+        for name in filenames:
+            path = repo_dir / name
+            if path.is_file():
+                archive.write(path, name)
+    return buffer.getvalue()
+
+
+def register_rdp_bundle_zip(zip_bytes: bytes) -> str:
+    token = secrets.token_urlsafe(32)
+    expires_at = time.time() + TTL_SECONDS
+    with _lock:
+        _prune_expired(time.time())
+        _entries[token] = (zip_bytes, expires_at)
+    return token
+
+
+def get_rdp_bundle_zip(token: str) -> bytes | None:
+    text = (token or "").strip()
+    if not text:
+        return None
+    now = time.time()
+    with _lock:
+        _prune_expired(now)
+        entry = _entries.get(text)
+        if entry is None:
+            return None
+        data, expires_at = entry
+        if expires_at <= now:
+            _entries.pop(text, None)
+            return None
+        return data

backend/app/services/winrm_connect.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import base64
 import re
 import socket
 from dataclasses import dataclass
@@ -11,6 +10,7 @@ from pathlib import Path
 from app.config import get_settings
 from app.models import Host
 from app.services.agent_git_release import _cache_dir, clone_or_update_repo
+from app.services.rdp_bundle_delivery import build_rdp_bundle_zip, register_rdp_bundle_zip
 
 _CLIXML_PROGRESS_NOISE = frozenset(
     {
@@ -38,7 +38,6 @@ RDP_BUNDLE_REQUIRED = frozenset(
         "Deploy-LoginMonitor.ps1",
     }
 )
-_WINRM_FILE_CHUNK_BYTES = 20_000
 
 
 class WinAdminNotConfiguredError(Exception):
@@ -344,64 +343,27 @@ def _deploy_from_staging_body(staging_path: str) -> str:
     )
 
 
-def _write_file_chunk_body(remote_path: str, b64_chunk: str, *, append: bool) -> str:
-    literal = _powershell_literal_path(remote_path)
-    if append:
-        return (
-            f"$path = '{literal}'\n"
-            f"$bytes = [Convert]::FromBase64String('{b64_chunk}')\n"
-            "$stream = [System.IO.File]::Open($path, [System.IO.FileMode]::Append)\n"
-            "try { $stream.Write($bytes, 0, $bytes.Length) } finally { $stream.Close() }\n"
-        )
+def _bundle_download_url(token: str) -> str:
+    base = get_settings().sac_public_url.rstrip("/")
+    return f"{base}/api/v1/agent/rdp-bundle/{token}"
+
+
+def _download_bundle_body(staging_path: str, bundle_url: str) -> str:
+    staging = _powershell_literal_path(staging_path)
+    url = _powershell_literal_path(bundle_url)
     return (
-        f"$path = '{literal}'\n"
-        f"$bytes = [Convert]::FromBase64String('{b64_chunk}')\n"
-        "$dir = Split-Path -LiteralPath $path -Parent\n"
-        "if (-not (Test-Path -LiteralPath $dir)) { New-Item -ItemType Directory -Path $dir -Force | Out-Null }\n"
-        "[System.IO.File]::WriteAllBytes($path, $bytes)\n"
+        f"$staging = '{staging}'\n"
+        f"$url = '{url}'\n"
+        "$zip = Join-Path $staging 'sac-rdp-bundle.zip'\n"
+        "[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12\n"
+        "Write-Output \"Downloading bundle: $url\"\n"
+        "Invoke-WebRequest -Uri $url -OutFile $zip -UseBasicParsing\n"
+        "Expand-Archive -LiteralPath $zip -DestinationPath $staging -Force\n"
+        "Remove-Item -LiteralPath $zip -Force\n"
+        "Write-Output \"Bundle extracted to $staging\"\n"
     )
 
 
-def _winrm_push_file(
-    *,
-    target: str,
-    user: str,
-    password: str,
-    remote_path: str,
-    data: bytes,
-) -> WinRmCmdResult:
-    last: WinRmCmdResult | None = None
-    if not data:
-        body = (
-            f"$path = '{_powershell_literal_path(remote_path)}'\n"
-            "$dir = Split-Path -LiteralPath $path -Parent\n"
-            "if (-not (Test-Path -LiteralPath $dir)) { New-Item -ItemType Directory -Path $dir -Force | Out-Null }\n"
-            "[System.IO.File]::WriteAllBytes($path, [byte[]]@())\n"
-        )
-        return run_winrm_ps(
-            target=target,
-            user=user,
-            password=password,
-            script=_wrap_powershell_body(body),
-            timeout_sec=120,
-        )
-    for offset in range(0, len(data), _WINRM_FILE_CHUNK_BYTES):
-        chunk = data[offset : offset + _WINRM_FILE_CHUNK_BYTES]
-        b64_chunk = base64.b64encode(chunk).decode("ascii")
-        body = _write_file_chunk_body(remote_path, b64_chunk, append=offset > 0)
-        last = run_winrm_ps(
-            target=target,
-            user=user,
-            password=password,
-            script=_wrap_powershell_body(body),
-            timeout_sec=120,
-        )
-        if not last.ok:
-            return last
-    assert last is not None
-    return last
-
-
 def _fetch_rdp_bundle_dir(repo_url: str, git_branch: str) -> Path:
     settings = get_settings()
     return clone_or_update_repo(
@@ -460,6 +422,16 @@ def run_winrm_rdp_monitor_update(
         )
 
     staging = RDP_REMOTE_STAGING
+    zip_bytes = build_rdp_bundle_zip(repo_dir, RDP_BUNDLE_FILES)
+    if not zip_bytes:
+        return WinRmCmdResult(
+            ok=False,
+            message="RDP bundle zip is empty on SAC",
+            target=target,
+        )
+    bundle_token = register_rdp_bundle_zip(zip_bytes)
+    bundle_url = _bundle_download_url(bundle_token)
+
     prep = run_winrm_ps(
         target=target,
         user=user,
@@ -470,31 +442,25 @@ def run_winrm_rdp_monitor_update(
     if not prep.ok:
         return prep
 
-    pushed: list[str] = []
-    log_parts = [prep.stdout]
-    for filename in RDP_BUNDLE_FILES:
-        local_file = repo_dir / filename
-        if not local_file.is_file():
-            continue
-        remote_file = f"{staging}\\{filename}"
-        push_result = _winrm_push_file(
+    download = run_winrm_ps(
+        target=target,
+        user=user,
+        password=password,
+        script=_wrap_powershell_body(_download_bundle_body(staging, bundle_url)),
+        timeout_sec=300,
+    )
+    if not download.ok:
+        return WinRmCmdResult(
+            ok=False,
+            message=(
+                f"Failed to download RDP bundle on client: {download.message} "
+                f"(URL: {bundle_url})"
+            ),
             target=target,
-            user=user,
-            password=password,
-            remote_path=remote_file,
-            data=local_file.read_bytes(),
+            stdout="\n\n".join(part.strip() for part in [prep.stdout, download.stdout] if part.strip()),
+            stderr=download.stderr,
+            exit_code=download.exit_code,
         )
-        log_parts.append(push_result.stdout)
-        if not push_result.ok:
-            return WinRmCmdResult(
-                ok=False,
-                message=f"Failed to push {filename} to client: {push_result.message}",
-                target=target,
-                stdout="\n".join(part.strip() for part in log_parts if part.strip()),
-                stderr=push_result.stderr,
-                exit_code=push_result.exit_code,
-            )
-        pushed.append(filename)
 
     deploy = run_winrm_ps(
         target=target,
@@ -503,10 +469,10 @@ def run_winrm_rdp_monitor_update(
         script=_wrap_powershell_body(_deploy_from_staging_body(staging)),
         timeout_sec=900,
     )
-    header = f"SAC pushed {len(pushed)} file(s) from git to {staging}"
+    header = f"SAC served bundle from git; client downloaded to {staging}"
     stdout = "\n\n".join(
         part.strip()
-        for part in [header, "\n".join(part.strip() for part in log_parts if part.strip()), deploy.stdout]
+        for part in [header, prep.stdout, download.stdout, deploy.stdout]
         if part.strip()
     )
     return WinRmCmdResult(