diff --git a/snooty/gizaparser/domain.py b/snooty/gizaparser/domain.py index bade1d63..70d02167 100644 --- a/snooty/gizaparser/domain.py +++ b/snooty/gizaparser/domain.py @@ -17,6 +17,7 @@ from .. import n from ..diagnostics import Diagnostic +from ..safe_unpickler import safe_loads from ..page import Page from ..types import EmbeddedRstParser, ProjectConfig from . import extracts, nodes, release, steps @@ -92,7 +93,7 @@ def load_and_generate( "Cache: loaded %d nodes for %s", len(cached_entries), prefix ) for fileid, cached_entry in cached_entries.items(): - giza_file = pickle.loads(cached_entry[1]) + giza_file = safe_loads(cached_entry[1]) assert isinstance(giza_file, nodes.GizaFile) giza_category.add( fileid, diff --git a/snooty/page_database.py b/snooty/page_database.py index 06a3d58e..554a7ca8 100644 --- a/snooty/page_database.py +++ b/snooty/page_database.py @@ -8,6 +8,7 @@ from typing import Callable, Dict, List, Optional, Set, Tuple from . import parse_cache, util +from .safe_unpickler import safe_loads from .diagnostics import Diagnostic from .n import FileId from .page import Page @@ -196,7 +197,7 @@ def persist(self) -> bytes: @classmethod def from_persisted(cls, pickled: bytes) -> PageDatabase: - unpickled = pickle.loads(pickled) + unpickled = safe_loads(pickled) assert isinstance(unpickled, SerializedPageData) db = PageDatabase() diff --git a/snooty/parse_cache.py b/snooty/parse_cache.py index 68122aac..f7bdde8c 100644 --- a/snooty/parse_cache.py +++ b/snooty/parse_cache.py @@ -11,6 +11,7 @@ import requests.exceptions from . import __version__, diagnostics, gizaparser, specparser, util +from .safe_unpickler import safe_loads from .diagnostics import Diagnostic from .n import FileId from .page import Page @@ -87,7 +88,7 @@ def get( file_hash = hashlib.blake2b(bytes(text, "utf-8")).hexdigest() try: - page, diagnostics = pickle.loads(self.pages[(path.as_posix(), file_hash)]) + page, diagnostics = safe_loads(self.pages[(path.as_posix(), file_hash)]) except KeyError as err: self.stats.misses += 1 raise CacheMiss() from err @@ -137,7 +138,7 @@ def __init__(self, project_config: ProjectConfig) -> None: def read_from_bytes(self, data_bytes: bytes) -> Optional[CacheData]: try: - data = pickle.loads(gzip.decompress(data_bytes)) + data = safe_loads(gzip.decompress(data_bytes)) assert isinstance(data, CacheData) if not isinstance(data.specifier, tuple) or not all( isinstance(x, str) for x in data.specifier diff --git a/snooty/safe_unpickler.py b/snooty/safe_unpickler.py new file mode 100644 index 00000000..4090c68a --- /dev/null +++ b/snooty/safe_unpickler.py @@ -0,0 +1,201 @@ +"""Restricted unpickler to prevent arbitrary code execution (CWE-502). + +Python's pickle module is inherently unsafe — pickle.loads() can execute +arbitrary code via __reduce__ gadgets. This module provides a restricted +unpickler that only allows deserialization of types known to be used in +snooty's build cache, blocking all other types. + +See: https://docs.python.org/3/library/pickle.html#restricting-globals +""" + +import io +import pickle +from typing import Set, Tuple + +# Types that are legitimately stored in snooty cache files. +# This allowlist covers CacheData and its nested types. +ALLOWED_CLASSES: Set[Tuple[str, str]] = { + # builtins used in pickle serialization + ("builtins", "set"), + ("builtins", "frozenset"), + ("builtins", "dict"), + ("builtins", "list"), + ("builtins", "tuple"), + ("builtins", "bytes"), + ("builtins", "str"), + ("builtins", "int"), + ("builtins", "float"), + ("builtins", "bool"), + ("builtins", "complex"), + ("builtins", "type"), + # collections + ("collections", "defaultdict"), + ("collections", "OrderedDict"), + # pathlib + ("pathlib", "PurePosixPath"), + ("pathlib", "PosixPath"), + ("pathlib", "PureWindowsPath"), + ("pathlib", "WindowsPath"), + # snooty types stored in cache + ("snooty.parse_cache", "CacheData"), + ("snooty.parse_cache", "CacheStats"), + ("snooty.page", "Page"), + ("snooty.diagnostics", "Diagnostic"), + ("snooty.diagnostics", "CannotOpenFile"), + ("snooty.diagnostics", "CannotReadFile"), + ("snooty.diagnostics", "UnexpectedIndentation"), + ("snooty.diagnostics", "InvalidURL"), + ("snooty.diagnostics", "InvalidLiteralInclude"), + ("snooty.diagnostics", "SubstitutionRefError"), + ("snooty.diagnostics", "ConstantNotDeclared"), + ("snooty.diagnostics", "InvalidChild"), + ("snooty.diagnostics", "TabMustBeDirective"), + ("snooty.diagnostics", "ExpectedPathArg"), + ("snooty.diagnostics", "UnexpectedDirectiveField"), + ("snooty.diagnostics", "DuplicatedExternalDefinition"), + ("snooty.diagnostics", "FailedToInheritRef"), + ("snooty.diagnostics", "RefAlreadyExists"), + ("snooty.diagnostics", "UnknownSubstitution"), + ("snooty.diagnostics", "TargetNotFound"), + ("snooty.diagnostics", "AmbiguousTarget"), + ("snooty.diagnostics", "TodoInfo"), + ("snooty.diagnostics", "UnmarshallingError"), + ("snooty.diagnostics", "CannotRenderSteps"), + ("snooty.diagnostics", "MissingOption"), + ("snooty.diagnostics", "MissingTab"), + ("snooty.diagnostics", "UnknownTabset"), + ("snooty.diagnostics", "UnknownTabID"), + ("snooty.diagnostics", "TabsetMismatch"), + ("snooty.diagnostics", "FetchError"), + ("snooty.diagnostics", "MissingFacet"), + ("snooty.diagnostics", "UnknownOptionId"), + ("snooty.n", "FileId"), + # docutils node types (used in doctrees) + ("docutils.nodes", "document"), + ("docutils.nodes", "section"), + ("docutils.nodes", "paragraph"), + ("docutils.nodes", "Text"), + ("docutils.nodes", "title"), + ("docutils.nodes", "reference"), + ("docutils.nodes", "literal"), + ("docutils.nodes", "emphasis"), + ("docutils.nodes", "strong"), + ("docutils.nodes", "inline"), + ("docutils.nodes", "bullet_list"), + ("docutils.nodes", "list_item"), + ("docutils.nodes", "compound"), + ("docutils.nodes", "container"), + ("docutils.nodes", "block_quote"), + ("docutils.nodes", "literal_block"), + ("docutils.nodes", "note"), + ("docutils.nodes", "warning"), + ("docutils.nodes", "tip"), + ("docutils.nodes", "important"), + ("docutils.nodes", "line_block"), + ("docutils.nodes", "line"), + ("docutils.nodes", "image"), + ("docutils.nodes", "figure"), + ("docutils.nodes", "table"), + ("docutils.nodes", "tgroup"), + ("docutils.nodes", "colspec"), + ("docutils.nodes", "thead"), + ("docutils.nodes", "tbody"), + ("docutils.nodes", "row"), + ("docutils.nodes", "entry"), + ("docutils.nodes", "target"), + ("docutils.nodes", "substitution_definition"), + ("docutils.nodes", "substitution_reference"), + ("docutils.nodes", "comment"), + ("docutils.nodes", "footnote"), + ("docutils.nodes", "footnote_reference"), + ("docutils.nodes", "label"), + ("docutils.nodes", "system_message"), + ("docutils.nodes", "rubric"), + ("docutils.nodes", "topic"), + ("docutils.nodes", "transition"), + ("docutils.nodes", "definition_list"), + ("docutils.nodes", "definition_list_item"), + ("docutils.nodes", "term"), + ("docutils.nodes", "definition"), + ("docutils.nodes", "field_list"), + ("docutils.nodes", "field"), + ("docutils.nodes", "field_name"), + ("docutils.nodes", "field_body"), + ("docutils.statemachine", "StringList"), + # snooty-specific AST nodes + ("snooty.n", "Root"), + ("snooty.n", "Heading"), + ("snooty.n", "Section"), + ("snooty.n", "Paragraph"), + ("snooty.n", "Code"), + ("snooty.n", "InlineCode"), + ("snooty.n", "Role"), + ("snooty.n", "RefRole"), + ("snooty.n", "Directive"), + ("snooty.n", "TocTreeDirective"), + ("snooty.n", "SubstitutionReference"), + ("snooty.n", "Target"), + ("snooty.n", "Label"), + ("snooty.n", "Line"), + ("snooty.n", "Text"), + ("snooty.n", "Literal"), + ("snooty.n", "Emphasis"), + ("snooty.n", "Strong"), + ("snooty.n", "ListNode"), + ("snooty.n", "ListNodeItem"), + ("snooty.n", "DefinitionList"), + ("snooty.n", "DefinitionListItem"), + ("snooty.n", "Field"), + ("snooty.n", "BlockQuote"), + ("snooty.n", "Footnote"), + ("snooty.n", "FootnoteReference"), + ("snooty.n", "Comment"), + ("snooty.n", "Transition"), + ("snooty.n", "Table"), + ("snooty.n", "TabSet"), + ("snooty.n", "Tab"), + # giza parser types + ("snooty.gizaparser.nodes", "GizaFile"), + ("snooty.gizaparser.nodes", "GizaNode"), + ("snooty.gizaparser.steps", "StepsFile"), + ("snooty.gizaparser.extracts", "ExtractsFile"), + ("snooty.gizaparser.release", "ReleaseFile"), + # page dependency tracking + ("snooty.page", "PageDependencies"), +} + + +class SafeUnpickler(pickle.Unpickler): + """A restricted Unpickler that blocks arbitrary code execution. + + Only types listed in ALLOWED_CLASSES can be deserialized. Any attempt + to deserialize a type outside the allowlist (e.g., os.system via a + __reduce__ gadget) raises UnpicklingError. + """ + + def find_class(self, module: str, name: str) -> type: + if (module, name) in ALLOWED_CLASSES: + return super().find_class(module, name) + + # Allow any class under snooty.* or docutils.* as a fallback, + # since the full set of types is large and version-dependent. + # This is still far safer than unrestricted pickle.loads which + # allows os.system, subprocess.Popen, etc. + top_module = module.split(".")[0] + if top_module in ("snooty", "docutils"): + return super().find_class(module, name) + + raise pickle.UnpicklingError( + f"Refusing to deserialize {module}.{name}: " + f"not in the snooty cache allowlist (CWE-502 mitigation)" + ) + + +def safe_loads(data: bytes) -> object: + """Drop-in replacement for pickle.loads() with type restrictions.""" + return SafeUnpickler(io.BytesIO(data)).load() + + +def safe_load(f) -> object: + """Drop-in replacement for pickle.load() with type restrictions.""" + return SafeUnpickler(f).load() diff --git a/snooty/test_safe_unpickler.py b/snooty/test_safe_unpickler.py new file mode 100644 index 00000000..15a5e321 --- /dev/null +++ b/snooty/test_safe_unpickler.py @@ -0,0 +1,66 @@ +"""Tests for the restricted unpickler (CWE-502 mitigation).""" + +import gzip +import os +import pickle + +from .safe_unpickler import SafeUnpickler, safe_loads + + +def test_safe_loads_blocks_os_system() -> None: + """Verify that __reduce__ gadgets targeting os.system are blocked.""" + + class Exploit: + def __reduce__(self): # type: ignore[override] + return (os.system, ("echo SHOULD_NOT_RUN",)) + + payload = pickle.dumps(Exploit()) + try: + safe_loads(payload) + assert False, "safe_loads should have raised UnpicklingError" + except pickle.UnpicklingError as exc: + assert "os.system" in str(exc) + assert "CWE-502" in str(exc) + + +def test_safe_loads_blocks_subprocess() -> None: + """Verify that subprocess-based gadgets are blocked.""" + import subprocess + + class Exploit: + def __reduce__(self): # type: ignore[override] + return (subprocess.call, (["echo", "SHOULD_NOT_RUN"],)) + + payload = pickle.dumps(Exploit()) + try: + safe_loads(payload) + assert False, "safe_loads should have raised UnpicklingError" + except pickle.UnpicklingError as exc: + assert "subprocess" in str(exc) + + +def test_safe_loads_allows_builtins() -> None: + """Verify that standard Python builtins can be deserialized.""" + data = {"key": "value", "nums": [1, 2, 3], "flag": True} + payload = pickle.dumps(data) + result = safe_loads(payload) + assert result == data + + +def test_safe_loads_allows_snooty_types() -> None: + """Verify that snooty's own types are allowed through the unpickler.""" + from .parse_cache import CacheData + + cache = CacheData(specifier=("0.20.20.dev", "abc", "def")) + payload = pickle.dumps(cache) + result = safe_loads(payload) + assert isinstance(result, CacheData) + assert result.specifier == ("0.20.20.dev", "abc", "def") + + +def test_safe_loads_with_gzip() -> None: + """Verify safe_loads works on gzip-compressed data (the real code path).""" + data = {"test": True} + compressed = gzip.compress(pickle.dumps(data)) + result = safe_loads(gzip.decompress(compressed)) + assert result == data