"""Execution hosts: this machine, or a remote server behind an SSH session."""
import logging
import re
import shutil
import stat as stat_module
import subprocess
import threading
import time
from collections import deque
from dataclasses import dataclass
from datetime import timedelta
from pathlib import Path
from uuid import uuid4
import paramiko
from remote_compression.command import FORBIDDEN_CHARS
from remote_compression.probe import probe
from remote_compression.ssh import (
RemoteExecError,
SSHConnectionError,
SSHSession,
is_localhost,
)
logger = logging.getLogger("rcomp")
#: Remote working directory, relative to the SFTP home (never absolutized:
#: OpenSSH on Windows reports homes as ``/C:/...``, invalid in a cmd argv).
WORKDIR = ".rcomp"
#: Seconds granted to get the connection back mid-batch.
RECONNECT_WINDOW = 300.0
#: A remote output whose size stalls this long is finished (or dead).
SETTLE_TIME = 60.0
#: Orphan polling interval.
POLL = 10.0
#: Seconds granted to a remote output to appear after a reconnection.
GRACE = 30.0
#: Encoder names safe to embed in a command line; anything else skips the
#: bootstrap encoder check (and will fail loudly at encode time instead).
_ENCODER_NAME = re.compile(r"[A-Za-z0-9_.-]+")
#: ffmpeg's message for an unknown encoder (exit codes vary across versions).
_UNKNOWN_ENCODER = "is not recognized"
[docs]
class BootstrapError(RuntimeError):
"""A prerequisite is missing (ffmpeg, workspace...); the batch cannot start."""
class _ConnectionLost(Exception):
"""Internal: the connection died during a phase of a remote compression."""
def __init__(self, phase, cause):
super().__init__(f"connection lost during {phase}: {cause}")
self.phase = phase
self.cause = cause
[docs]
@dataclass(frozen=True, slots=True)
class HostResult:
"""
Outcome of one compression attempt on a host.
Attributes
----------
ok: :class:`bool`
True iff the target file was produced and retrieved.
exit_status: :class:`int`, optional
ffmpeg exit status when known.
stderr_tail: :class:`str`
Bounded tail of ffmpeg stderr (diagnosis).
reason: :class:`str`
Failure category: ``ffmpeg-failed``, ``upload-failed``,
``download-failed`` or ``exec-timeout``. Empty on success.
"""
ok: bool
exit_status: int | None = None
stderr_tail: str = ""
reason: str = ""
[docs]
@dataclass(frozen=True, slots=True)
class WorkspaceEntry:
"""A regular file in the remote workspace."""
name: str
size: int
mtime: float
[docs]
def age(self, now=None):
""":class:`datetime.timedelta`: age of the file (relative to local `now`)."""
now = time.time() if now is None else now
return timedelta(seconds=max(0.0, now - self.mtime))
[docs]
@dataclass(slots=True)
class WorkspaceStatus:
"""Content of the remote workspace."""
entries: list
@property
def count(self):
""":class:`int`: number of files."""
return len(self.entries)
@property
def total_bytes(self):
""":class:`int`: cumulated size."""
return sum(e.size for e in self.entries)
[docs]
@dataclass(slots=True)
class PurgeReport:
"""Outcome of a workspace purge."""
removed: list
failed: list
dry_run: bool
@property
def freed_bytes(self):
""":class:`int`: cumulated size of the removed files."""
return sum(e.size for e in self.removed)
def _safe_suffix(path, fallback):
"""Lower-cased suffix of `path` if safe for a remote command line, else `fallback`."""
suffix = Path(path).suffix.lower()
if suffix.isascii() and not (FORBIDDEN_CHARS & set(suffix)) and " " not in suffix:
return suffix
return fallback
[docs]
def open_host(settings, config_path=None):
"""
Build the right host for the given settings.
Parameters
----------
settings: :class:`~remote_compression.settings.Settings`
Resolved settings (`hostname`, `auto_local`, `auto_purge_days`).
config_path: :class:`str`, optional
Alternative openSSH config file (mostly for tests).
Returns
-------
:class:`LocalHost` or :class:`RemoteHost`
Not yet connected: use it as a context manager, whose ``__enter__``
performs connection and fail-fast bootstrap checks.
"""
hostname = settings.hostname
if hostname.lower() == "local":
return LocalHost(encoder=settings.codec)
if settings.auto_local and is_localhost(hostname, config_path):
logger.info("Host '%s' resolves to this machine: switching to local mode.", hostname)
return LocalHost(encoder=settings.codec)
return RemoteHost(
hostname,
config_path=config_path,
auto_purge_days=settings.auto_purge_days,
encoder=settings.codec,
remote_ffmpeg=settings.remote_ffmpeg,
workdir=settings.remote_workdir,
workdir_sftp=settings.remote_workdir_sftp,
)
[docs]
class LocalHost:
"""
Compression on this very machine.
Parameters
----------
encoder: :class:`str`, optional
Encoder whose availability is verified at bootstrap.
"""
hostname = "local"
is_remote = False
def __init__(self, encoder=None):
self.encoder = encoder
def __enter__(self):
self.bootstrap()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
pass
[docs]
def bootstrap(self):
"""
Raises
------
BootstrapError
If ffmpeg is not in the local PATH, or lacks the requested encoder.
"""
if shutil.which("ffmpeg") is None:
msg = "ffmpeg not found on this machine. Install it or add it to the PATH."
raise BootstrapError(msg)
if self.encoder and _ENCODER_NAME.fullmatch(self.encoder):
result = subprocess.run(
["ffmpeg", "-hide_banner", "-h", f"encoder={self.encoder}"],
capture_output=True, check=False,
)
output = (result.stdout + result.stderr).decode("utf-8", errors="replace")
if result.returncode != 0 or _UNKNOWN_ENCODER in output:
msg = (
f"Encoder '{self.encoder}' is not available in the local ffmpeg build. "
"Pick another codec (-C) or install a fuller ffmpeg."
)
raise BootstrapError(msg)
[docs]
def compress_file(self, source, target, plan, on_progress=None):
"""
Run ffmpeg locally.
Parameters
----------
source: :class:`~pathlib.Path`
Input video.
target: :class:`~pathlib.Path`
Output video (the temporary ``comp_*`` file).
plan: :class:`~remote_compression.command.FfmpegPlan`
Command to run.
on_progress: callable, optional
Receives decoded chunks of the ffmpeg ``-progress`` stream.
Returns
-------
:class:`HostResult`
"""
argv = plan.argv(source, target)
logger.debug("Running locally: %s", subprocess.list2cmdline(argv))
proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stderr_chunks = deque(maxlen=64)
def drain_stderr():
for chunk in iter(lambda: proc.stderr.read(32768), b""):
stderr_chunks.append(chunk)
thread = threading.Thread(target=drain_stderr, daemon=True)
thread.start()
try:
for raw_line in proc.stdout:
if on_progress is not None:
on_progress(raw_line.decode("utf-8", errors="replace"))
finally:
exit_status = proc.wait()
thread.join(timeout=5)
stderr_tail = b"".join(stderr_chunks)[-4096:].decode("utf-8", errors="replace")
if exit_status != 0:
return HostResult(ok=False, exit_status=exit_status, stderr_tail=stderr_tail, reason="ffmpeg-failed")
return HostResult(ok=True, exit_status=0, stderr_tail=stderr_tail)
[docs]
class RemoteHost:
"""
Compression on a remote server, over a single SSH session.
All file operations go through SFTP (platform-agnostic: Linux and Windows
servers behave the same); only ``ffmpeg -version`` and the ffmpeg command
itself are executed remotely, with a syntax common to sh and cmd.exe.
Parameters
----------
hostname: :class:`str`
SSH alias of the server.
config_path: :class:`str`, optional
Alternative openSSH config file.
auto_purge_days: :class:`float`, optional
Age threshold of the opportunistic purge run at connection time
(None disables it).
encoder: :class:`str`, optional
Encoder whose availability is verified at bootstrap.
remote_ffmpeg: :class:`str`, optional
Explicit path of the ffmpeg binary on the server (space-free ASCII;
for hosts whose non-interactive PATH misses ffmpeg).
workdir: :class:`str`, optional
Workspace directory as seen by the shell running ffmpeg (default:
``.rcomp``, relative to the login home).
workdir_sftp: :class:`str`, optional
The same directory as seen by the SFTP channel, when the two views
differ (Synology chroots SFTP: its ``/home`` is the login home, so
``/home/.rcomp`` pairs with the default `workdir`). Defaults to
`workdir`.
session: :class:`~remote_compression.ssh.SSHSession`, optional
Injected session (tests).
"""
is_remote = True
def __init__(self, hostname, config_path=None, auto_purge_days=7.0,
encoder=None, remote_ffmpeg=None, workdir=None, workdir_sftp=None, session=None):
self.hostname = hostname
self.session = session if session is not None else SSHSession(hostname, config_path=config_path)
self.auto_purge_days = auto_purge_days
self.encoder = encoder
self.remote_ffmpeg = remote_ffmpeg
self.workdir = workdir if workdir is not None else WORKDIR
self.workdir_sftp = workdir_sftp if workdir_sftp is not None else self.workdir
# -- lifecycle -----------------------------------------------------------
def __enter__(self):
self.session.open()
self.bootstrap()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.session.close()
[docs]
def bootstrap(self):
"""Fail-fast checks: workspace, ffmpeg, then best-effort auto-purge."""
self.ensure_workspace()
self.check_ffmpeg()
if self.auto_purge_days is not None:
try:
report = self.purge_workspace(older_than=timedelta(days=self.auto_purge_days))
if report.removed:
logger.info(
"Auto-purged %d stale file(s), %d bytes, from %s:%s.",
len(report.removed), report.freed_bytes, self.hostname, self.workdir_sftp,
)
except Exception as e: # noqa: BLE001 - opportunistic purge is best-effort by design
logger.warning("Opportunistic purge of %s:%s failed: %s", self.hostname, self.workdir_sftp, e)
[docs]
def ensure_workspace(self):
"""
Create the remote workspace when missing (pure SFTP).
Raises
------
BootstrapError
If it can neither be accessed nor created.
"""
sftp = self.session.sftp
try:
sftp.stat(self.workdir_sftp)
except FileNotFoundError:
try:
sftp.mkdir(self.workdir_sftp)
logger.info("Created workspace %s on %s.", self.workdir_sftp, self.hostname)
except OSError as e:
msg = (
f"Cannot create the workspace {self.workdir_sftp} on {self.hostname}: {e}. "
"The parent folder must exist and be writable over SFTP. On Synology, "
"SFTP is chrooted to the shared folders and its /home is the login home: "
'set remote_workdir_sftp = "/home/.rcomp" in the configuration.'
)
raise BootstrapError(msg) from e
except OSError as e:
msg = f"Cannot access the workspace {self.workdir_sftp} on {self.hostname}: {e}"
raise BootstrapError(msg) from e
[docs]
def check_ffmpeg(self):
"""
Raises
------
BootstrapError
If ffmpeg fails on the server, or lacks the requested encoder.
"""
binary = self.remote_ffmpeg or "ffmpeg"
result = self.session.run(f"{binary} -version", inactivity_timeout=30.0)
if not result.ok:
detail = result.stderr_tail.strip() or result.stdout_tail.strip() or "no output"
msg = (
f"ffmpeg not found on {self.hostname} (tried '{binary}', exit {result.exit_status}): "
f"{detail}. Install ffmpeg on the server, add it to its PATH, or set "
"remote_ffmpeg in the configuration."
)
raise BootstrapError(msg)
if self.encoder and _ENCODER_NAME.fullmatch(self.encoder):
result = self.session.run(
f"{binary} -hide_banner -h encoder={self.encoder}", inactivity_timeout=30.0
)
if not result.ok or _UNKNOWN_ENCODER in result.stdout_tail + result.stderr_tail:
msg = (
f"Encoder '{self.encoder}' is not available in the ffmpeg build on "
f"{self.hostname}. Pick another codec (-C) or install a fuller ffmpeg there."
)
raise BootstrapError(msg)
# -- workspace inspection --------------------------------------------------
[docs]
def workspace_status(self):
"""
Returns
-------
:class:`WorkspaceStatus` or None
Regular files of the workspace (sub-directories and symlinks are
ignored), sorted by mtime; None when the workspace does not exist.
"""
try:
attrs = self.session.sftp.listdir_attr(self.workdir_sftp)
except FileNotFoundError:
return None
entries = [
WorkspaceEntry(name=a.filename, size=a.st_size or 0, mtime=a.st_mtime or 0.0)
for a in attrs
if stat_module.S_ISREG(a.st_mode or 0)
]
entries.sort(key=lambda e: e.mtime)
return WorkspaceStatus(entries=entries)
[docs]
def purge_workspace(self, older_than=timedelta(0), dry_run=False):
"""
Remove workspace files older than a threshold.
Parameters
----------
older_than: :class:`datetime.timedelta`
Minimal age (local clock vs remote mtimes; skew is negligible at
the thresholds in use). Zero purges everything.
dry_run: :class:`bool`
List without removing.
Returns
-------
:class:`PurgeReport`
Individual removal failures are collected, never raised.
"""
status = self.workspace_status()
if status is None:
return PurgeReport(removed=[], failed=[], dry_run=dry_run)
now = time.time()
threshold = older_than.total_seconds()
removed, failed = [], []
for entry in status.entries:
if now - entry.mtime <= threshold:
continue
if dry_run:
removed.append(entry)
continue
try:
self.session.sftp.remove(f"{self.workdir_sftp}/{entry.name}")
removed.append(entry)
except OSError as e:
failed.append((entry.name, str(e)))
return PurgeReport(removed=removed, failed=failed, dry_run=dry_run)
# -- compression -----------------------------------------------------------
[docs]
def compress_file(self, source, target, plan, on_progress=None):
"""
Upload, encode remotely, download.
On a connection loss, reconnects patiently then tries to *re-attach* to
the orphan encode (ffmpeg may well survive the disconnection on the
server) instead of blindly relaunching. At most one relaunch, always
with fresh temporary names — a zombie may still be writing to the old
target.
Parameters
----------
source: :class:`~pathlib.Path`
Local input video.
target: :class:`~pathlib.Path`
Local output (the temporary ``comp_*`` file).
plan: :class:`~remote_compression.command.FfmpegPlan`
Command to run.
on_progress: callable, optional
Receives decoded chunks of the ffmpeg ``-progress`` stream.
Returns
-------
:class:`HostResult`
Raises
------
SSHConnectionError
When the connection cannot be recovered, or is lost twice on the
same file: the batch cannot reasonably continue.
"""
for attempt in range(2):
# same physical files, two path views: SFTP operations vs the shell
# command (they differ on chrooted-SFTP servers like Synology)
src_name = f"{uuid4().hex}{_safe_suffix(source, '.src')}"
tgt_name = f"{uuid4().hex}{_safe_suffix(target, '.mkv')}"
r_source = f"{self.workdir_sftp}/{src_name}"
r_target = f"{self.workdir_sftp}/{tgt_name}"
exec_paths = (f"{self.workdir}/{src_name}", f"{self.workdir}/{tgt_name}")
try:
return self._attempt(source, target, plan, r_source, r_target, exec_paths, on_progress)
except _ConnectionLost as lost:
logger.warning(
"Connection to %s lost while processing %s (%s phase).",
self.hostname, Path(source).name, lost.phase,
)
self._patient_reconnect()
if lost.phase == "run" and self._reattach(r_target, target):
return HostResult(ok=True, exit_status=0)
self._cleanup(r_source, r_target)
if attempt == 1:
msg = (
f"Lost the connection to {self.hostname} twice on "
f"{Path(source).name}; giving up on the batch."
)
raise SSHConnectionError(msg) from lost.cause
return None # pragma: no cover - unreachable
def _attempt(self, source, target, plan, r_source, r_target, exec_paths, on_progress):
"""One full upload / encode / download cycle."""
session = self.session
try:
session.ensure_connected()
except SSHConnectionError as e:
raise _ConnectionLost("connect", e) from e
try:
session.sftp.put(str(source), r_source)
except Exception as e: # noqa: BLE001 - sorted into connection-loss vs file failure
self._raise_if_lost(e, "put")
return HostResult(ok=False, stderr_tail=str(e), reason="upload-failed")
try:
result = session.run(
plan.command(*exec_paths, binary=self.remote_ffmpeg), on_stdout=on_progress
)
except RemoteExecError as e:
self._cleanup(r_source, r_target)
return HostResult(ok=False, stderr_tail=str(e), reason="exec-timeout")
except SSHConnectionError as e:
raise _ConnectionLost("run", e) from e
try:
if result.ok and self._remote_size(r_target) > 0:
session.sftp.get(r_target, str(target))
return HostResult(ok=True, exit_status=result.exit_status, stderr_tail=result.stderr_tail)
detail = result.stderr_tail or result.stdout_tail
return HostResult(ok=False, exit_status=result.exit_status, stderr_tail=detail, reason="ffmpeg-failed")
except Exception as e: # noqa: BLE001 - sorted into connection-loss vs file failure
self._raise_if_lost(e, "get")
return HostResult(ok=False, exit_status=result.exit_status, stderr_tail=str(e), reason="download-failed")
finally:
self._cleanup(r_source, r_target)
def _raise_if_lost(self, exc, phase):
"""Translate an exception into :class:`_ConnectionLost` when the transport died."""
if isinstance(exc, _ConnectionLost):
raise exc
lost_types = (paramiko.ssh_exception.SSHException, SSHConnectionError, EOFError, ConnectionError)
if isinstance(exc, lost_types) or not self.session.is_alive():
raise _ConnectionLost(phase, exc) from exc
def _remote_size(self, r_path):
try:
return self.session.sftp.stat(r_path).st_size or 0
except FileNotFoundError:
return 0
def _cleanup(self, *r_paths):
"""Best-effort removal of remote temporaries (auto-purge catches leftovers)."""
for r_path in r_paths:
try:
self.session.sftp.remove(r_path)
except Exception: # noqa: BLE001, S110 - cleanup is best-effort by design
pass
def _patient_reconnect(self):
"""
Reconnect with a long patience window (a wifi/VPN blip mid-batch should
not kill hours of work).
Raises
------
SSHConnectionError
When :data:`RECONNECT_WINDOW` is exhausted.
"""
deadline = time.monotonic() + RECONNECT_WINDOW
delay = 2.0
while True:
try:
self.session.reconnect()
return
except SSHConnectionError as e:
if time.monotonic() + delay > deadline:
msg = f"Could not reconnect to {self.hostname} within {RECONNECT_WINDOW:.0f} s."
raise SSHConnectionError(msg) from e
logger.warning("Reconnection to %s failed, retrying in %.0f s.", self.hostname, delay)
time.sleep(delay)
delay = min(delay * 2, 30.0)
def _reattach(self, r_target, target):
"""
Try to recover the output of an orphan encode after a reconnection.
Polls the remote target: a growing size means the orphan is still
encoding (wait); a size stable for :data:`SETTLE_TIME` means it
finished or died — since its exit status is unknowable, the file is
downloaded and validated with a *local* ffprobe.
Returns
-------
:class:`bool`
True when a valid output was recovered into `target`.
Raises
------
SSHConnectionError
On a second connection loss (no third chance).
"""
logger.info("Checking whether the remote encode of %s survived...", Path(target).name)
last_size = -1
last_change = time.monotonic()
absent_deadline = time.monotonic() + GRACE
while True:
try:
size = self._remote_size(r_target)
except Exception as e: # noqa: BLE001 - sorted into connection-loss vs give-up
self._raise_if_lost_for_good(e)
return False
now = time.monotonic()
if size <= 0:
if now > absent_deadline:
logger.info("No remote output: the orphan encode died early. Relaunching.")
return False
elif size != last_size:
last_size = size
last_change = now
logger.debug("Orphan encode still writing (%d bytes).", size)
elif now - last_change > SETTLE_TIME:
break
time.sleep(POLL)
try:
self.session.sftp.get(r_target, str(target))
except Exception as e: # noqa: BLE001 - sorted into connection-loss vs give-up
self._raise_if_lost_for_good(e)
return False
finally:
self._cleanup(r_target)
if probe(target) is None:
Path(target).unlink(missing_ok=True)
logger.info("Recovered remote output is invalid; relaunching.")
return False
logger.info("Recovered the orphan encode output (%d bytes).", last_size)
return True
def _raise_if_lost_for_good(self, exc):
"""During re-attachment, a second connection loss aborts the batch."""
lost_types = (paramiko.ssh_exception.SSHException, SSHConnectionError, EOFError, ConnectionError)
if isinstance(exc, lost_types) or not self.session.is_alive():
msg = f"Lost the connection to {self.hostname} again while re-attaching; giving up on the batch."
raise SSHConnectionError(msg) from exc