DataOps-platform/scripts/p3_wp14_secure_io.py

"""Fail-closed descriptor I/O: trusted root dirfd, openat parent chain and atomic output."""
from __future__ import annotations
import fcntl, json, os, secrets, stat, threading
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator
import hashlib
MAX_JSON_BYTES = 2 * 1024 * 1024
TRUSTED_ROOT = Path(__file__).resolve().parents[1]
_DIRECTORY_LOCKS: dict[tuple[int, int], threading.RLock] = {}
_DIRECTORY_LOCKS_GUARD = threading.Lock()
_DIRECTORY_LOCK_STATE = threading.local()

def _safe(st: os.stat_result, name: str) -> None:
    if not stat.S_ISREG(st.st_mode) or st.st_nlink != 1: raise ValueError(f"unsafe path: {name}")

def _parent(path: Path) -> tuple[int, str]:
    if ".." in path.parts: raise ValueError("parent-traversal paths are forbidden")
    absolute = path if path.is_absolute() else TRUSTED_ROOT / path
    try: absolute.relative_to(TRUSTED_ROOT)
    except ValueError: raise ValueError("path outside trusted root")
    # Refuse a user-controlled leaf/parent before opening; every opened parent is O_NOFOLLOW.
    fd = os.open("/", os.O_RDONLY | os.O_DIRECTORY)
    parts = absolute.parts[1:-1]
    try:
        for part in parts:
            nxt = os.open(part, os.O_RDONLY | os.O_DIRECTORY | getattr(os, "O_NOFOLLOW", 0), dir_fd=fd)
            os.close(fd); fd = nxt
        return fd, absolute.name
    except Exception:
        os.close(fd); raise


@contextmanager
def _locked_trusted_parent(parent: int) -> Iterator[None]:
    """Serialize a trusted directory across processes and re-entrantly in-process."""
    directory = os.fstat(parent)
    key = (directory.st_dev, directory.st_ino)
    with _DIRECTORY_LOCKS_GUARD:
        process_lock = _DIRECTORY_LOCKS.setdefault(key, threading.RLock())
    process_lock.acquire()
    held = getattr(_DIRECTORY_LOCK_STATE, "held", {})
    outermost = held.get(key, 0) == 0
    try:
        if outermost:
            fcntl.flock(parent, fcntl.LOCK_EX)
        held[key] = held.get(key, 0) + 1
        _DIRECTORY_LOCK_STATE.held = held
        yield
    finally:
        remaining = held.get(key, 1) - 1
        if remaining:
            held[key] = remaining
        else:
            held.pop(key, None)
            if outermost:
                fcntl.flock(parent, fcntl.LOCK_UN)
        process_lock.release()

def read_bytes_once(path: Path, *, limit: int = MAX_JSON_BYTES) -> bytes:
    parent, leaf = _parent(path)
    try:
        with _locked_trusted_parent(parent):
            before = os.stat(leaf, dir_fd=parent, follow_symlinks=False); _safe(before, str(path))
            fd = os.open(leaf, os.O_RDONLY | getattr(os,"O_NOFOLLOW",0), dir_fd=parent)
            try:
                opened=os.fstat(fd); _safe(opened,str(path))
                if (opened.st_dev,opened.st_ino)!=(before.st_dev,before.st_ino): raise ValueError("leaf swapped")
                chunks=[]; total=0
                while (chunk:=os.read(fd,65536)):
                    total+=len(chunk)
                    if total>limit: raise ValueError("input too large")
                    chunks.append(chunk)
                after=os.fstat(fd)
                if any(getattr(opened,k)!=getattr(after,k) for k in ("st_dev","st_ino","st_mode","st_nlink","st_size","st_mtime_ns","st_ctime_ns")): raise ValueError("input changed")
                return b"".join(chunks)
            finally: os.close(fd)
    finally: os.close(parent)

def read_json_once(path: Path, *, limit: int=MAX_JSON_BYTES)->object: return json.loads(read_bytes_once(path,limit=limit))

def sha256_file_once(path: Path, *, limit: int = MAX_JSON_BYTES) -> str:
    """Stream a stable regular file through SHA-256 with the same FD checks."""
    parent, leaf = _parent(path)
    try:
        with _locked_trusted_parent(parent):
            before = os.stat(leaf, dir_fd=parent, follow_symlinks=False); _safe(before, str(path))
            fd = os.open(leaf, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), dir_fd=parent)
            try:
                opened = os.fstat(fd); _safe(opened, str(path))
                if (opened.st_dev, opened.st_ino) != (before.st_dev, before.st_ino): raise ValueError("leaf swapped")
                digest = hashlib.sha256(); total = 0
                while (chunk := os.read(fd, 65536)):
                    total += len(chunk)
                    if total > limit: raise ValueError("input too large")
                    digest.update(chunk)
                after = os.fstat(fd)
                if any(getattr(opened,k)!=getattr(after,k) for k in ("st_dev","st_ino","st_mode","st_nlink","st_size","st_mtime_ns","st_ctime_ns")): raise ValueError("input changed")
                return digest.hexdigest()
            finally: os.close(fd)
    finally: os.close(parent)

