TG3/tg3_local_teleop/delete_ready_episode.py

#!/usr/bin/env python3
"""Safely and idempotently remove one verified recorder ``ready`` episode.

This helper intentionally has no command-line option for either root directory.
It can only move a direct child of the fixed recorder ``ready`` directory into
the fixed ``deleting`` staging directory and remove that tombstone.  The caller
must provide both the episode id and the SHA-256 of the exact manifest that was
verified and durably published on the PS workstation.

Before recursive removal, an authorization ledger is fsynced in the fixed
``delete_ledger`` sibling directory. It survives partial tombstone deletion and
is removed only after the tree cleanup itself is durable.
"""

from __future__ import annotations

import argparse
import fcntl
import hashlib
import json
import os
import re
import stat
import sys
import uuid
from pathlib import Path
from typing import Any


BASE_ROOT = Path("/home/nvidia/tg3_data_collection")
READY_ROOT = BASE_ROOT / "ready"
DELETING_ROOT = BASE_ROOT / "deleting"
LEDGER_ROOT = BASE_ROOT / "delete_ledger"
SAFE_EPISODE_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")
SAFE_SHA256 = re.compile(r"^[0-9a-f]{64}$")
MAX_MANIFEST_BYTES = 16 * 1024 * 1024
MAX_LEDGER_BYTES = 64 * 1024
LEDGER_SCHEMA_VERSION = 1


class DeleteSafetyError(RuntimeError):
    """The requested deletion did not satisfy the fixed safety contract."""


class DeleteCollisionError(DeleteSafetyError):
    """A conflicting ready/deleting episode exists and was left untouched."""


def safe_episode_name(value: str) -> bool:
    return bool(SAFE_EPISODE_NAME.fullmatch(value)) and value not in (".", "..")


def _open_directory(path: Path) -> int:
    flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(path, flags)
    except OSError as exc:
        raise DeleteSafetyError(f"unsafe or missing directory: {path}: {exc}") from exc
    try:
        metadata = os.fstat(descriptor)
        if not stat.S_ISDIR(metadata.st_mode):
            raise DeleteSafetyError(f"not a directory: {path}")
        return descriptor
    except Exception:
        os.close(descriptor)
        raise


def _open_child_directory(parent_fd: int, name: str) -> int:
    flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(name, flags, dir_fd=parent_fd)
    except OSError as exc:
        raise DeleteSafetyError(f"unsafe or missing episode directory: {name}: {exc}") from exc
    try:
        if not stat.S_ISDIR(os.fstat(descriptor).st_mode):
            raise DeleteSafetyError(f"episode is not a directory: {name}")
        return descriptor
    except Exception:
        os.close(descriptor)
        raise


def _read_regular_file(parent_fd: int, name: str, *, max_bytes: int) -> bytes:
    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(name, flags, dir_fd=parent_fd)
    except OSError as exc:
        raise DeleteSafetyError(f"missing or unsafe regular file {name!r}: {exc}") from exc
    try:
        metadata = os.fstat(descriptor)
        if not stat.S_ISREG(metadata.st_mode):
            raise DeleteSafetyError(f"not a regular file: {name}")
        if metadata.st_size > max_bytes:
            raise DeleteSafetyError(f"file is too large: {name}")
        chunks: list[bytes] = []
        remaining = max_bytes + 1
        while remaining:
            block = os.read(descriptor, min(1024 * 1024, remaining))
            if not block:
                break
            chunks.append(block)
            remaining -= len(block)
        payload = b"".join(chunks)
        if len(payload) > max_bytes:
            raise DeleteSafetyError(f"file is too large: {name}")
        return payload
    finally:
        os.close(descriptor)


def _validate_tree_entries(directory_fd: int) -> None:
    """Require the quarantined recorder tree to contain only dirs/files."""

    for child in os.listdir(directory_fd):
        if child in (".", "..") or "/" in child or "\x00" in child:
            raise DeleteSafetyError(f"unsafe directory entry: {child!r}")
        metadata = os.stat(child, dir_fd=directory_fd, follow_symlinks=False)
        if stat.S_ISDIR(metadata.st_mode):
            child_fd = _open_child_directory(directory_fd, child)
            try:
                _validate_tree_entries(child_fd)
            finally:
                os.close(child_fd)
        elif not stat.S_ISREG(metadata.st_mode):
            raise DeleteSafetyError(f"episode contains a symlink/special entry: {child}")


def _ledger_name(episode: str, manifest_sha256: str) -> str:
    return f"{episode}.{manifest_sha256}.authorized.json"


