scrypt/sandbox/vfs.py

"""The sandboxed virtual filesystem.

A pure in-memory tree. Nothing here ever touches the host filesystem —
that invariant is what makes the Warden's threats safe to carry out.
"""

from __future__ import annotations

from dataclasses import dataclass, field


class VfsError(Exception):
    pass


@dataclass
class FileNode:
    name: str
    content: str = ""
    # Locked files refuse to open until unlocked (puzzle artifacts).
    locked: bool = False
    password: str = ""
    # Archives: contents revealed by `unzip` with the right password.
    archive: dict[str, str] | None = None
    # Set by `chmod +x` (the daemon puzzle checks this).
    executable: bool = False


@dataclass
class DirNode:
    name: str
    children: dict[str, "FileNode | DirNode"] = field(default_factory=dict)


Node = FileNode | DirNode


class VFS:
    def __init__(self) -> None:
        self.root = DirNode(name="/")
        self.cwd: list[str] = ["home", "drifter"]
        self.mkdir("/home/drifter")

    # ----------------------------------------------------------- resolving

    def _parts(self, path: str) -> list[str]:
        if path.startswith("~"):
            path = "/home/drifter" + path[1:]
        parts = list(self.cwd) if not path.startswith("/") else []
        for piece in path.split("/"):
            if piece in ("", "."):
                continue
            if piece == "..":
                if parts:
                    parts.pop()
            else:
                parts.append(piece)
        return parts

    def _node_at(self, parts: list[str]) -> Node | None:
        node: Node = self.root
        for piece in parts:
            if not isinstance(node, DirNode) or piece not in node.children:
                return None
            node = node.children[piece]
        return node

    def resolve(self, path: str) -> Node | None:
        return self._node_at(self._parts(path))

    def path_of(self, parts: list[str]) -> str:
        return "/" + "/".join(parts)

    @property
    def cwd_path(self) -> str:
        return self.path_of(self.cwd)

    # ----------------------------------------------------------- mutation

    def mkdir(self, path: str) -> DirNode:
        parts = self._parts(path)
        node: Node = self.root
        for piece in parts:
            assert isinstance(node, DirNode)
            if piece not in node.children:
                node.children[piece] = DirNode(name=piece)
            node = node.children[piece]
            if not isinstance(node, DirNode):
                raise VfsError(f"{piece}: not a directory")
        return node

    def write(self, path: str, content: str, **kwargs) -> FileNode:
        parts = self._parts(path)
        if not parts:
            raise VfsError("cannot write to /")
        parent = self.mkdir(self.path_of(parts[:-1]))
        f = FileNode(name=parts[-1], content=content, **kwargs)
        parent.children[parts[-1]] = f
        return f

    def remove(self, path: str, recursive: bool = False) -> int:
        """Delete a node. Returns number of files removed (for the horror)."""
        parts = self._parts(path)
        if not parts:
            raise VfsError("cannot remove /")
        parent = self._node_at(parts[:-1])
        if not isinstance(parent, DirNode) or parts[-1] not in parent.children:
            raise VfsError(f"{path}: no such file or directory")
        node = parent.children[parts[-1]]
        if isinstance(node, DirNode):
            if not recursive:
                raise VfsError(f"{path}: is a directory")
            count = sum(1 for _ in self.iter_files(node))
        else:
            count = 1
        del parent.children[parts[-1]]
        return count

    def chdir(self, path: str) -> None:
        parts = self._parts(path)
        node = self._node_at(parts)
        if node is None:
            raise VfsError(f"{path}: no such file or directory")
        if not isinstance(node, DirNode):
            raise VfsError(f"{path}: not a directory")
        self.cwd = parts

    # ----------------------------------------------------------- reading

    def read(self, path: str) -> str:
        node = self.resolve(path)
        if node is None:
            raise VfsError(f"{path}: no such file or directory")
        if isinstance(node, DirNode):
            raise VfsError(f"{path}: is a directory")
        if node.locked:
            raise VfsError(f"{path}: permission denied")
        return node.content

    def listing(self, path: str = ".", show_hidden: bool = False) -> list[Node]:
        node = self.resolve(path)
        if node is None:
            raise VfsError(f"{path}: no such file or directory")
        if isinstance(node, FileNode):
            return [node]
        items = [
            child
            for name, child in sorted(node.children.items())
            if show_hidden or not name.startswith(".")
        ]
        return items

    def iter_files(self, node: Node | None = None, prefix: str = ""):
        """Yield (path, FileNode) for every file under node (default: root)."""
        node = node or self.root
        if isinstance(node, FileNode):
            yield prefix or node.name, node
            return
        for name, child in sorted(node.children.items()):
            child_path = f"{prefix}/{name}"
            if isinstance(child, FileNode):
                yield child_path, child
            else:
                yield from self.iter_files(child, child_path)

    def total_files(self) -> int:
        return sum(1 for _ in self.iter_files())