258 lines
8.0 KiB
Python
258 lines
8.0 KiB
Python
"""
|
|
Safe pickle unpickling utilities.
|
|
|
|
WARNING: Python's pickle format can execute arbitrary code during
|
|
deserialization. This module provides restricted unpicklers that only
|
|
allow known-safe classes to be loaded, preventing malicious pickle
|
|
files from executing harmful code.
|
|
|
|
See: https://docs.python.org/3/library/pickle.html#module-pickle
|
|
"""
|
|
|
|
import io
|
|
import logging
|
|
import pickle
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Whitelist definitions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Standard library modules that are safe to import during unpickling.
|
|
SAFE_MODULES: set[str] = {
|
|
"collections",
|
|
"collections.abc",
|
|
}
|
|
|
|
# Builtin names that are safe to unpickle (data types only).
|
|
# Dangerous names like eval, exec, __import__, getattr, etc. are excluded.
|
|
SAFE_BUILTINS: set[str] = {
|
|
"set",
|
|
"frozenset",
|
|
"list",
|
|
"dict",
|
|
"tuple",
|
|
"bytes",
|
|
"bytearray",
|
|
"complex",
|
|
"True",
|
|
"False",
|
|
"None",
|
|
"object",
|
|
"int",
|
|
"float",
|
|
"str",
|
|
"bool",
|
|
"slice",
|
|
"range",
|
|
"type",
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DummyClass — safe placeholder for unknown classes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class DummyClass(object):
|
|
"""
|
|
Safe placeholder class for unpickling unknown classes.
|
|
|
|
Handles all pickle reconstruction patterns:
|
|
- NEWOBJ/REDUCE with arguments (stored as _new_args/_new_kwargs)
|
|
- BUILD with dict state, (state, slotstate) tuples, or opaque state
|
|
- List-like append operations (APPENDS opcode)
|
|
"""
|
|
|
|
_state: Any
|
|
_new_args: tuple[Any, ...] | None
|
|
_new_kwargs: dict[str, Any] | None
|
|
|
|
def __new__(cls, *args: Any, **kwargs: Any) -> "DummyClass":
|
|
self = object.__new__(cls)
|
|
self._state = None
|
|
self._new_args = args if args else None
|
|
self._new_kwargs = kwargs if kwargs else None
|
|
return self
|
|
|
|
def __init__(self, *args: Any, **kwargs: Any):
|
|
# __new__ already handled storage; avoid overwriting
|
|
pass
|
|
|
|
def append(self, value: Any) -> None:
|
|
if self._state is None:
|
|
self._state = []
|
|
self._state.append(value)
|
|
|
|
def __getitem__(self, key: str) -> Any:
|
|
return self.__dict__[key]
|
|
|
|
def __eq__(self, other: object) -> bool:
|
|
if not isinstance(other, DummyClass):
|
|
return NotImplemented
|
|
return self.__dict__ == other.__dict__
|
|
|
|
def __setitem__(self, key: str, value: Any) -> None:
|
|
self.__dict__[key] = value
|
|
|
|
def __getstate__(self) -> Any:
|
|
if self._state is not None:
|
|
return self._state
|
|
# Exclude internal fields from state
|
|
return {k: v for k, v in self.__dict__.items() if not k.startswith("_new_")}
|
|
|
|
def __setstate__(self, state: Any) -> None:
|
|
# Handle (state, slotstate) tuple pattern from __slots__ classes
|
|
slotstate = None
|
|
if (
|
|
isinstance(state, tuple)
|
|
and len(state) == 2
|
|
and (state[0] is None or isinstance(state[0], dict))
|
|
and (state[1] is None or isinstance(state[1], dict))
|
|
):
|
|
state, slotstate = state
|
|
|
|
if isinstance(state, dict):
|
|
self.__dict__.update(state)
|
|
elif state is not None:
|
|
self._state = state
|
|
|
|
if slotstate:
|
|
self.__dict__.update(slotstate)
|
|
|
|
def __repr__(self) -> str:
|
|
cls = type(self)
|
|
module = getattr(cls, "__module__", "?")
|
|
name = getattr(cls, "__name__", "?")
|
|
if self._state is not None:
|
|
return f"<{module}.{name}: state={self._state!r}>"
|
|
attrs = {
|
|
k: v
|
|
for k, v in self.__dict__.items()
|
|
if not k.startswith("_new_") and k != "_state"
|
|
}
|
|
return f"<{module}.{name}: {attrs}>"
|
|
|
|
def get_code(self, **kwargs) -> str:
|
|
name = getattr(type(self), "__name__", "<?>")
|
|
logger.warning("Unrecognized node %s in decompile output", name)
|
|
return f"# <unrecognized: {name}>"
|
|
|
|
|
|
def make_dummy_class(module: str, name: str) -> type:
|
|
"""Create a named DummyClass subclass for a given module.name."""
|
|
return type(name, (DummyClass,), {"__module__": module})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# SafeUnpickler — base restricted unpickler
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class SafeUnpickler(pickle.Unpickler):
|
|
"""Base restricted unpickler with whitelist-based class loading.
|
|
|
|
Security model:
|
|
- ``renpy.*`` / ``store.*``: allowed (our fake packages)
|
|
- ``builtins`` / ``__builtin__``: only safe data types
|
|
- ``SAFE_MODULES``: standard library modules known to be safe
|
|
- Everything else: handled by ``_on_unknown_class()``
|
|
|
|
Subclasses can override ``_on_unknown_class()`` to customise behaviour
|
|
for classes outside the whitelist (raise, create placeholder, etc.).
|
|
"""
|
|
|
|
def find_class(self, module: str, name: str) -> Any:
|
|
# store.* modules -> DummyClass placeholder
|
|
if module.startswith("store"):
|
|
return make_dummy_class(module, name)
|
|
|
|
# Our fake renpy package is safe to import
|
|
if module.startswith("renpy"):
|
|
try:
|
|
return super().find_class(module, name)
|
|
except (ImportError, AttributeError):
|
|
return self._on_unknown_class(module, name)
|
|
|
|
# Whitelisted standard library modules
|
|
if module in SAFE_MODULES:
|
|
return super().find_class(module, name)
|
|
|
|
# Safe builtins only (excludes eval, exec, __import__, etc.)
|
|
if module in ("builtins", "__builtin__") and name in SAFE_BUILTINS:
|
|
return super().find_class(module, name)
|
|
|
|
return self._on_unknown_class(module, name)
|
|
|
|
def _on_unknown_class(self, module: str, name: str) -> Any:
|
|
"""Handle a class reference outside the whitelist.
|
|
|
|
Default: log a warning and return a DummyClass placeholder.
|
|
Subclasses may override to raise or return different placeholders.
|
|
"""
|
|
logger.warning(
|
|
f"Unknown class encountered: {module}.{name}. "
|
|
f"Substituting with DummyClass placeholder. "
|
|
f"If this is a legitimate Ren'Py class, please report an issue."
|
|
)
|
|
return make_dummy_class(module, name)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# RestrictedUnpickler — only allows primitive/builtin data types
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Primitive builtin types that are always safe.
|
|
_PRIMITIVE_BUILTINS = {
|
|
"dict",
|
|
"list",
|
|
"tuple",
|
|
"set",
|
|
"frozenset",
|
|
"bytes",
|
|
"bytearray",
|
|
"str",
|
|
"int",
|
|
"float",
|
|
"bool",
|
|
"complex",
|
|
"True",
|
|
"False",
|
|
"None",
|
|
}
|
|
|
|
|
|
class RestrictedUnpickler(pickle.Unpickler):
|
|
"""Strictly restricted unpickler that only allows primitive data types.
|
|
|
|
Suitable for data that should only contain basic Python types
|
|
(dict, list, tuple, int, str, bytes, etc.) with no custom classes.
|
|
|
|
Used by RPA index parsing and any other context where only
|
|
primitive data structures are expected.
|
|
"""
|
|
|
|
def find_class(self, module: str, name: str) -> Any:
|
|
if module in ("builtins", "__builtin__") and name in _PRIMITIVE_BUILTINS:
|
|
return super().find_class(module, name)
|
|
|
|
raise pickle.UnpicklingError(
|
|
f"Refusing to unpickle non-primitive class: {module}.{name}"
|
|
)
|
|
|
|
|
|
def safe_loads(data: bytes, **kwargs: Any) -> Any:
|
|
"""Safely unpickle data using SafeUnpickler.
|
|
|
|
Unknown classes are substituted with DummyClass placeholders.
|
|
"""
|
|
return SafeUnpickler(io.BytesIO(data), **kwargs).load()
|
|
|
|
|
|
def rpa_loads(data: bytes) -> Any:
|
|
"""Safely unpickle RPA index data (very restricted, primitives only)."""
|
|
return RestrictedUnpickler(io.BytesIO(data), encoding="bytes").load()
|