lodum.pickle

  1# SPDX-FileCopyrightText: 2025-present Michael R. Bernstein <zopemaven@gmail.com>
  2#
  3# SPDX-License-Identifier: Apache-2.0
  4import pickle
  5import builtins
  6import io
  7from typing import Any, Type, TypeVar
  8
  9from .core import Dumper
 10from .internal import (
 11    dump as validate_lodum_structure,
 12    DEFAULT_MAX_SIZE,
 13)
 14from .exception import DeserializationError
 15
 16T = TypeVar("T")
 17
 18# --- Safe Encoding ---
 19
 20
 21class ValidationDumper(Dumper):
 22    """A no-op dumper used only for validation."""
 23
 24    def dump_int(self, value: int) -> None:
 25        pass
 26
 27    def dump_str(self, value: str) -> None:
 28        pass
 29
 30    def dump_float(self, value: float) -> None:
 31        pass
 32
 33    def dump_bool(self, value: bool) -> None:
 34        pass
 35
 36    def dump_bytes(self, value: bytes) -> None:
 37        pass
 38
 39    def dump_list(self, value: list[Any]) -> None:
 40        pass
 41
 42    def dump_dict(self, value: dict[str, Any]) -> None:
 43        pass
 44
 45    def begin_struct(self, cls: Type[Any]) -> dict[str, Any]:
 46        return {}  # Return a dummy dict
 47
 48    def end_struct(self) -> None:
 49        pass
 50
 51
 52def dumps(obj: Any) -> bytes:
 53    """
 54    Encodes a Python object to a pickle byte string, ensuring it is safe.
 55    """
 56    validator = ValidationDumper()
 57    validate_lodum_structure(obj, validator)
 58    return pickle.dumps(obj)
 59
 60
 61# --- Safe Decoding ---
 62
 63
 64class SafeUnpickler(pickle.Unpickler):
 65    """
 66    A custom unpickler that only allows safe, lodum-enabled classes to be loaded.
 67    """
 68
 69    def find_class(self, module_name: str, class_name: str) -> Type:
 70        if "os" in module_name or "sys" in module_name or "subprocess" in module_name:
 71            raise pickle.UnpicklingError(f"Unsafe module '{module_name}' is forbidden.")
 72
 73        SAFE_BUILTINS = {
 74            "int",
 75            "float",
 76            "str",
 77            "bool",
 78            "bytes",
 79            "bytearray",
 80            "list",
 81            "tuple",
 82            "dict",
 83            "set",
 84            "frozenset",
 85            "complex",
 86            "NoneType",
 87            "type",
 88        }
 89
 90        if module_name == "builtins":
 91            if class_name in SAFE_BUILTINS and hasattr(builtins, class_name):
 92                return getattr(builtins, class_name)
 93            raise pickle.UnpicklingError(f"Unsafe builtin '{class_name}' is forbidden.")
 94
 95        if module_name == "collections" and class_name in (
 96            "defaultdict",
 97            "OrderedDict",
 98            "Counter",
 99        ):
100            import collections
101
102            return getattr(collections, class_name)
103
104        if module_name == "array" and class_name in ("array", "_array_reconstructor"):
105            import array
106
107            return getattr(array, class_name)
108
109        cls = super().find_class(module_name, class_name)
110
111        if getattr(cls, "_lodum_enabled", False):
112            return cls
113
114        raise pickle.UnpicklingError(
115            f"Attempted to unpickle a non-lodum type: {module_name}.{class_name}"
116        )
117
118
119def loads(cls: Type[T], data: bytes, max_size: int = DEFAULT_MAX_SIZE) -> T:
120    """
121    Decodes a pickle byte string to a Python object, ensuring it is safe.
122    """
123    if len(data) > max_size:
124        raise DeserializationError(
125            f"Input size ({len(data)}) exceeds maximum allowed ({max_size})"
126        )
127
128    with io.BytesIO(data) as f:
129        unpickler = SafeUnpickler(f)
130        try:
131            obj = unpickler.load()
132        except (
133            pickle.UnpicklingError,
134            AttributeError,
135            ImportError,
136            IndexError,
137            TypeError,
138        ) as e:
139            raise DeserializationError(f"Failed to unpickle data: {e}")
140
141    if not isinstance(obj, cls):
142        raise DeserializationError(
143            f"Deserialized object is of type {type(obj).__name__}, but expected {cls.__name__}"
144        )
145
146    return obj