PTah
70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
"""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
|
|
_UTF8_BOM = b"\xef\xbb\xbf"
|
|
_BOM_SUFFIXES = {".ps1", ".txt"}
|
|
|
|
_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 _bundle_file_bytes(path: Path) -> bytes:
|
|
data = path.read_bytes()
|
|
if path.suffix.lower() in _BOM_SUFFIXES and not data.startswith(_UTF8_BOM):
|
|
return _UTF8_BOM + data
|
|
return data
|
|
|
|
|
|
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():
|
|
info = zipfile.ZipInfo(name)
|
|
info.compress_type = zipfile.ZIP_DEFLATED
|
|
info.flag_bits |= 0x800
|
|
archive.writestr(info, _bundle_file_bytes(path))
|
|
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
|