def _ledger_payload(
    episode: str, manifest_sha256: str, tombstone: str
) -> dict[str, Any]:
    return {
        "schema_version": LEDGER_SCHEMA_VERSION,
        "state": "AUTHORIZED_DELETE",
        "episode_id": episode,
        "manifest_sha256": manifest_sha256,
        "tombstone": tombstone,
    }


def _validate_ledger(
    ledger_fd: int,
    ledger_name: str,
    *,
    episode: str,
    manifest_sha256: str,
    tombstone: str,
) -> None:
    raw = _read_regular_file(ledger_fd, ledger_name, max_bytes=MAX_LEDGER_BYTES)
    try:
        payload: Any = json.loads(raw.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise DeleteSafetyError(f"invalid deletion authorization ledger: {exc}") from exc
    if payload != _ledger_payload(episode, manifest_sha256, tombstone):
        raise DeleteSafetyError("deletion authorization ledger identity mismatch")


def _write_ledger(
    ledger_fd: int,
    ledger_name: str,
    *,
    episode: str,
    manifest_sha256: str,
    tombstone: str,
) -> None:
    payload = (
        json.dumps(
            _ledger_payload(episode, manifest_sha256, tombstone),
            sort_keys=True,
            separators=(",", ":"),
        )
        + "\n"
    ).encode("utf-8")
    temporary = f".{ledger_name}.{uuid.uuid4().hex}.tmp"
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
    descriptor = os.open(temporary, flags, 0o600, dir_fd=ledger_fd)
    try:
        view = memoryview(payload)
        while view:
            written = os.write(descriptor, view)
            if written <= 0:
                raise OSError("short write while creating deletion ledger")
            view = view[written:]
        os.fsync(descriptor)
    finally:
        os.close(descriptor)
    try:
        # A hard-link publish is atomic and refuses to replace any existing
        # sidecar. The helper flock serializes legitimate writers as well.
        os.link(
            temporary,
            ledger_name,
            src_dir_fd=ledger_fd,
            dst_dir_fd=ledger_fd,
            follow_symlinks=False,
        )
        os.fsync(ledger_fd)
    finally:
        try:
            os.unlink(temporary, dir_fd=ledger_fd)
            os.fsync(ledger_fd)
        except FileNotFoundError:
            pass


def _validate_episode(
    parent_fd: int,
    directory_name: str,
    *,
    episode: str,
    expected_manifest_sha256: str,
) -> None:
    episode_fd = _open_child_directory(parent_fd, directory_name)
    try:
        # READY must be a real regular file.  Its contents are deliberately not
        # trusted as a deletion credential; the exact manifest digest is.
        _read_regular_file(episode_fd, "READY", max_bytes=4096)
        raw_manifest = _read_regular_file(
            episode_fd, "manifest.json", max_bytes=MAX_MANIFEST_BYTES
        )
        actual_digest = hashlib.sha256(raw_manifest).hexdigest()
        if actual_digest != expected_manifest_sha256:
            raise DeleteSafetyError(
                "manifest SHA-256 mismatch: "
                f"expected {expected_manifest_sha256}, got {actual_digest}"
            )
        try:
            manifest: Any = json.loads(raw_manifest.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise DeleteSafetyError(f"invalid manifest JSON: {exc}") from exc
        if not isinstance(manifest, dict):
            raise DeleteSafetyError("manifest root is not an object")
        if manifest.get("state") != "complete":
            raise DeleteSafetyError("manifest state is not complete")
        if manifest.get("episode_id") != episode:
            raise DeleteSafetyError("manifest episode_id does not match target")
        _validate_tree_entries(episode_fd)
    finally:
        os.close(episode_fd)


def _remove_tree_at(parent_fd: int, name: str) -> None:
    """Remove one child tree without following any symbolic link."""

    directory_fd = _open_child_directory(parent_fd, name)
    try:
        for child in os.listdir(directory_fd):
            if child in (".", "..") or "/" in child or "\x00" in child:
                raise DeleteSafetyError(f"unsafe directory entry: {child!r}")
            metadata = os.stat(child, dir_fd=directory_fd, follow_symlinks=False)
            if stat.S_ISDIR(metadata.st_mode):
                _remove_tree_at(directory_fd, child)
            else:
                # Symlinks and all non-directory entries are unlinked, never
                # followed. Recorder output is expected to contain only files.
                os.unlink(child, dir_fd=directory_fd)
        os.fsync(directory_fd)
    finally:
        os.close(directory_fd)
    os.rmdir(name, dir_fd=parent_fd)


def _child_kind(parent_fd: int, name: str) -> str:
    try:
        metadata = os.stat(name, dir_fd=parent_fd, follow_symlinks=False)
    except FileNotFoundError:
        return "absent"
    if stat.S_ISDIR(metadata.st_mode):
        return "directory"
    if stat.S_ISLNK(metadata.st_mode):
        return "symlink"
    if stat.S_ISREG(metadata.st_mode):
        return "regular"
    return "other"


def delete_ready_episode(
    episode: str,
    manifest_sha256: str,
    *,
    ready_root: Path = READY_ROOT,
    deleting_root: Path = DELETING_ROOT,
    ledger_root: Path | None = None,
) -> dict[str, str]:
    """Compare, atomically quarantine, and remove exactly one ready episode.

    Root arguments are injectable only for unit tests. The CLI below always
    uses the immutable production constants, including a ledger directory that
    is outside the recursively removed ``deleting`` tree.
    """

    if not safe_episode_name(episode):
        raise DeleteSafetyError(f"unsafe episode id: {episode!r}")
    if not SAFE_SHA256.fullmatch(manifest_sha256):
        raise DeleteSafetyError("manifest SHA-256 must be 64 lowercase hex characters")

    ready_root = Path(ready_root)
    deleting_root = Path(deleting_root)
    ledger_root = Path(ledger_root or (deleting_root.parent / LEDGER_ROOT.name))
    if not (
        ready_root.parent == deleting_root.parent == ledger_root.parent
    ):
        raise DeleteSafetyError("ready, deleting, and ledger roots must be siblings")
    base_root = ready_root.parent

    # Refuse a symlinked/wrong parent even in tests. resolve(strict=True) also
    # proves that all existing parent components resolve to this exact path.
    try:
        if base_root.resolve(strict=True) != base_root.absolute():
            raise DeleteSafetyError("base root must not contain symlinks")
        if ready_root.resolve(strict=True) != ready_root.absolute():
            raise DeleteSafetyError("ready root must not contain symlinks")
    except FileNotFoundError as exc:
        raise DeleteSafetyError(f"fixed recorder root is missing: {exc}") from exc

    base_fd = _open_directory(base_root)
    ready_fd = _open_directory(ready_root)
    deleting_fd: int | None = None
    ledger_fd: int | None = None
    lock_fd: int | None = None
    try:
        base_stat = os.fstat(base_fd)
        ready_stat = os.fstat(ready_fd)
        if base_stat.st_dev != ready_stat.st_dev:
            raise DeleteSafetyError("ready root is not on the recorder filesystem")

        deleting_kind = _child_kind(base_fd, deleting_root.name)
        if deleting_kind == "absent":
            os.mkdir(deleting_root.name, mode=0o750, dir_fd=base_fd)
            os.fsync(base_fd)
        elif deleting_kind != "directory":
            raise DeleteSafetyError("deleting root is not a real directory")
        deleting_fd = _open_directory(deleting_root)
        if os.fstat(deleting_fd).st_dev != ready_stat.st_dev:
            raise DeleteSafetyError("deleting root is not on the ready filesystem")

        ledger_kind = _child_kind(base_fd, ledger_root.name)
        if ledger_kind == "absent":
            os.mkdir(ledger_root.name, mode=0o750, dir_fd=base_fd)
            os.fsync(base_fd)
        elif ledger_kind != "directory":
            raise DeleteSafetyError("delete ledger root is not a real directory")
        ledger_fd = _open_directory(ledger_root)
        if os.fstat(ledger_fd).st_dev != ready_stat.st_dev:
            raise DeleteSafetyError("delete ledger is not on the ready filesystem")

        lock_flags = os.O_RDWR | os.O_CREAT | getattr(os, "O_NOFOLLOW", 0)
        lock_fd = os.open(".delete_ready.lock", lock_flags, 0o600, dir_fd=deleting_fd)
        if not stat.S_ISREG(os.fstat(lock_fd).st_mode):
            raise DeleteSafetyError("deletion lock is not a regular file")
        fcntl.flock(lock_fd, fcntl.LOCK_EX)

        tombstone = f"{episode}.{manifest_sha256}.deleting"
        authorization = _ledger_name(episode, manifest_sha256)
        ready_kind = _child_kind(ready_fd, episode)
        tombstone_kind = _child_kind(deleting_fd, tombstone)
        authorization_kind = _child_kind(ledger_fd, authorization)

        # Any other tombstone for this episode is a manifest-identity collision.
        prefix = f"{episode}."
        collisions = [
            name
            for name in os.listdir(deleting_fd)
            if name.startswith(prefix)
            and name.endswith(".deleting")
            and name != tombstone
        ]
        if collisions:
            raise DeleteCollisionError(
                f"conflicting deleting tombstone(s) exist: {sorted(collisions)!r}"
            )
        ledger_collisions = [
            name
            for name in os.listdir(ledger_fd)
            if name.startswith(prefix)
            and name.endswith(".authorized.json")
            and name != authorization
        ]
        if ledger_collisions:
            raise DeleteCollisionError(
                "conflicting deletion authorization(s) exist: "
                f"{sorted(ledger_collisions)!r}"
            )
        if ready_kind not in ("absent", "directory"):
            raise DeleteSafetyError(f"ready target is unsafe: {ready_kind}")
        if tombstone_kind not in ("absent", "directory"):
            raise DeleteSafetyError(f"deleting target is unsafe: {tombstone_kind}")
        if authorization_kind not in ("absent", "regular"):
            raise DeleteSafetyError(
                f"deletion authorization is unsafe: {authorization_kind}"
            )
        if ready_kind == "directory" and tombstone_kind == "directory":
            raise DeleteCollisionError("ready target and deleting tombstone both exist")
        if ready_kind == "directory" and authorization_kind == "regular":
            raise DeleteCollisionError(
                "ready target exists alongside a stale deletion authorization"
            )

        if ready_kind == "directory":
            _validate_episode(
                ready_fd,
                episode,
                episode=episode,
                expected_manifest_sha256=manifest_sha256,
            )
            os.rename(
                episode,
                tombstone,
                src_dir_fd=ready_fd,
                dst_dir_fd=deleting_fd,
            )
            os.fsync(ready_fd)
            os.fsync(deleting_fd)
            # Re-validate after the rename.  If the source name was exchanged
            # between the first comparison and rename, the unexpected tree is
            # quarantined but is never unlinked.
            _validate_episode(
                deleting_fd,
                tombstone,
                episode=episode,
                expected_manifest_sha256=manifest_sha256,
            )
            _write_ledger(
                ledger_fd,
                authorization,
                episode=episode,
                manifest_sha256=manifest_sha256,
                tombstone=tombstone,
            )
            authorization_kind = "regular"
            result_state = "deleted"
        elif tombstone_kind == "directory":
            if authorization_kind == "regular":
                # The durable out-of-tree ledger remains usable even if an
                # interrupted recursive removal already removed READY and the
                # manifest from inside the tombstone.
                _validate_ledger(
                    ledger_fd,
                    authorization,
                    episode=episode,
                    manifest_sha256=manifest_sha256,
                    tombstone=tombstone,
                )
            else:
                # A crash between rename and ledger publication is recoverable
                # only while the tombstone still proves its original identity.
                _validate_episode(
                    deleting_fd,
                    tombstone,
                    episode=episode,
                    expected_manifest_sha256=manifest_sha256,
                )
                _write_ledger(
                    ledger_fd,
                    authorization,
                    episode=episode,
                    manifest_sha256=manifest_sha256,
                    tombstone=tombstone,
                )
                authorization_kind = "regular"
            result_state = "resumed_delete"
        else:
            if authorization_kind == "regular":
                # The tree was fully removed and only the final ledger cleanup
                # (or its ACK) was interrupted.
                _validate_ledger(
                    ledger_fd,
                    authorization,
                    episode=episode,
                    manifest_sha256=manifest_sha256,
                    tombstone=tombstone,
                )
                os.unlink(authorization, dir_fd=ledger_fd)
                os.fsync(ledger_fd)
            return {
                "state": "already_absent",
                "episode_id": episode,
                "manifest_sha256": manifest_sha256,
            }

        _remove_tree_at(deleting_fd, tombstone)
        os.fsync(deleting_fd)
        # Delete authorization is deliberately the final artifact removed.
        # If this unlink/ACK is lost, the next invocation validates it and
        # safely returns already_absent.
        _validate_ledger(
            ledger_fd,
            authorization,
            episode=episode,
            manifest_sha256=manifest_sha256,
            tombstone=tombstone,
        )
        os.unlink(authorization, dir_fd=ledger_fd)
        os.fsync(ledger_fd)
        return {
            "state": result_state,
            "episode_id": episode,
            "manifest_sha256": manifest_sha256,
        }
    finally:
        if lock_fd is not None:
            os.close(lock_fd)
        if deleting_fd is not None:
            os.close(deleting_fd)
        if ledger_fd is not None:
            os.close(ledger_fd)
        os.close(ready_fd)
        os.close(base_fd)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("episode_id")
    parser.add_argument("manifest_sha256")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        result = delete_ready_episode(args.episode_id, args.manifest_sha256)
    except DeleteCollisionError as exc:
        print(json.dumps({"state": "collision", "error": str(exc)}))
        return 3
    except (DeleteSafetyError, OSError) as exc:
        print(json.dumps({"state": "refused", "error": str(exc)}))
        return 2
    print(json.dumps(result, sort_keys=True))
    return 0


if __name__ == "__main__":
    sys.exit(main())