← Back
Editing: fd_ops.py
"""fd-based file operations for symlink-attack mitigation. All helpers in this module use O_NOFOLLOW and dir_fd-relative syscalls so that no path-based resolution can be redirected by a concurrent symlink swap. This module is intentionally kept separate from utils/__init__.py to avoid loading these OS-specific helpers into every agent component. """ import errno import logging import os import stat from contextlib import contextmanager, suppress from pathlib import Path logger = logging.getLogger(__name__) def rmtree_fd(dir_fd) -> None: """Remove all contents of a directory using fd-relative operations. Every entry is opened with ``O_NOFOLLOW`` so symlinks inside the tree are unlinked rather than followed. The directory referenced by *dir_fd* itself is **not** removed — the caller should ``os.rmdir()`` the parent entry after this call returns. Uses an iterative approach with an explicit stack to avoid hitting Python's recursion limit on adversarial deeply-nested trees. *dir_fd* must be an open ``O_RDONLY | O_DIRECTORY`` descriptor. """ # Each stack frame is (fd, name_to_rmdir_after_close) where # name_to_rmdir_after_close is the entry name that should be # rmdir'd from the parent once this fd is fully processed. # The initial fd is managed by the caller, so its rmdir entry is None. stack = [(dir_fd, None)] try: while stack: current_fd, _ = stack[-1] pushed = False with os.scandir(current_fd) as entries: for entry in entries: if entry.is_dir(follow_symlinks=False): child_fd = os.open( entry.name, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=current_fd, ) stack.append((child_fd, entry.name)) pushed = True break # restart scan from the new directory else: os.unlink(entry.name, dir_fd=current_fd) if not pushed: # All entries in current directory have been removed. fd, name = stack.pop() if name is not None: # Close the child fd and rmdir it from the parent. os.close(fd) parent_fd, _ = stack[-1] os.rmdir(name, dir_fd=parent_fd) except BaseException: # On error, close any fds we opened (but not the caller's dir_fd). for fd, name in stack: if name is not None: os.close(fd) raise def open_dir_no_symlinks(path) -> int: """Open a directory, refusing symlinks at every path component. Walks the absolute *path* one component at a time, opening each with ``O_NOFOLLOW | O_DIRECTORY`` relative to the parent fd. This guards against symlink attacks at *any* depth in the hierarchy, not just the leaf. Returns an ``O_RDONLY`` file descriptor for the final directory. The caller is responsible for closing it. """ path = os.path.abspath(os.fspath(path)) parts = Path(path).parts # ('/', 'home', 'user', ...) fd = os.open(parts[0], os.O_RDONLY | os.O_DIRECTORY) try: for part in parts[1:]: new_fd = os.open( part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=fd, ) os.close(fd) fd = new_fd return fd except BaseException: os.close(fd) raise @contextmanager def open_nofollow(path, flags=os.O_RDONLY, *, dir_fd=None): """Open a file with O_NOFOLLOW, closing the fd on exit. Yields the raw file descriptor. Rejects symlinks at the leaf component (raises ELOOP). When *dir_fd* is provided, *path* is resolved relative to that directory descriptor. """ kw = {"dir_fd": dir_fd} if dir_fd is not None else {} fd = os.open(str(path), flags | os.O_NOFOLLOW, **kw) try: yield fd finally: os.close(fd) @contextmanager def safe_dir(path): """Open a directory with symlink protection, closing the fd on exit. Walks every path component with O_NOFOLLOW via open_dir_no_symlinks and yields the resulting fd. """ fd = open_dir_no_symlinks(path) try: yield fd finally: os.close(fd) def atomic_rewrite_fd( filename, data: bytes, *, uid, gid, allow_empty_content, permissions, dir_fd: int, ) -> bool: """dir_fd-relative implementation of atomic_rewrite. The caller opens the directory with O_NOFOLLOW before any file I/O begins. All file operations use dir_fd so that a concurrent rename of the directory to a symlink cannot redirect writes to a privileged path. """ _, basename = os.path.split(filename) # Read current content without following symlinks. try: content_fd = os.open( basename, os.O_RDONLY | os.O_NOFOLLOW, dir_fd=dir_fd ) with os.fdopen(content_fd, "rb") as f: old_content = f.read(len(data) + 1) if old_content == data: return False except FileNotFoundError: pass # file does not exist yet; will be created except OSError as exc: if exc.errno == errno.ELOOP: pass # existing entry is a symlink; overwrite it else: raise if not allow_empty_content and not data: logger.error("empty content: %r for file: %s", data, filename) return False if permissions is None: try: st = os.stat(basename, dir_fd=dir_fd, follow_symlinks=False) if stat.S_ISLNK(st.st_mode): raise OSError(errno.ELOOP, os.strerror(errno.ELOOP), basename) permissions = stat.S_IMODE(st.st_mode) except FileNotFoundError: current_umask = os.umask(0) os.umask(current_umask) permissions = 0o666 & ~current_umask # Create temp file atomically inside the directory referenced by dir_fd. # O_NOFOLLOW + O_EXCL ensures the name cannot be a pre-existing symlink. tmp_basename = None tmp_fd = -1 for _ in range(100): tmp_basename = f"{basename}_{os.urandom(4).hex()}.i360edit" try: tmp_fd = os.open( tmp_basename, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600, dir_fd=dir_fd, ) break except FileExistsError: continue else: raise FileExistsError("Could not create temporary file (100 attempts)") try: view = memoryview(data) written = 0 while written < len(data): written += os.write(tmp_fd, view[written:]) if uid is not None and gid is not None: os.chown(tmp_fd, uid, gid) os.chmod(tmp_fd, permissions) os.fsync(tmp_fd) os.close(tmp_fd) tmp_fd = -1 # Atomic rename entirely within the directory we hold open. os.rename(tmp_basename, basename, src_dir_fd=dir_fd, dst_dir_fd=dir_fd) tmp_basename = None # rename succeeded; no cleanup needed finally: if tmp_fd >= 0: os.close(tmp_fd) if tmp_basename is not None: with suppress(FileNotFoundError): os.unlink(tmp_basename, dir_fd=dir_fd) return True
Save File
Cancel