mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
100 lines
3.6 KiB
Python
100 lines
3.6 KiB
Python
"""Pure validation and command contracts for remote execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
import shlex
|
|
|
|
|
|
_REMOTE_USER = re.compile(r"[A-Za-z0-9_.-]+")
|
|
|
|
|
|
def _require_text(value: object, label: str) -> str:
|
|
if not isinstance(value, str) or not value:
|
|
raise ValueError(f"{label} must be a non-empty string")
|
|
if any(ord(character) < 32 or ord(character) == 127 for character in value):
|
|
raise ValueError(f"{label} may not contain control characters")
|
|
return value
|
|
|
|
|
|
def normalize_gpu_selection(value: str | None) -> str | None:
|
|
"""Return a canonical CUDA device selection or reject malformed input."""
|
|
if value is None:
|
|
return None
|
|
value = _require_text(value, "remote gpus").strip()
|
|
if value.lower() == "all":
|
|
return "all"
|
|
parts = value.split(",")
|
|
if any(not part.strip().isdigit() for part in parts):
|
|
raise ValueError("remote gpus must be 'all' or comma-separated non-negative integers")
|
|
return ",".join(str(int(part.strip())) for part in parts)
|
|
|
|
|
|
def _validate_remote_identity(host: str, user: str) -> None:
|
|
"""Validate the host and user portions shared by CLI and YAML inputs."""
|
|
host = _require_text(host, "remote host")
|
|
user = _require_text(user, "remote user")
|
|
if host.startswith("-") or "@" in host or any(character.isspace() for character in host):
|
|
raise ValueError("remote host must be a host name or address without user or options")
|
|
if user.startswith("-") or _REMOTE_USER.fullmatch(user) is None:
|
|
raise ValueError("remote user contains unsupported characters")
|
|
|
|
|
|
def validate_remote_settings(
|
|
*,
|
|
host: str,
|
|
user: str,
|
|
port: int,
|
|
remote_dir: str,
|
|
python: str,
|
|
gpus: str | None,
|
|
install_source: str,
|
|
) -> str | None:
|
|
"""Validate public remote settings and return canonical GPU selection."""
|
|
_validate_remote_identity(host, user)
|
|
remote_dir = _require_text(remote_dir, "remote directory")
|
|
_require_text(python, "remote Python")
|
|
_require_text(install_source, "remote install source")
|
|
|
|
if isinstance(port, bool) or not isinstance(port, int) or not 1 <= port <= 65535:
|
|
raise ValueError("remote port must be an integer from 1 through 65535")
|
|
if not remote_dir.startswith("/"):
|
|
raise ValueError("remote directory must be an absolute POSIX path")
|
|
return normalize_gpu_selection(gpus)
|
|
|
|
|
|
def parse_remote_target(remote: str) -> tuple[str, str]:
|
|
"""Parse and validate ``[USER@]HOST`` without accepting SSH options."""
|
|
remote = _require_text(remote, "remote target").strip()
|
|
if remote.count("@") > 1:
|
|
raise ValueError("remote target may contain at most one user separator")
|
|
user, separator, host = remote.partition("@")
|
|
if not separator:
|
|
user, host = "root", user
|
|
_validate_remote_identity(host, user)
|
|
return user, host
|
|
|
|
|
|
def remote_python_command(
|
|
python: str,
|
|
arguments: list[str],
|
|
*,
|
|
gpus: str | None = None,
|
|
) -> str:
|
|
"""Build one shell-safe command for the remote SSH shell."""
|
|
tokens: list[str] = []
|
|
normalized_gpus = normalize_gpu_selection(gpus)
|
|
if normalized_gpus not in (None, "all"):
|
|
tokens.extend(["env", f"CUDA_VISIBLE_DEVICES={normalized_gpus}"])
|
|
tokens.extend([_require_text(python, "remote Python"), *arguments])
|
|
return shlex.join(tokens)
|
|
|
|
|
|
def remote_scp_spec(target: str, path: str, *, directory: bool = False) -> str:
|
|
"""Quote a remote SCP path while preserving the host/path separator."""
|
|
target = _require_text(target, "SSH target")
|
|
path = _require_text(path, "remote SCP path")
|
|
if directory:
|
|
path = f"{path.rstrip('/')}/"
|
|
return f"{target}:{shlex.quote(path)}"
|