PTah
70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
"""Short-lived RDP agent bundles for WinRM clients to download."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import re
|
|
import secrets
|
|
import time
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
_UTF8_BOM = b"\xef\xbb\xbf"
|
|
_BOM_SUFFIXES = {".ps1", ".txt"}
|
|
_TOKEN_RE = re.compile(r"^[A-Za-z0-9_-]{16,128}$")
|
|
TTL_SECONDS = 600
|
|
BUNDLE_DIR = Path("/tmp/sac-rdp-bundles")
|
|
|
|
|
|
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 _prune_expired_bundles(now: float | None = None) -> None:
|
|
if not BUNDLE_DIR.is_dir():
|
|
return
|
|
cutoff = (now or time.time()) - TTL_SECONDS
|
|
for path in BUNDLE_DIR.glob("*.zip"):
|
|
try:
|
|
if path.stat().st_mtime < cutoff:
|
|
path.unlink(missing_ok=True)
|
|
except OSError:
|
|
continue
|
|
|
|
|
|
def register_rdp_bundle_zip(zip_bytes: bytes) -> str:
|
|
token = secrets.token_urlsafe(32)
|
|
BUNDLE_DIR.mkdir(parents=True, exist_ok=True)
|
|
_prune_expired_bundles()
|
|
(BUNDLE_DIR / f"{token}.zip").write_bytes(zip_bytes)
|
|
return token
|
|
|
|
|
|
def get_rdp_bundle_zip(token: str) -> bytes | None:
|
|
text = (token or "").strip()
|
|
if not _TOKEN_RE.fullmatch(text):
|
|
return None
|
|
path = BUNDLE_DIR / f"{text}.zip"
|
|
if not path.is_file():
|
|
return None
|
|
if time.time() - path.stat().st_mtime > TTL_SECONDS:
|
|
path.unlink(missing_ok=True)
|
|
return None
|
|
return path.read_bytes()
|