diff --git a/lib/utils.py b/lib/utils.py index 1bb5ead..486e0fc 100644 --- a/lib/utils.py +++ b/lib/utils.py @@ -32,7 +32,7 @@ import tarfile import threading import types -from typing import BinaryIO, Dict, FrozenSet, Generator, IO, Set, Tuple, cast +from typing import BinaryIO, Callable, Dict, FrozenSet, Generator, IO, Set, Tuple, cast import zipfile from absl import logging @@ -114,6 +114,67 @@ def create_pattern_for_unknowns(strings: FrozenSet[str]) -> re.Pattern[str]: }) +PICKLEMAGIC_PATTERNS = [ + ( + "FakeWarning.__new__", + re.compile( + r"FakeWarning\.__new__ called on .*?)'> with" + r" args=(?P.*), kwargs=(?P.*)" + ), + False, + ), + ( + "FakeWarning.__setstate__", + re.compile( + r"FakeWarning\.__setstate__ called on .*?)'>" + r" with unexpected state=(?P.*)" + ), + True, + ), + ( + "FakeClass.__getattr__", + re.compile( + r"FakeClass\.__getattr__ called on .*?)'>" + r" with name=(?P.*)" + ), + False, + ), + ( + "FakeClass method", + re.compile( + r"FakeClass method (?P.*?) called on .*?)'> with args=(?P.*)," + r" kwargs=(?P.*)" + ), + False, + ), + ( + "FakeClass.__call__", + re.compile( + r"FakeClass\.__call__ called on .*?)'> with" + r" args=(?P.*), kwargs=(?P.*)" + ), + False, + ), + ( + "FakeModule.__getattr__", + re.compile( + r"FakeModule\.__getattr__ called on (?P.*?) with" + r" name=(?P.*)" + ), + False, + ), + ( + "Failed to ", + re.compile( + r"Failed to (?:reduce|set state|newobj|newobj_ex|instantiate)" + r" .*?)'> with (?P.*?):" + ), + False, + ), +] + + # Creates a copy of the module def copy_module(original_name: str, new_name: str) -> types.ModuleType | None: """Copies a module and creates a new module with the same attributes. @@ -790,3 +851,142 @@ def is_sys_executable_patched() -> bool: logging.warning("Warning: sys.executable is not set to a valid interpreter.") return False + + +def _classify_item(item: str) -> Classification | None: + """Classifies a single item string into a Classification enum.""" + if not item: + return None + if item in constants.UNSAFE_STRINGS: + return Classification.UNSAFE + if item in constants.SUSPICIOUS_STRINGS: + return Classification.SUSPICIOUS + if item in constants.SAFE_STRINGS: + return Classification.SAFE + return classify_class_name(item) + + +def _parse_and_process_pattern( + line: str, + pattern: re.Pattern[str], + register_item: Callable[..., None], + is_suspicious_override: bool = False, +) -> bool: + """Parses a log line using a named group regex and processes matches directly.""" + match = pattern.search(line) + if not match: + return False + + groups = match.groupdict() + class_name = groups.get("class_name") + method_name = groups.get("method_name") + attr_name = groups.get("attr_name") + module_name = groups.get("module_name") + args = groups.get("args", "") + kwargs = groups.get("kwargs", "") + state = groups.get("state", "") + + if class_name: + register_item( + class_name, + Classification.SUSPICIOUS if is_suspicious_override else None, + ) + if method_name: + register_item(f"{class_name}.{method_name}") + if attr_name: + register_item(attr_name) + if module_name: + register_item(module_name) + + combined_args = (args or state) + " " + kwargs + if combined_args.strip(): + for group in re.findall( + r"['\"](.*?)['\"]|([a-zA-Z_][a-zA-Z0-9_.]*(?:\(.*?\))?)", combined_args + ): + for token in group: + if token: + register_item(token) + + return True + + +def categorize_picklemagic( + filtered_output: io.StringIO, +) -> Tuple[Set[str], Set[str], Set[str], Set[str]]: + """Parses and categorizes picklemagic log output.""" + results: dict[Classification, Set[str]] = { + Classification.SAFE: set(), + Classification.UNSAFE: set(), + Classification.SUSPICIOUS: set(), + Classification.UNKNOWN: set(), + } + + def register_item(item: str, override: Classification | None = None): + cls = override or _classify_item(item) + if cls is not None: + results[cls].add(item) + + lines = filtered_output.getvalue().split("\n") + for line in lines: + if not line: + continue + + if "Unsafe module/class invoked:" in line: + match = re.search( + r"Unsafe module/class invoked:\s*([a-zA-Z0-9_.]+)", line + ) + if match: + full_name = match.group(1) + results[Classification.UNSAFE].add(full_name) + if "." in full_name: + results[Classification.UNSAFE].add(full_name.split(".", 1)[0]) + continue + + if "Unknown module/class imported:" in line: + match = re.search( + r"Unknown module/class imported:\s*([a-zA-Z0-9_.]+)", line + ) + if match: + results[Classification.UNKNOWN].add(match.group(1)) + continue + + matched = False + for keyword, pattern, is_override in PICKLEMAGIC_PATTERNS: + if keyword in line: + if _parse_and_process_pattern( + line, + pattern, + register_item, + is_suspicious_override=is_override, + ): + matched = True + break + + if matched: + continue + + # Legacy fallback parsing + if line.lower().startswith("warning"): + match = re.search( + r"Unsafe module/class invoked:\s*([a-zA-Z0-9_.]+)", line + ) + if match: + full_name = match.group(1) + results[Classification.UNSAFE].add(full_name) + if "." in full_name: + results[Classification.UNSAFE].add(full_name.split(".", 1)[0]) + elif line.lower().startswith("<"): + class_args_match = ARGS_REGEX.search(line.lower()) + if class_args_match: + register_item(class_args_match.group(1)) + class_args = class_args_match.group(2) + for method_pattern in PYTHON_METHOD_PATTERNS: + for argument_find in method_pattern.findall(class_args): + register_item(argument_find) + + return ( + results[Classification.SAFE], + results[Classification.UNSAFE], + results[Classification.SUSPICIOUS], + results[Classification.UNKNOWN], + ) diff --git a/saferpickle.py b/saferpickle.py index bfcae48..ef29938 100644 --- a/saferpickle.py +++ b/saferpickle.py @@ -20,6 +20,7 @@ import functools import importlib import io +import logging as std_logging import lzma import math from multiprocessing import shared_memory @@ -460,7 +461,7 @@ def generate_ops( def get_class_instantiations( pickle_bytes: bytes | BinaryIO, -) -> tuple[io.StringIO, bool]: +) -> tuple[io.StringIO, bool, bool]: """Gets the class instantiations from a pickle file/stream. Args: @@ -471,9 +472,22 @@ def get_class_instantiations( - picklemagic_output: Suspicious function calls from picklemagic. - was_unsafe_build_blocked: A boolean indicating if a dangerous state assignment was blocked by the custom load_build hook. + - has_scan_error: A boolean indicating if the sandbox unpickler raised + an exception. """ picklemagic_output = io.StringIO() unpickler = None + has_scan_error = False + + # Configure temporary log handler to capture picklemagic logs safely + handler = std_logging.StreamHandler(picklemagic_output) + handler.setFormatter(std_logging.Formatter("%(message)s")) + logger = std_logging.getLogger("corrupy.picklemagic") + logger.addHandler(handler) + original_level = logger.level + logger.setLevel(std_logging.WARNING) + original_propagate = logger.propagate + logger.propagate = False # Handle stream seek-back if necessary original_pos = None @@ -483,7 +497,7 @@ def get_class_instantiations( original_pos = pickle_bytes.tell() pickle_stream = pickle_bytes - with contextlib.redirect_stdout(picklemagic_output): + try: try: factory = picklemagic.FakeClassFactory([], picklemagic.FakeWarning) @@ -496,22 +510,8 @@ def get_class_instantiations( unsafe_modules=constants.UNSAFE_STRINGS, ) factory.default.unpickler = unpickler - - # Monkey-patch load_build so that we don't miss - # BUILD instructions due to differing Pickle implementations. - original_load_build = unpickler.load_build - - def fixed_load_build(*unused_args): - return original_load_build() - - unpickler.load_build = fixed_load_build - unpickler.dispatch[pickle.BUILD[0]] = unpickler.load_build - unpickler.load() - # These errors are expected and should not be raised. - # Even if errors are encountered, we still get the class instantiations - # before errors occur. except ( ValueError, AttributeError, @@ -522,136 +522,91 @@ def fixed_load_build(*unused_args): EOFError, KeyError, struct.error, - ): - pass - finally: - if original_pos is not None: - pickle_bytes.seek(original_pos) + ) as e: + logging.warning("Sandbox unpickling failed: %s", e) + has_scan_error = True + finally: + logger.removeHandler(handler) + logger.setLevel(original_level) + logger.propagate = original_propagate + if original_pos is not None: + pickle_bytes.seek(original_pos) - is_build_instr_blocked = False + was_unsafe_build_blocked = False if unpickler: - is_build_instr_blocked = getattr( + was_unsafe_build_blocked = getattr( unpickler, "has_blocked_unsafe_build_instr", False ) - return picklemagic_output, is_build_instr_blocked + + return picklemagic_output, was_unsafe_build_blocked, has_scan_error def categorize_strings( filtered_output: Set[str] | io.StringIO, use_picklemagic: bool = False, ) -> ScanResults: - """Counts strings from filtered output and categorizes them. + """Counts strings from filtered output and categorizes them.""" + if use_picklemagic and isinstance(filtered_output, io.StringIO): + safe, unsafe, suspicious, unknown = utils.categorize_picklemagic( + filtered_output + ) + else: + safe, unsafe, suspicious, unknown = _categorize_genops(filtered_output) + return _reclassify_with_resolution(safe, unsafe, suspicious, unknown) - Args: - filtered_output: The series of statements filtered by string declarations. - use_picklemagic: If True, the filtered output is from picklemagic, otherwise - it is from genops or disassembly. - Returns: - A ScanResults object. - """ - - unsafe_results: Set[str] = set() +def _categorize_genops( + filtered_output: Set[str], +) -> Tuple[Set[str], Set[str], Set[str], Set[str]]: + """Helper to categorize genops output.""" safe_results: Set[str] = set() + unsafe_results: Set[str] = set() suspicious_results: Set[str] = set() unknown_results: Set[str] = set() - allow_list = config.get_allow_list() - deny_list = config.get_deny_list() - - if use_picklemagic and isinstance(filtered_output, io.StringIO): - filtered_output = filtered_output.getvalue().split("\n") # pyrefly: ignore[bad-assignment] - for picklemagic_warning in filtered_output: - if not picklemagic_warning: - continue - - picklemagic_warning_lower = picklemagic_warning.lower() - - # Printable warning sourced from every suspicious invocation of - # find_class() - if picklemagic_warning_lower.startswith("warning"): - unsafe_module_match = utils.EXTRACT_UNSAFE_MODULE_REGEX.search( - picklemagic_warning_lower - ) - if unsafe_module_match: - unsafe_results.add(unsafe_module_match.group(1)) - - # Printable warning for suspicious class instantiations - if picklemagic_warning_lower.startswith("<"): - class_args_match = utils.ARGS_REGEX.search(picklemagic_warning_lower) - - if not class_args_match: - continue - - class_name = class_args_match.group(1) - class_name_classification = utils.classify_class_name(class_name) - match class_name_classification: - case utils.Classification.SAFE: - safe_results.add(class_name) - case utils.Classification.UNSAFE: - unsafe_results.add(class_name) - case utils.Classification.SUSPICIOUS: - suspicious_results.add(class_name) - case utils.Classification.UNKNOWN: - unknown_results.add(class_name) + for line in filtered_output: + line_in_lowercase = line.lower() + unsafe_match = any( + unsafe_string in line_in_lowercase + for unsafe_string in constants.UNSAFE_STRINGS + ) and re.findall(utils.unsafe_pattern, line_in_lowercase) + safe_match = any( + safe_string in line_in_lowercase + for safe_string in constants.SAFE_STRINGS + ) and re.findall(utils.safe_pattern, line_in_lowercase) + suspicious_match = any( + suspicious_string in line_in_lowercase + for suspicious_string in constants.SUSPICIOUS_STRINGS + ) and re.findall(utils.suspicious_pattern, line_in_lowercase) + + if unsafe_match: + for match in unsafe_match: + unsafe_results.add(match) + elif safe_match: + for match in safe_match: + safe_results.add(match) + elif suspicious_match: + for match in suspicious_match: + suspicious_results.add(match) + else: + # Only check for unknown if no other categories matched + unknown_match = re.findall(utils.unknown_pattern, line_in_lowercase) + if unknown_match: + for match in unknown_match: + unknown_results.add(match) - class_args = class_args_match.group(2) + return safe_results, unsafe_results, suspicious_results, unknown_results - for method_pattern in utils.PYTHON_METHOD_PATTERNS: - argument_finds = method_pattern.findall(class_args) - if not argument_finds: - continue - for argument_find in argument_finds: - found_match = False - for unsafe_string in constants.UNSAFE_STRINGS: - if unsafe_string in argument_find: - unsafe_results.add(argument_find) - found_match = True - for safe_string in constants.SAFE_STRINGS: - if safe_string in argument_find: - safe_results.add(argument_find) - found_match = True - for suspicious_string in constants.SUSPICIOUS_STRINGS: - if suspicious_string in argument_find: - suspicious_results.add(argument_find) - found_match = True - - if not found_match and re.search( - utils.unknown_pattern, argument_find - ): - unknown_results.add(argument_find) - else: - for line in filtered_output: - line_in_lowercase = line.lower() - unsafe_match = any( - unsafe_string in line_in_lowercase - for unsafe_string in constants.UNSAFE_STRINGS - ) and re.findall(utils.unsafe_pattern, line_in_lowercase) - safe_match = any( - safe_string in line_in_lowercase - for safe_string in constants.SAFE_STRINGS - ) and re.findall(utils.safe_pattern, line_in_lowercase) - suspicious_match = any( - suspicious_string in line_in_lowercase - for suspicious_string in constants.SUSPICIOUS_STRINGS - ) and re.findall(utils.suspicious_pattern, line_in_lowercase) - - if unsafe_match: - for match in unsafe_match: - unsafe_results.add(match) - elif safe_match: - for match in safe_match: - safe_results.add(match) - elif suspicious_match: - for match in suspicious_match: - suspicious_results.add(match) - else: - # Only check for unknown if no other categories matched - unknown_match = re.findall(utils.unknown_pattern, line_in_lowercase) - if unknown_match: - for match in unknown_match: - unknown_results.add(match) +def _reclassify_with_resolution( + safe_results: Set[str], + unsafe_results: Set[str], + suspicious_results: Set[str], + unknown_results: Set[str], +) -> ScanResults: + """Helper to resolve modules and re-classify results.""" + allow_list = config.get_allow_list() + deny_list = config.get_deny_list() # Combine results for `resolve_library_modules_from_results` call. all_results = safe_results.union( @@ -744,11 +699,11 @@ def strict_security_scan(pickle_bytes: bytes | BinaryIO) -> bool: logging.debug("Failed to seek back stream before picklemagic scan.") # The below handles catching cases of unknown imports and state attacks. - instantiations_output, was_unsafe_build_blocked = get_class_instantiations( - pickle_bytes + instantiations_output, was_unsafe_build_blocked, has_scan_error = ( + get_class_instantiations(pickle_bytes) ) - if was_unsafe_build_blocked: + if was_unsafe_build_blocked or has_scan_error: return True instantiations = instantiations_output.getvalue().split("\n") @@ -819,16 +774,18 @@ def picklemagic_scan( Returns: A ScanResults object. """ - picklemagic_output, was_unsafe_build_blocked = get_class_instantiations( - pickle_bytes + picklemagic_output, was_unsafe_build_blocked, has_scan_error = ( + get_class_instantiations(pickle_bytes) ) results = categorize_strings(picklemagic_output, use_picklemagic=True) if was_unsafe_build_blocked: - # Temporary addition to increase number of suspicious results given the + # Temporary addition to increase suspicious results count given the # current scoring implementation. This will be removed in the future. results.suspicious_results.add("unsafe_state_assignment") + if has_scan_error: + results.suspicious_results.add("sandbox_unpickling_error") return results diff --git a/third_party/corrupy/picklemagic.py b/third_party/corrupy/picklemagic.py index 14fb285..75119c3 100644 --- a/third_party/corrupy/picklemagic.py +++ b/third_party/corrupy/picklemagic.py @@ -2,6 +2,9 @@ # This module provides tools for safely analyizing pickle files programmatically +import importlib +import importlib.util +import logging import sys PY3 = sys.version_info >= (3, 0) @@ -14,7 +17,7 @@ try: # only available (and needed) from 3.4 onwards. from importlib.machinery import ModuleSpec -except: +except ImportError: pass @@ -23,6 +26,9 @@ else: from cStringIO import StringIO +logger = logging.getLogger("corrupy.picklemagic") +logger.propagate = True + __all__ = [ "load", "loads", @@ -135,14 +141,45 @@ def __subclasscheck__(self, subclass): ) +def _fake_class_getattr(self, name): + logger.info( + "FakeClass.__getattr__ called on %s with name=%s", self.__class__, name + ) + return FakeClassType( + name, + (type(self),), + {}, + module=f"{self.__class__.__module__}.{self.__class__.__name__}" + if self.__class__.__name__ != "FakeClass" + else self.__class__.__module__, + )() + + +def _fake_class_call(self, *args, **kwargs): + logger.warning( + "FakeClass method %s called on %s with args=%s, kwargs=%s", + self.__class__.__name__, + self.__class__.__bases__[0] if self.__class__.__bases__ else "unknown", + args, + kwargs, + ) + return self + + # PY2 doesn't like the PY3 way of metaclasses and PY3 doesn't support the PY2 way # so we call the metaclass directly FakeClass = FakeClassType( "FakeClass", (), - {"__doc__": """ + { + "__doc__": ( + """ A barebones instance of :class:`FakeClassType`. Inherit from this to create fake classes. -"""}, +""" + ), + "__getattr__": _fake_class_getattr, + "__call__": _fake_class_call, + }, module=__name__, ) @@ -190,10 +227,11 @@ class FakeWarning(FakeClass, object): def __new__(cls, *args, **kwargs): self = FakeClass.__new__(cls) if args or kwargs: - print( - "{0} was instantiated with unexpected arguments {1}, {2}".format( - cls, args, kwargs - ) + logger.warning( + "FakeWarning.__new__ called on %s with args=%s, kwargs=%s", + cls, + args, + kwargs, ) self._new_args = args return self @@ -212,10 +250,10 @@ def __setstate__(self, state): if state: # Don't have to check for slotstate here since it's either None or a dict if not isinstance(state, dict): - print( - "{0}.__setstate__() got unexpected arguments {1}".format( - self.__class__, state - ) + logger.warning( + "FakeWarning.__setstate__ called on %s with unexpected state=%s", + self.__class__, + state, ) self._setstate_args = state else: @@ -415,6 +453,12 @@ def _remove(self): del self.__dict__[i] del sys.modules[self.__name__] + def __getattr__(self, name): + logger.info( + "FakeModule.__getattr__ called on %s with name=%s", self.__name__, name + ) + return FakeClassType(name, (FakeClass,), {}, module=self.__name__)() + def __eq__(self, other): if not hasattr(other, "__name__"): return False @@ -657,27 +701,74 @@ def __init__( self.use_copyreg = use_copyreg self.has_blocked_unsafe_build_instr = False - # Hook the BUILD opcode to our custom method. + # Hook the opcodes to our custom methods. self.dispatch[pickle.BUILD[0]] = self.load_build + self.dispatch[pickle.REDUCE[0]] = self.load_reduce + if hasattr(pickle, "INST"): + self.dispatch[pickle.INST[0]] = self.load_inst + if hasattr(pickle, "OBJ"): + self.dispatch[pickle.OBJ[0]] = self.load_obj + if hasattr(pickle, "NEWOBJ"): + self.dispatch[pickle.NEWOBJ[0]] = self.load_newobj + if hasattr(pickle, "NEWOBJ_EX"): + self.dispatch[pickle.NEWOBJ_EX[0]] = self.load_newobj_ex + self.dispatch[pickle.GLOBAL[0]] = self.load_global + if hasattr(pickle, "STACK_GLOBAL"): + self.dispatch[pickle.STACK_GLOBAL[0]] = self.load_stack_global def find_class(self, module, name): + if isinstance(module, bytes): + module = module.decode("utf-8", errors="replace") + elif not isinstance(module, str): + module = str(module) + + if isinstance(name, bytes): + name = name.decode("utf-8", errors="replace") + elif not isinstance(name, str): + name = str(name) + # __main__ can be manipulated so it's # never safe to load real classes from it. if module == "__main__": return self.class_factory(name, module) - if ( - module in self.unsafe_modules - or f"{module}.{name}" in self.unsafe_modules - ): - print(f"Warning: {module}.{name} is unsafe") + # Support submodule prefix matching for unsafe modules + is_unsafe_mod = module in self.unsafe_modules or any( + module.startswith(unsafe + ".") for unsafe in self.unsafe_modules + ) + if is_unsafe_mod or f"{module}.{name}" in self.unsafe_modules: + logger.warning("Unsafe module/class invoked: %s.%s", module, name) - if module in self.safe_modules: + sorted_safe_modules = sorted(self.safe_modules, key=len, reverse=True) + is_safe_mod = module in self.safe_modules or any( + module.startswith(safe + ".") for safe in sorted_safe_modules + ) + is_safe_class = f"{module}.{name}" in self.safe_modules + + if not is_safe_mod and not is_safe_class: + # Check if module exists spec-wise without executing arbitrary __import__ if not sys.modules.get(module): - return self.class_factory(name, module) - mod = sys.modules[module] - if not hasattr(mod, "__all__") or name in mod.__all__: - klass = getattr(mod, name) + try: + spec = importlib.util.find_spec(module) + if spec is None: + logger.warning("Unknown module/class imported: %s", module) + except ( + AttributeError, + TypeError, + ValueError, + ImportError, + ModuleNotFoundError, + ): + logger.warning("Unknown module/class imported: %s", module) + return self.class_factory(name, module) + + mod = sys.modules.get(module) + if not mod: + return self.class_factory(name, module) + + if not hasattr(mod, "__all__") or name in mod.__all__ or is_safe_class: + klass = getattr(mod, name, None) + if klass is not None: return klass return self.class_factory(name, module) @@ -688,39 +779,172 @@ def get_extension(self, code): else: return self.class_factory("extension_code_{0}".format(code), "copyreg") - def _state_contains_fake_class(self, obj): - """Recursively check if an object or its contents are FakeClass instances.""" - if isinstance(obj, FakeClass): + def _state_contains_fake_class(self, obj, visited=None): + """Recursively check if an object or its contents are FakeClass instances or types.""" + if isinstance(obj, (int, float, str, bytes, bool, type(None))): + return False + + if visited is None: + visited = set() + obj_id = id(obj) + if obj_id in visited: + return False + visited.add(obj_id) + + if isinstance(obj, (FakeClass, FakeClassType)) or ( + isinstance(obj, type) and issubclass(obj, FakeClass) + ): return True if isinstance(obj, (list, tuple, set)): - return any(self._state_contains_fake_class(item) for item in obj) + return any(self._state_contains_fake_class(item, visited) for item in obj) if isinstance(obj, dict): return any( - self._state_contains_fake_class(k) - or self._state_contains_fake_class(v) + self._state_contains_fake_class(k, visited) + or self._state_contains_fake_class(v, visited) for k, v in obj.items() ) + if hasattr(obj, "__dict__") and isinstance( + getattr(obj, "__dict__", None), dict + ): + if any( + self._state_contains_fake_class(k, visited) + or self._state_contains_fake_class(v, visited) + for k, v in obj.__dict__.items() + ): + return True + if hasattr(obj, "__slots__"): + slots = obj.__slots__ + if isinstance(slots, str): + slots = [slots] + for slot in slots: + if hasattr(obj, slot): + if self._state_contains_fake_class(getattr(obj, slot), visited): + return True return False - def load_build(self): + def load_build(self, *unused_args): """Custom handler for the BUILD opcode to prevent setting state of + a real object with a fake object (potentially dangerous). """ state = self.stack.pop() - if not state: + if state is None: return inst = self.stack[-1] + contains_fake = self._state_contains_fake_class(state) # Prevent a real object from being configured with a fake one. - if not isinstance(inst, FakeClass) and self._state_contains_fake_class( - state - ): + if not isinstance(inst, FakeClass) and contains_fake: self.has_blocked_unsafe_build_instr = True # Return to prevent inst.__setstate__(state) from being called. return - inst.__setstate__(state) + try: + inst.__setstate__(state) + except AttributeError: + logger.warning( + "Attribute __setstate__ is not available for object %s", type(inst) + ) + # Standard pickle fallback: if __setstate__ is not defined, + # update __dict__ or slots. + if isinstance(state, tuple) and len(state) == 2: + dict_state, slots_state = state + if isinstance(dict_state, dict): + inst.__dict__.update(dict_state) + if isinstance(slots_state, dict): + for slot, val in slots_state.items(): + setattr(inst, slot, val) + elif isinstance(state, dict): + logger.info("Updating __dict__ of %s with state", type(inst)) + inst.__dict__.update(state) + else: + logger.warning( + "Cannot update state of %s with state of type %s", + type(inst), + type(state), + ) + except Exception as e: # pylint: disable=broad-exception-caught + logger.warning("Failed to set state on %s: %s", type(inst), e) + logger.info("Proceeding after state failure on %s", type(inst)) + + def load_reduce(self, *unused_args): + stack = self.stack + args = stack.pop() + func = stack[-1] + + if self._state_contains_fake_class(args): + self.has_blocked_unsafe_build_instr = True + + try: + stack[-1] = func(*args) + except Exception as e: # pylint: disable=broad-exception-caught + logger.warning("Failed to reduce %s with %s: %s", func, args, e) + stack[-1] = self.class_factory("failed_reduce", "picklemagic")() + + def load_inst(self, *unused_args): + module = self.readline()[:-1].decode("utf-8", errors="replace") + name = self.readline()[:-1].decode("utf-8", errors="replace") + klass = self.find_class(module, name) + self._instantiate(klass, self.pop_mark()) + + def load_obj(self, *unused_args): + mark = self.pop_mark() + klass = mark.pop(0) + self._instantiate(klass, mark) + + def load_newobj(self, *unused_args): + args = self.stack.pop() + cls = self.stack.pop() + self._instantiate_newobj(cls, args) + + def load_newobj_ex(self, *unused_args): + kwargs = self.stack.pop() + args = self.stack.pop() + cls = self.stack.pop() + self._instantiate_newobj(cls, args, kwargs) + + def _instantiate_newobj(self, cls, args, kwargs=None): + """Internal helper for newobj/newobj_ex instantiation.""" + try: + if kwargs is not None: + obj = cls.__new__(cls, *args, **kwargs) + else: + obj = cls.__new__(cls, *args) + except Exception as e: # pylint: disable=broad-exception-caught + logger.warning("Failed to instantiate newobj %s: %s", cls, e) + obj = self.class_factory("failed_newobj", "picklemagic")() + self.stack.append(obj) + + def _instantiate(self, klass, args): + """Internal helper to instantiate a class and append to stack.""" + try: + value = klass(*args) + except Exception as e: # pylint: disable=broad-exception-caught + logger.warning("Failed to instantiate %s with %s: %s", klass, args, e) + value = self.class_factory("failed_instantiate", "picklemagic")() + self.stack.append(value) + + def load_global(self, *unused_args): + module = self.readline()[:-1].decode("utf-8", errors="replace") + name = self.readline()[:-1].decode("utf-8", errors="replace") + klass = self.find_class(module, name) + self.stack.append(klass) + + def load_stack_global(self, *unused_args): + if len(self.stack) < 2: + self.stack.append( + self.class_factory("failed_stack_global", "picklemagic")() + ) + return + name = self.stack.pop() + module = self.stack.pop() + if isinstance(name, bytes): + name = name.decode("utf-8", errors="replace") + if isinstance(module, bytes): + module = module.decode("utf-8", errors="replace") + klass = self.find_class(module, name) + self.stack.append(klass) class SafePickler(pickle.Pickler if PY2 else pickle._Pickler):