def _verify_temp_matches_open_fd(parent: int, temporary: str, fd: int) -> None:
    """Refuse a name swap before publishing a temp file by rename."""
    expected = os.fstat(fd)
    _safe(expected, temporary)
    verify = os.open(temporary, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), dir_fd=parent)
    try:
        observed = os.fstat(verify)
        _safe(observed, temporary)
        if (expected.st_dev, expected.st_ino, expected.st_nlink, stat.S_IFMT(expected.st_mode)) != (
            observed.st_dev, observed.st_ino, observed.st_nlink, stat.S_IFMT(observed.st_mode)
        ):
            raise ValueError("temp swapped before atomic replace")
    finally:
        os.close(verify)

def _verify_published_leaf(parent: int, leaf: str, expected: os.stat_result, expected_digest: str) -> None:
    """Verify that rename published the exact checked temp inode, not a swapped name."""
    fd = os.open(leaf, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), dir_fd=parent)
    try:
        observed = os.fstat(fd); _safe(observed, leaf)
        if (expected.st_dev, expected.st_ino, stat.S_IFMT(expected.st_mode)) != (observed.st_dev, observed.st_ino, stat.S_IFMT(observed.st_mode)):
            raise ValueError("published leaf swapped after atomic replace")
        digest = hashlib.sha256()
        while (chunk := os.read(fd, 65536)): digest.update(chunk)
        if digest.hexdigest() != expected_digest:
            raise ValueError("published leaf digest mismatch")
    finally:
        os.close(fd)

def atomic_write_bytes(path: Path,data: bytes)->None:
    """Cooperatively publish only a fully fsynced new or previous record.

    Readers and writers take the same trusted-parent lock.  A temporary inode
    is therefore never observable by cooperative readers.  An uncooperative
    same-UID pathname mutation is detected after publication; it is not a
    trusted publisher, and the previous verified inode is restored before the
    lock is released.
    """
    parent, leaf = _parent(path)
    temporary = f".{leaf}.{secrets.token_hex(16)}.tmp"
    backup = f".{leaf}.{secrets.token_hex(16)}.bak"
    temp_exists = False
    backup_exists = False
    published = False
    old_expected: os.stat_result | None = None
    old_digest: str | None = None
    try:
        with _locked_trusted_parent(parent):
            try:
                old_expected = os.stat(leaf, dir_fd=parent, follow_symlinks=False)
                _safe(old_expected, leaf)
                old_fd = os.open(leaf, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0), dir_fd=parent)
                try:
                    if (os.fstat(old_fd).st_dev, os.fstat(old_fd).st_ino) != (old_expected.st_dev, old_expected.st_ino):
                        raise ValueError("previous leaf swapped before publication")
                    digest = hashlib.sha256()
                    while chunk := os.read(old_fd, 65536): digest.update(chunk)
                    old_digest = digest.hexdigest()
                finally:
                    os.close(old_fd)
            except FileNotFoundError:
                old_expected = None

            fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), 0o600, dir_fd=parent)
            temp_exists = True
            try:
                offset = 0
                while offset < len(data):
                    written = os.write(fd, data[offset:])
                    if written <= 0: raise OSError("short write")
                    offset += written
                os.fsync(fd)
                expected = os.fstat(fd); _safe(expected, temporary)
                _verify_temp_matches_open_fd(parent, temporary, fd)
            finally:
                os.close(fd)

            try:
                if old_expected is not None:
                    os.replace(leaf, backup, src_dir_fd=parent, dst_dir_fd=parent)
                    backup_exists = True
                    _verify_published_leaf(parent, backup, old_expected, old_digest or "")
                os.replace(temporary, leaf, src_dir_fd=parent, dst_dir_fd=parent)
                temp_exists = False
                published = True
                _verify_published_leaf(parent, leaf, expected, hashlib.sha256(data).hexdigest())
                os.fsync(parent)
            except Exception:
                if backup_exists and old_expected is not None:
                    os.replace(backup, leaf, src_dir_fd=parent, dst_dir_fd=parent)
                    backup_exists = False
                    _verify_published_leaf(parent, leaf, old_expected, old_digest or "")
                    os.fsync(parent)
                elif published:
                    os.unlink(leaf, dir_fd=parent)
                    os.fsync(parent)
                raise
            if backup_exists:
                os.unlink(backup, dir_fd=parent)
                backup_exists = False
                os.fsync(parent)
    finally:
        # A failure before publication leaves the old final record untouched.
        if temp_exists:
            try:
                os.unlink(temporary, dir_fd=parent)
                os.fsync(parent)
            except FileNotFoundError:
                pass
        if backup_exists:
            try:
                os.replace(backup, leaf, src_dir_fd=parent, dst_dir_fd=parent)
                os.fsync(parent)
            except FileNotFoundError:
                pass
        os.close(parent)
@contextmanager
def exclusive_lock(path: Path)->Iterator[str]:
    parent,_leaf=_parent(path)
    try:
        # The directory inode, never a replaceable marker filename, is the lock identity.
        with _locked_trusted_parent(parent):
            yield secrets.token_urlsafe(24)
    finally:
        os.close(parent)