From a879b0e459aef8e1de1694d5a2bc5294eadce3e7 Mon Sep 17 00:00:00 2001 From: Vijay Panchal Date: Wed, 26 Aug 2026 00:27:04 +0530 Subject: [PATCH 1/5] Add .gitignore and commit uv project metadata --- .gitignore | 47 +++++++++++++++++++++++++++++++++++++++++++++++ .python-version | 1 + pyproject.toml | 10 ++++++++++ uv.lock | 36 ++++++++++++++++++++++++++++++++++++ 4 files changed, 94 insertions(+) create mode 100644 .gitignore create mode 100644 .python-version create mode 100644 pyproject.toml create mode 100644 uv.lock diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9b9921e --- /dev/null +++ b/.gitignore @@ -0,0 +1,47 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Python bytecode / caches +__pycache__/ +*.py[cod] + +# Virtual environments +.venv/ +venv/ +env/ +ENV/ + +# Build / distribution artifacts +build/ +dist/ +*.egg-info/ +.eggs/ + +# Test / lint / type-check caches +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +.pyre/ +.pyrefly/ +.coverage +htmlcov/ + +# IDE / OS +.vscode/ +.idea/ +*.swp +.DS_Store + +# NOTE: pyproject.toml, uv.lock, and .python-version are intentionally NOT +# ignored — commit them for reproducible uv-based development. diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..24ee5b1 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.13 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..00ce422 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,10 @@ +[project] +name = "saferpickle" +version = "0.1.0" +description = "Add your description here" +readme = "README.md" +requires-python = ">=3.13" +dependencies = [ + "absl-py>=2.5.0", + "immutabledict>=4.3.1", +] diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000..62a53d7 --- /dev/null +++ b/uv.lock @@ -0,0 +1,36 @@ +version = 1 +revision = 3 +requires-python = ">=3.13" + +[[package]] +name = "absl-py" +version = "2.5.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d0/4f/d79676ab82f2e42fc3611618139f13a9c4c31d0cff4b486982047679a802/absl_py-2.5.0.tar.gz", hash = "sha256:0c996f25c0490700fadabe6351630f6111534fa0ae252cc6d2014ea3b141135f", size = 118119, upload-time = "2026-07-03T10:57:48.157Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl", hash = "sha256:0f17b89f2a4eaaedc4f28c622998aa690564b3012a396a4ffad0821007fe03ba", size = 137410, upload-time = "2026-07-03T10:57:46.735Z" }, +] + +[[package]] +name = "immutabledict" +version = "4.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1d/e6/718471048fea0366c3e3d1df3acfd914ca66d571cdffcf6d37bbcd725708/immutabledict-4.3.1.tar.gz", hash = "sha256:f844a669106cfdc73f47b1a9da003782fb17dc955a54c80972e0d93d1c63c514", size = 7806, upload-time = "2026-02-15T10:32:34.668Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/ce/f9018bf69ae91b273b6391a095e7c93fa5e1617f25b6ba81ad4b20c9df10/immutabledict-4.3.1-py3-none-any.whl", hash = "sha256:c9facdc0ff30fdb8e35bd16532026cac472a549e182c94fa201b51b25e4bf7bf", size = 5000, upload-time = "2026-02-15T10:32:33.672Z" }, +] + +[[package]] +name = "saferpickle" +version = "0.1.0" +source = { virtual = "." } +dependencies = [ + { name = "absl-py" }, + { name = "immutabledict" }, +] + +[package.metadata] +requires-dist = [ + { name = "absl-py", specifier = ">=2.5.0" }, + { name = "immutabledict", specifier = ">=4.3.1" }, +] From 586c8263f8ab6fd65d67a42115a67a2838ea6f02 Mon Sep 17 00:00:00 2001 From: Vijay Panchal Date: Wed, 26 Aug 2026 00:55:28 +0530 Subject: [PATCH 2/5] build: migrate to Poetry and rename package to saferpickle - Replace requirements.txt and setup.py with Poetry configuration - Rename safer_pickle to saferpickle in README and code - Update CLI dependencies and type imports --- requirements.txt | 2 -- setup.py | 42 ------------------------------------------ 2 files changed, 44 deletions(-) delete mode 100644 requirements.txt delete mode 100644 setup.py diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 86adcd0..0000000 --- a/requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ -absl-py -immutabledict \ No newline at end of file diff --git a/setup.py b/setup.py deleted file mode 100644 index b1f5303..0000000 --- a/setup.py +++ /dev/null @@ -1,42 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Setup script for the saferpickle package.""" - -import setuptools - -with open("README.md", "r") as fh: - long_description = fh.read() - -setuptools.setup( - name="saferpickle", - version="0.1.0", - author="The SaferPickle Authors", - author_email="saferpickle-dev@google.com", - description="A safer alternative to Python's pickle module.", - long_description=long_description, - long_description_content_type="text/markdown", - url="https://github.com/google/saferpickle", - packages=setuptools.find_packages(), - classifiers=[ - "Programming Language :: Python :: 3", - "License :: OSI Approved :: Apache Software License", - "Operating System :: OS Independent", - ], - python_requires=">=3.6", - install_requires=[ - "absl-py", - "immutabledict", - ], -) From 907c7977c45f786cb1149267618cfbde05904d62 Mon Sep 17 00:00:00 2001 From: Vijay Panchal Date: Wed, 26 Aug 2026 00:55:46 +0530 Subject: [PATCH 3/5] refactor(package): rename to saferpickle and migrate CLI to click Update all imports from safer_pickle to saferpickle for consistent naming. Replace absl-based CLI with click and add new dependencies (click, colorama). --- README.md | 30 +- cli.py | 381 ++--- lib/__init__.py | 1 + lib/config.py | 216 +-- lib/constants.py | 386 ++--- lib/exceptions.py | 40 +- lib/utils.py | 1421 +++++++++-------- pyproject.toml | 47 +- saferpickle.py | 3026 +++++++++++++++++++------------------ tests/conftest.py | 36 + tests/test_cli.py | 25 + tests/test_load.py | 67 + tests/test_saferpickle.py | 58 + third_party/__init__.py | 1 + uv.lock | 178 ++- 15 files changed, 3222 insertions(+), 2691 deletions(-) create mode 100644 lib/__init__.py create mode 100644 tests/conftest.py create mode 100644 tests/test_cli.py create mode 100644 tests/test_load.py create mode 100644 tests/test_saferpickle.py create mode 100644 third_party/__init__.py diff --git a/README.md b/README.md index d41dea3..29950b4 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ You can use the `security_scan` function to scan a pickle file and get a report of the findings. ```py -import safer_pickle +import saferpickle import pickle class MyObject: @@ -46,7 +46,7 @@ class MyObject: my_object = MyObject("some data") pickle_bytes = pickle.dumps(my_object) -scan_results = safer_pickle.security_scan(pickle_bytes) +scan_results = saferpickle.security_scan(pickle_bytes) if scan_results["unsafe"] > 0: print("Unsafe content found!") @@ -67,47 +67,47 @@ security scan to all the standard pickle-like libraries (`pickle`, `_pickle`, application from unsafe pickles. ```py -import safer_pickle +import saferpickle import pickle -safer_pickle.hook_pickle() +saferpickle.hook_pickle() # Now, any call to pickle.load() or pickle.loads() will be protected. # For example, if you try to load a malicious pickle file, it will raise -# a safer_pickle.UnsafePickleDetectedError. +# a saferpickle.UnsafePickleDetectedError. try: # malicious_pickle_bytes is a pickle file that contains malicious code pickle.loads(malicious_pickle_bytes) -except safer_pickle.UnsafePickleDetectedError as e: +except saferpickle.UnsafePickleDetectedError as e: print(f"Blocked malicious pickle file: {e}") ``` -### 3. Use `safer_pickle.load()` and `safer_pickle.loads()` +### 3. Use `saferpickle.load()` and `saferpickle.loads()` -You can also use `safer_pickle.load()` and `safer_pickle.loads()` as direct +You can also use `saferpickle.load()` and `saferpickle.loads()` as direct replacements for `pickle.load()` and `pickle.loads()`. These functions provide more control over the security scan. ```py -import safer_pickle +import saferpickle -# This will raise a safer_pickle.UnsafePickleDetectedError if the pickle is unsafe +# This will raise a saferpickle.UnsafePickleDetectedError if the pickle is unsafe try: - obj = safer_pickle.loads(malicious_pickle_bytes) -except safer_pickle.UnsafePickleDetectedError as e: + obj = saferpickle.loads(malicious_pickle_bytes) +except saferpickle.UnsafePickleDetectedError as e: print(f"Blocked malicious pickle file: {e}") # You can also use a strict check, which is more aggressive in detecting # potentially malicious content. try: - obj = safer_pickle.loads(malicious_pickle_bytes, strict_check=True) -except safer_pickle.StrictCheckError as e: + obj = saferpickle.loads(malicious_pickle_bytes, strict_check=True) +except saferpickle.StrictCheckError as e: print(f"Blocked by strict check: {e}") # If you trust the source of the pickle file, you can bypass the security scan. -obj = safer_pickle.loads(pickle_bytes, allow_unsafe=True) +obj = saferpickle.loads(pickle_bytes, allow_unsafe=True) ``` ### 4. Command-Line Interface (CLI) diff --git a/cli.py b/cli.py index 432a879..3c3f14d 100644 --- a/cli.py +++ b/cli.py @@ -16,189 +16,215 @@ import json import os -from typing import Any, Dict, List, Sequence -from absl import app -from absl import flags +from typing import Any, Dict, List, Sequence, Set, Tuple + +from absl import app, flags + import saferpickle from lib import utils +def _scan_pickle_payload( + pickle_bytes: bytes, file_path: str | None = None +) -> Tuple[Set[str], Set[str], Set[str], Set[str]]: + """Runs picklemagic and genops scans on a single pickle payload. + + Args: + pickle_bytes: The pickle byte content to analyze. + file_path: The path to the pickle file, for streaming scan (or None when the + payload is held in memory, e.g. a zip member). + + Returns: + A tuple of (safe, unsafe, suspicious, unknown) result sets. + """ + safe_results: Set[str] = set() + unsafe_results: Set[str] = set() + suspicious_results: Set[str] = set() + unknown_results: Set[str] = set() + + # Picklemagic scan + picklemagic_results = saferpickle.picklemagic_scan(pickle_bytes) + safe_results.update(picklemagic_results.safe_results) + unsafe_results.update(picklemagic_results.unsafe_results) + suspicious_results.update(picklemagic_results.suspicious_results) + unknown_results.update(picklemagic_results.unknown_results) + + # Genops scan + genops_results = saferpickle.genops_scan( + pickle_bytes, pickle_file_path=file_path, fail_fast=False + ) + safe_results.update(genops_results.safe_results) + unsafe_results.update(genops_results.unsafe_results) + suspicious_results.update(genops_results.suspicious_results) + unknown_results.update(genops_results.unknown_results) + + return safe_results, unsafe_results, suspicious_results, unknown_results + + def security_scan_with_justifications( pickle_bytes: bytes, file_path: str | None = None ) -> Dict[str, Any]: - """Analyzes pickle byte content and returns a detailed analysis result. - - Args: - pickle_bytes: The bytes of the pickle file to analyze. - file_path: The path to the pickle file, for streaming scan. - - Returns: - A dictionary containing the analysis result. It includes: - - "classification": A string indicating if the pickle is malicious, - suspicious or benign. - - "justification": (Optional) A string providing a list of appropriate - keywords for the classification. - """ - classification = "Not supported" - justification = "Not a pickle file" - safe_results, unsafe_results, suspicious_results, unknown_results = ( - set(), - set(), - set(), - set(), - ) - - if utils.is_zip_bytes(pickle_bytes): - unzipped_files = utils.extract_zip_contents(pickle_bytes) - for unzipped_file in unzipped_files: - _, file_stream = unzipped_file - file_bytes = file_stream.read() - - if not utils.is_pickle_file(file_bytes) or not file_bytes: - continue - - pickle_bytes = file_bytes - break - - if not utils.is_pickle_file(pickle_bytes): - return {"classification": classification, "justification": justification} - - # Call the individual scan functions from SaferPickle to sets of results. - # Picklemagic Scan - picklemagic_results = saferpickle.picklemagic_scan(pickle_bytes) - - safe_results.update(picklemagic_results.safe_results) - unsafe_results.update(picklemagic_results.unsafe_results) - suspicious_results.update(picklemagic_results.suspicious_results) - unknown_results.update(picklemagic_results.unknown_results) - - # Genops Scan - genops_results = saferpickle.genops_scan( - pickle_bytes, pickle_file_path=file_path, fail_fast=False - ) - safe_results.update(genops_results.safe_results) - unsafe_results.update(genops_results.unsafe_results) - suspicious_results.update(genops_results.suspicious_results) - unknown_results.update(genops_results.unknown_results) - - final_safe_results = utils.resolve_library_modules_from_results(safe_results) - final_unsafe_results = utils.resolve_library_modules_from_results( - unsafe_results - ) - final_suspicious_results = utils.resolve_library_modules_from_results( - suspicious_results - ) - final_unknown_results = utils.resolve_library_modules_from_results( - unknown_results - ) - - # Score the results - ( - num_safe, - num_unsafe, - num_suspicious, - _, # The unknown_score is not used for classification, only reporting - ) = saferpickle.score_results( - final_safe_results, - final_unsafe_results, - final_suspicious_results, - final_unknown_results, - ) - - # Check for safety and return the results with justifications. - if saferpickle.is_unsafe(num_safe, num_unsafe, num_suspicious): - if num_unsafe > num_suspicious: - classification = "unsafe" - all_results = [] - if unsafe_results: - all_results.append( - f"malicious results: {', '.join(map(str, final_unsafe_results))}" + """Analyzes pickle byte content and returns a detailed analysis result. + + Args: + pickle_bytes: The bytes of the pickle file to analyze. + file_path: The path to the pickle file, for streaming scan. + + Returns: + A dictionary containing the analysis result. It includes: + - "classification": A string indicating if the pickle is malicious, + suspicious or benign. + - "justification": (Optional) A string providing a list of appropriate + keywords for the classification. + """ + classification = "Not supported" + justification = "Not a pickle file" + + # Collect the pickle payloads to scan. For zip archives, scan every member + # that looks like a pickle instead of only the first one. + payloads: List[Tuple[bytes, str | None]] = [] + if utils.is_zip_bytes(pickle_bytes): + for _, file_stream in utils.extract_zip_contents(pickle_bytes): + file_bytes = file_stream.read() + if file_bytes and utils.is_pickle_file(file_bytes): + # Zip members are scanned from memory, so no streaming path is passed. + payloads.append((file_bytes, None)) + elif utils.is_pickle_file(pickle_bytes): + payloads.append((pickle_bytes, file_path)) + + if not payloads: + return {"classification": classification, "justification": justification} + + safe_results: Set[str] = set() + unsafe_results: Set[str] = set() + suspicious_results: Set[str] = set() + unknown_results: Set[str] = set() + + for payload, payload_path in payloads: + payload_safe, payload_unsafe, payload_suspicious, payload_unknown = ( + _scan_pickle_payload(payload, payload_path) ) - if suspicious_results: - all_results.append( - "suspicious results:" - f" {', '.join(map(str, final_suspicious_results))}" - ) - justification = f"Found {' and '.join(all_results)}" + safe_results.update(payload_safe) + unsafe_results.update(payload_unsafe) + suspicious_results.update(payload_suspicious) + unknown_results.update(payload_unknown) + + final_safe_results = utils.resolve_library_modules_from_results(safe_results) + final_unsafe_results = utils.resolve_library_modules_from_results(unsafe_results) + final_suspicious_results = utils.resolve_library_modules_from_results( + suspicious_results + ) + final_unknown_results = utils.resolve_library_modules_from_results(unknown_results) + + # Score the results + ( + num_safe, + num_unsafe, + num_suspicious, + _, # The unknown_score is not used for classification, only reporting + ) = saferpickle.score_results( + final_safe_results, + final_unsafe_results, + final_suspicious_results, + final_unknown_results, + ) + + # Check for safety and return the results with justifications. + if saferpickle.is_unsafe(num_safe, num_unsafe, num_suspicious): + if num_unsafe > num_suspicious: + classification = "unsafe" + all_results = [] + if unsafe_results: + all_results.append( + f"malicious results: {', '.join(map(str, final_unsafe_results))}" + ) + if suspicious_results: + all_results.append( + "suspicious results:" + f" {', '.join(map(str, final_suspicious_results))}" + ) + justification = f"Found {' and '.join(all_results)}" + else: + classification = "suspicious" + justification = ( + "Found suspicious results:" + f" {', '.join(map(str, final_suspicious_results))}" + ) else: - classification = "suspicious" - justification = ( - "Found suspicious results:" - f" {', '.join(map(str, final_suspicious_results))}" - ) - else: - justification_parts = [] - if safe_results: - justification_parts.append( - f"Found safe results: {', '.join(map(str, final_safe_results))}" - ) - if unknown_results: - justification_parts.append( - f"Found unknown results: {', '.join(map(str, final_unknown_results))}" - ) - justification = " and ".join(justification_parts) - classification = "benign" - - return {"classification": classification, "justification": justification} + justification_parts = [] + if safe_results: + justification_parts.append( + f"Found safe results: {', '.join(map(str, final_safe_results))}" + ) + if unknown_results: + justification_parts.append( + f"Found unknown results: {', '.join(map(str, final_unknown_results))}" + ) + justification = " and ".join(justification_parts) + classification = "benign" + + return {"classification": classification, "justification": justification} def scan_directory(directory_path: str) -> List[Dict[str, Any]]: - """Recursively scans all files in a given directory, analyzes them for potential malicious pickles, and returns a list of results in JSON format. - - Args: - directory_path: The path to the directory to scan. - - Returns: - A list of dictionaries, where each dictionary represents the analysis - result of a file. Each dictionary includes: - - "status": "success" if the file was analyzed successfully, "error" - otherwise. - - "filename": The path to the analyzed file. - - "classification": (Optional) "benign", "suspicious" or "unsafe" - indicating the result of the analysis. - - "justification": (Optional) A string providing a reason for the - "suspicious" or "unsafe" classification. - - "error_msg": (Optional) A string describing the error if status is - "error". - """ - results = [] - if not os.path.isdir(directory_path): - results.append({ - "status": "error", - "filename": directory_path, - "error_msg": "Input path is not a valid directory.", - }) + """Recursively scans all files in a given directory, analyzes them for potential malicious pickles, and returns a list of results in JSON format. + + Args: + directory_path: The path to the directory to scan. + + Returns: + A list of dictionaries, where each dictionary represents the analysis + result of a file. Each dictionary includes: + - "status": "success" if the file was analyzed successfully, "error" + otherwise. + - "filename": The path to the analyzed file. + - "classification": (Optional) "benign", "suspicious" or "unsafe" + indicating the result of the analysis. + - "justification": (Optional) A string providing a reason for the + "suspicious" or "unsafe" classification. + - "error_msg": (Optional) A string describing the error if status is + "error". + """ + results = [] + if not os.path.isdir(directory_path): + results.append( + { + "status": "error", + "filename": directory_path, + "error_msg": "Input path is not a valid directory.", + } + ) + return results + + for root, _, files in os.walk(directory_path): + for filename in files: + file_path = os.path.join(root, filename) + try: + with open(file_path, "rb") as f: + content = f.read() + + if not content: + analysis_result = {"classification": "Not supported"} + else: + analysis_result = security_scan_with_justifications( + content, file_path=file_path + ) + + file_result = { + "status": "success", + "filename": file_path, + **analysis_result, + } + + except Exception as e: + file_result = { + "status": "error", + "filename": file_path, + "error_msg": f"Exception during parsing: {type(e).__name__} - {e}", + } + results.append(file_result) return results - for root, _, files in os.walk(directory_path): - for filename in files: - file_path = os.path.join(root, filename) - try: - with open(file_path, "rb") as f: - content = f.read() - - if not content: - analysis_result = {"classification": "Not supported"} - else: - analysis_result = security_scan_with_justifications( - content, file_path=file_path - ) - - file_result = { - "status": "success", - "filename": file_path, - **analysis_result, - } - - except (IOError, ValueError) as e: - file_result = { - "status": "error", - "filename": file_path, - "error_msg": f"Exception during parsing: {type(e).__name__} - {e}", - } - results.append(file_result) - return results - _DIRECTORY = flags.DEFINE_string( "directory", @@ -209,13 +235,18 @@ def scan_directory(directory_path: str) -> List[Dict[str, Any]]: def main(argv: Sequence[str]) -> None: - """Defines the command-line interface and executes the scan.""" - if len(argv) > 1: - raise app.UsageError("Too many command-line arguments.") + """Defines the command-line interface and executes the scan.""" + if len(argv) > 1: + raise app.UsageError("Too many command-line arguments.") + + scan_results = scan_directory(_DIRECTORY.value) + print(json.dumps(scan_results, indent=2)) + - scan_results = scan_directory(_DIRECTORY.value) - print(json.dumps(scan_results, indent=2)) +def run() -> None: + """Entry point for the safer_pickle_cli console script.""" + app.run(main) if __name__ == "__main__": - app.run(main) + app.run(main) diff --git a/lib/__init__.py b/lib/__init__.py new file mode 100644 index 0000000..e3c3f90 --- /dev/null +++ b/lib/__init__.py @@ -0,0 +1 @@ +"""Support library for saferpickle.""" diff --git a/lib/config.py b/lib/config.py index b65ce96..57664c8 100644 --- a/lib/config.py +++ b/lib/config.py @@ -21,108 +21,108 @@ class _ConfigManager: - """Manages the configuration for the safer_pickle library.""" - - def __init__(self): - self._config_path: Optional[str] = None - self._allow_list_cache: Optional[Set[str]] = None - self._deny_list_cache: Optional[Set[str]] = None - - def set_path(self, path: Optional[str]) -> None: - """Sets the path for the configuration file and resets the cache.""" - self._config_path = path - # Reset the cache to None to force a reload on the next scan. - self._allow_list_cache = None - self._deny_list_cache = None - - def get_allow_list(self) -> Set[str]: - """Loads the allow-list from the configured JSON file and caches it.""" - # Return the cached result immediately if available. - if self._allow_list_cache is not None: - return self._allow_list_cache - - # If no config path is set, cache and return an empty set. - if not self._config_path: - self._allow_list_cache = set() - return self._allow_list_cache - - try: - with open(self._config_path, "r", encoding="utf-8") as f: - config_data = json.load(f) - # Safely access the nested allow_list. - allow_list = config_data.get("safer_pickle", {}).get("allow_list", []) - if not isinstance(allow_list, list): - raise TypeError("The 'allow_list' in the config must be an array.") - self._allow_list_cache = set(allow_list) - except FileNotFoundError: - logging.warning( - "SaferPickle config file not found at %s. ", - self._config_path, - ) - self._allow_list_cache = set() - except json.JSONDecodeError as e: - raise json.JSONDecodeError( - f"Could not parse SaferPickle config at {self._config_path}: {e.msg}", - e.doc, - e.pos, - ) from e - except TypeError as e: - raise TypeError( - f"Invalid SaferPickle config at {self._config_path}: {e}" - ) from e - except IOError as e: - logging.warning( - "Could not read SaferPickle config at %s: %s. ", - self._config_path, - e, - ) - self._allow_list_cache = set() - - return self._allow_list_cache - - def get_deny_list(self) -> Set[str]: - """Loads the deny-list from the configured JSON file and caches it.""" - # Return the cached result immediately if available. - if self._deny_list_cache is not None: - return self._deny_list_cache - - # If no config path is set, cache and return an empty set. - if not self._config_path: - self._deny_list_cache = set() - return self._deny_list_cache - - try: - with open(self._config_path, "r", encoding="utf-8") as f: - config_data = json.load(f) - # Safely access the nested deny_list. - deny_list = config_data.get("safer_pickle", {}).get("deny_list", []) - if not isinstance(deny_list, list): - raise TypeError("The 'deny_list' in the config must be an array.") - self._deny_list_cache = set(deny_list) - except FileNotFoundError: - logging.warning( - "SaferPickle config file not found at %s. ", - self._config_path, - ) - self._deny_list_cache = set() - except json.JSONDecodeError as e: - raise json.JSONDecodeError( - f"Could not parse SaferPickle config at {self._config_path}: {e.msg}", - e.doc, - e.pos, - ) from e - except TypeError as e: - raise TypeError( - f"Invalid SaferPickle config at {self._config_path}: {e}" - ) from e - except IOError as e: - logging.warning( - "Could not read SaferPickle config at %s: %s. ", - self._config_path, - e, - ) - self._deny_list_cache = set() - return self._deny_list_cache + """Manages the configuration for the safer_pickle library.""" + + def __init__(self): + self._config_path: Optional[str] = None + self._allow_list_cache: Optional[Set[str]] = None + self._deny_list_cache: Optional[Set[str]] = None + + def set_path(self, path: Optional[str]) -> None: + """Sets the path for the configuration file and resets the cache.""" + self._config_path = path + # Reset the cache to None to force a reload on the next scan. + self._allow_list_cache = None + self._deny_list_cache = None + + def get_allow_list(self) -> Set[str]: + """Loads the allow-list from the configured JSON file and caches it.""" + # Return the cached result immediately if available. + if self._allow_list_cache is not None: + return self._allow_list_cache + + # If no config path is set, cache and return an empty set. + if not self._config_path: + self._allow_list_cache = set() + return self._allow_list_cache + + try: + with open(self._config_path, "r", encoding="utf-8") as f: + config_data = json.load(f) + # Safely access the nested allow_list. + allow_list = config_data.get("safer_pickle", {}).get("allow_list", []) + if not isinstance(allow_list, list): + raise TypeError("The 'allow_list' in the config must be an array.") + self._allow_list_cache = set(allow_list) + except FileNotFoundError: + logging.warning( + "SaferPickle config file not found at %s. ", + self._config_path, + ) + self._allow_list_cache = set() + except json.JSONDecodeError as e: + raise json.JSONDecodeError( + f"Could not parse SaferPickle config at {self._config_path}: {e.msg}", + e.doc, + e.pos, + ) from e + except TypeError as e: + raise TypeError( + f"Invalid SaferPickle config at {self._config_path}: {e}" + ) from e + except IOError as e: + logging.warning( + "Could not read SaferPickle config at %s: %s. ", + self._config_path, + e, + ) + self._allow_list_cache = set() + + return self._allow_list_cache + + def get_deny_list(self) -> Set[str]: + """Loads the deny-list from the configured JSON file and caches it.""" + # Return the cached result immediately if available. + if self._deny_list_cache is not None: + return self._deny_list_cache + + # If no config path is set, cache and return an empty set. + if not self._config_path: + self._deny_list_cache = set() + return self._deny_list_cache + + try: + with open(self._config_path, "r", encoding="utf-8") as f: + config_data = json.load(f) + # Safely access the nested deny_list. + deny_list = config_data.get("safer_pickle", {}).get("deny_list", []) + if not isinstance(deny_list, list): + raise TypeError("The 'deny_list' in the config must be an array.") + self._deny_list_cache = set(deny_list) + except FileNotFoundError: + logging.warning( + "SaferPickle config file not found at %s. ", + self._config_path, + ) + self._deny_list_cache = set() + except json.JSONDecodeError as e: + raise json.JSONDecodeError( + f"Could not parse SaferPickle config at {self._config_path}: {e.msg}", + e.doc, + e.pos, + ) from e + except TypeError as e: + raise TypeError( + f"Invalid SaferPickle config at {self._config_path}: {e}" + ) from e + except IOError as e: + logging.warning( + "Could not read SaferPickle config at %s: %s. ", + self._config_path, + e, + ) + self._deny_list_cache = set() + return self._deny_list_cache # A single, global instance of the configuration manager. @@ -130,15 +130,15 @@ def get_deny_list(self) -> Set[str]: def set_config_path(path: Optional[str]) -> None: - """Sets the path for the configuration file and resets the cache.""" - _config_manager.set_path(path) + """Sets the path for the configuration file and resets the cache.""" + _config_manager.set_path(path) def get_allow_list() -> Set[str]: - """Loads the allow-list from the configured JSON file and caches it.""" - return _config_manager.get_allow_list() + """Loads the allow-list from the configured JSON file and caches it.""" + return _config_manager.get_allow_list() def get_deny_list() -> Set[str]: - """Loads the deny-list from the configured JSON file and caches it.""" - return _config_manager.get_deny_list() + """Loads the deny-list from the configured JSON file and caches it.""" + return _config_manager.get_deny_list() diff --git a/lib/constants.py b/lib/constants.py index 2148ff1..d16a1bd 100644 --- a/lib/constants.py +++ b/lib/constants.py @@ -18,194 +18,201 @@ import pickletools import string from typing import FrozenSet + import immutabledict # List of globals below can be updated if additional modules are identified # Note: The strings in this list should be in the format of # "library.member" or "library" or "member" # Eg. "os.system", "loads", "system" -SUSPICIOUS_STRINGS: FrozenSet[str] = frozenset([ - "__builtin__.", - "__builtins__", - "__call__", - "__class__", - "__code__", - "__getattribute__", - "__getitem__", - "__globals__", - "__import__", - "__setstate__", - "__sub__", - "__subclasses__", - "attr", - "builtin", - "builtins", - "cgitb.lookup", - "copy_reg", - "doctest.debug_script", - "fnmatch", - "getattr", - "glob", - "hasattr", - "itertools", - "lib2to3.fixer_util.attr_chain", - "linecache", - "numpy.f2py.capi_maps.getinit", - "operator", - "pathlib", - "print", - "setattr", - "shlex", - "str.join", - "sympy.utilities.lambdify.lambdify", - "tempfile", - "test.support.get_attribute", - "types.FunctionType", - "unittest.mock._dot_lookup", - "unittest.mock._importer", - "xml.etree", - "xmlrpc.server.resolve_dotted_attribute", -]) +SUSPICIOUS_STRINGS: FrozenSet[str] = frozenset( + [ + "__builtin__.", + "__builtins__", + "__call__", + "__class__", + "__code__", + "__getattribute__", + "__getitem__", + "__globals__", + "__import__", + "__setstate__", + "__sub__", + "__subclasses__", + "attr", + "builtin", + "builtins", + "cgitb.lookup", + "copy_reg", + "doctest.debug_script", + "fnmatch", + "getattr", + "glob", + "hasattr", + "itertools", + "lib2to3.fixer_util.attr_chain", + "linecache", + "numpy.f2py.capi_maps.getinit", + "operator", + "pathlib", + "print", + "setattr", + "shlex", + "str.join", + "sympy.utilities.lambdify.lambdify", + "tempfile", + "test.support.get_attribute", + "types.FunctionType", + "unittest.mock._dot_lookup", + "unittest.mock._importer", + "xml.etree", + "xmlrpc.server.resolve_dotted_attribute", + ] +) -UNSAFE_STRINGS: FrozenSet[str] = frozenset([ - "CreateThread", - "Crypto", - "RtlMoveMemory", - "VirtualAlloc", - "WaitForSingleObject", - "_codecs.decode", - "_compat_pickle", - "_pickle", - "aiohttp", - "apply", - "asyncio", - "base64", - "bdb", - "breakpoint", - "cProfile", - "cloudpickle.load", - "cloudpickle.loads", - "code.interact", - "code.InteractiveConsole", - "code.InteractiveInterpreter", - "codecs.decode", - "codeop.compile_command", - "commands", - "compile", - "config.set_config_path", - "copyreg", - "corrupy", - "cryptography", - "ctorch", - "ctypes", - "decode", - "dill", - "eval", - "exec", - "execfile", - "fileinput", - "get_type_hints", - "gzip", - "hashlib", - "httplib", - "importlib", - "itemgetter", - "joblib", - "load", - "load_module", - "loads", - "lzma.open", - "malicious", - "marshal", - "msvcrt", - "open", - "os", - "pexpect", - "pickle", - "picklemagic", - "posix", - "profile", - "psutil", - "pty", - "numpy.lib.npyio.loadtxt", - "pycrypto", - "pydoc.locate", - "pydoc.pipepager", - "pydoc.replace", - "pandas.read_pickle", - "python", - "pywin32_system32", - "pyyaml", - "raise", - "read", - "requests", - "runpy", - "safer_pickle", - "saferpickle", - "shutil", - "socket", - "ssl", - "stdin", - "subprocess", - "sys", - "system", - "timeit", - "torch.load", - "torch.unsupported_tensor_ops", - "trace", - "txwinrm", - "urllib", - "webbrowser", - "winapi", - "winreg", - "write", - "zlib", -]) +UNSAFE_STRINGS: FrozenSet[str] = frozenset( + [ + "CreateThread", + "Crypto", + "RtlMoveMemory", + "VirtualAlloc", + "WaitForSingleObject", + "_codecs.decode", + "_compat_pickle", + "_pickle", + "aiohttp", + "apply", + "asyncio", + "base64", + "bdb", + "breakpoint", + "cProfile", + "cloudpickle.load", + "cloudpickle.loads", + "code.interact", + "code.InteractiveConsole", + "code.InteractiveInterpreter", + "codecs.decode", + "codeop.compile_command", + "commands", + "compile", + "config.set_config_path", + "copyreg", + "corrupy", + "cryptography", + "ctorch", + "ctypes", + "decode", + "dill", + "eval", + "exec", + "execfile", + "fileinput", + "get_type_hints", + "gzip", + "hashlib", + "httplib", + "importlib", + "itemgetter", + "joblib", + "load", + "load_module", + "loads", + "lzma.open", + "malicious", + "marshal", + "msvcrt", + "open", + "os", + "pexpect", + "pickle", + "picklemagic", + "posix", + "profile", + "psutil", + "pty", + "numpy.lib.npyio.loadtxt", + "pycrypto", + "pydoc.locate", + "pydoc.pipepager", + "pydoc.replace", + "pandas.read_pickle", + "python", + "pywin32_system32", + "pyyaml", + "raise", + "read", + "requests", + "runpy", + "safer_pickle", + "saferpickle", + "shutil", + "socket", + "ssl", + "stdin", + "subprocess", + "sys", + "system", + "timeit", + "torch.load", + "torch.unsupported_tensor_ops", + "trace", + "txwinrm", + "urllib", + "webbrowser", + "winapi", + "winreg", + "write", + "zlib", + ] +) -SAFE_STRINGS: FrozenSet[str] = frozenset([ - "PIL", - "__builtin__.print", - "__builtin__.set", - "cloudpickle.cloudpickle", - "collections", - "complex", - "ctorch.nn", - "cv2", - "dtypes", - "epoch_loop", - "functools.partial", - "gensim", - "google3", - "jax", - "keras", - "layer", - "lightning", - "nltk", - "nn", - "numpy", - "opacus", - "pandas", - "pillow", - "pydoc", - "python.v2.dataclasses", - "reconstruct", - "scipy", - "set", - "shutil.disk_usage", - "sklearn", - "spacy", - "str", - "tensorflow", - "theano", - "torch", - "torch_frame", - "torchmetrics", - "torchvision", - "training_step", - "transformers", - "usage", -]) +SAFE_STRINGS: FrozenSet[str] = frozenset( + [ + "PIL", + "__builtin__.print", + "__builtin__.set", + "cloudpickle.cloudpickle", + "collections", + "complex", + "ctorch.nn", + "cv2", + "dtypes", + "epoch_loop", + "functools.partial", + "gensim", + "google3", + "jax", + "keras", + "layer", + "lightning", + "nltk", + "nn", + "numpy", + "opacus", + "pandas", + "pillow", + "pydoc", + "python.v2.dataclasses", + "reconstruct", + "scipy", + "set", + "shutil.disk_usage", + "sklearn", + "spacy", + "str", + "tensorflow", + "theano", + "torch", + "torch_frame", + "torchmetrics", + "torchvision", + "training_step", + "transformers", + "usage", + ] +) # Error template for the hook ERROR_STRING = string.Template( @@ -289,12 +296,14 @@ # These are substrings of opcodes that declare strings often before # REDUCE, BUILD, and MEMOIZE opcodes -OPCODE_SUBSTRS_THAT_DECLARE_STRINGS: FrozenSet[str] = frozenset([ - "GLOBAL", - "INST", - "UNICODE", - "STRING", -]) +OPCODE_SUBSTRS_THAT_DECLARE_STRINGS: FrozenSet[str] = frozenset( + [ + "GLOBAL", + "INST", + "UNICODE", + "STRING", + ] +) OPCODES_INFO = immutabledict.immutabledict( {opcode.code: opcode for opcode in pickletools.opcodes} @@ -326,10 +335,11 @@ # High score to indicate definite unsafety due to a detected zip slip. HIGH_SEVERITY_ZIPSLIP = 1337 +# High score to indicate that an archive could not be processed safely. +HIGH_SEVERITY_ARCHIVE_ERROR = 1337 + OPCODES_4BYTE_LEN = (b"B",) OPCODES_8BYTE_LEN = (b"\x8e", b"\x96") OPCODES_1BYTE_LEN = (b"C",) -LENGTH_PREFIXED_OPCODES = ( - OPCODES_4BYTE_LEN + OPCODES_8BYTE_LEN + OPCODES_1BYTE_LEN -) +LENGTH_PREFIXED_OPCODES = OPCODES_4BYTE_LEN + OPCODES_8BYTE_LEN + OPCODES_1BYTE_LEN diff --git a/lib/exceptions.py b/lib/exceptions.py index 59c2622..25ff414 100644 --- a/lib/exceptions.py +++ b/lib/exceptions.py @@ -2,40 +2,40 @@ class IllegalArgumentCombinationError(Exception): - """Custom exception for using allow_unsafe and strict_check together.""" + """Custom exception for using allow_unsafe and strict_check together.""" - def __init__(self, m: str) -> None: - self.message = m + def __init__(self, m: str) -> None: + self.message = m - def __str__(self) -> str: - return self.message + def __str__(self) -> str: + return self.message class StrictCheckError(Exception): - """Custom exception for strict check failures.""" + """Custom exception for strict check failures.""" - def __init__(self, m: str) -> None: - self.message = m + def __init__(self, m: str) -> None: + self.message = m - def __str__(self) -> str: - return self.message + def __str__(self) -> str: + return self.message class UnsafePickleDetectedError(Exception): - """Custom exception for unsafe pickle files.""" + """Custom exception for unsafe pickle files.""" - def __init__(self, m: str) -> None: - self.message = m + def __init__(self, m: str) -> None: + self.message = m - def __str__(self) -> str: - return self.message + def __str__(self) -> str: + return self.message class MaxRecursionDepthExceededError(Exception): - """Custom exception for exceeding maximum recursion depth in archives.""" + """Custom exception for exceeding maximum recursion depth in archives.""" - def __init__(self, m: str) -> None: - self.message = m + def __init__(self, m: str) -> None: + self.message = m - def __str__(self) -> str: - return self.message + def __str__(self) -> str: + return self.message diff --git a/lib/utils.py b/lib/utils.py index da97476..c48b428 100644 --- a/lib/utils.py +++ b/lib/utils.py @@ -32,58 +32,58 @@ import tarfile import threading import types -from typing import BinaryIO, Callable, Dict, FrozenSet, Generator, IO, Set, Tuple, cast import zipfile +from typing import IO, BinaryIO, Callable, Dict, FrozenSet, Generator, Set, Tuple, cast from absl import logging -from lib import config -from lib import constants + +from lib import config, constants @enum.unique class Classification(enum.Enum): - """Classification of a class name.""" + """Classification of a class name.""" - SAFE = "SAFE" - UNSAFE = "UNSAFE" - SUSPICIOUS = "SUSPICIOUS" - UNKNOWN = "UNKNOWN" + SAFE = "SAFE" + UNSAFE = "UNSAFE" + SUSPICIOUS = "SUSPICIOUS" + UNKNOWN = "UNKNOWN" def create_pattern(strings: FrozenSet[str]) -> re.Pattern[str]: - """Creates a pattern for matching method calls from a list of strings. + """Creates a pattern for matching method calls from a list of strings. - Args: - strings: The strings to match. + Args: + strings: The strings to match. - Returns: - A pattern for matching method calls. - """ - included = "|".join(set(map(re.escape, strings))) - return re.compile( - rf""" + Returns: + A pattern for matching method calls. + """ + included = "|".join(set(map(re.escape, strings))) + return re.compile( + rf""" (? re.Pattern[str]: - """Creates a pattern for matching method calls excluding substrings from the list of strings. + """Creates a pattern for matching method calls excluding substrings from the list of strings. - Args: - strings: The strings to exclude in matching. + Args: + strings: The strings to exclude in matching. - Returns: - A pattern string for matching unknown method calls. - """ + Returns: + A pattern string for matching unknown method calls. + """ - excluded = "|".join(set(map(re.escape, strings))) - return re.compile( - rf""" + excluded = "|".join(set(map(re.escape, strings))) + return re.compile( + rf""" \b(? re.Pattern[str]: ) \b """, - re.VERBOSE, - ) + re.VERBOSE, + ) safe_pattern = create_pattern(constants.SAFE_STRINGS) @@ -103,16 +103,16 @@ def create_pattern_for_unknowns(strings: FrozenSet[str]) -> re.Pattern[str]: # Precompiled regex patterns for categorize_strings EXTRACT_UNSAFE_MODULE_REGEX = re.compile(r"warning: (.*?) is unsafe") -ARGS_REGEX = re.compile( - r"\.*?unexpected arguments [({](.*)[)}]" +ARGS_REGEX = re.compile(r"\.*?unexpected arguments [({](.*)[)}]") + +PYTHON_METHOD_PATTERNS = frozenset( + { + re.compile(r"(\w+\.\w+\(\))"), # Method Calls (a.b()) + re.compile(r"(\w+)\("), # Function Calls (a()) + re.compile(r"[b]?['\"](\w+)['\"]"), # String Arguments (like 'system') + } ) -PYTHON_METHOD_PATTERNS = frozenset({ - re.compile(r"(\w+\.\w+\(\))"), # Method Calls (a.b()) - re.compile(r"(\w+)\("), # Function Calls (a()) - re.compile(r"[b]?['\"](\w+)['\"]"), # String Arguments (like 'system') -}) - PICKLEMAGIC_PATTERNS = [ ( @@ -177,43 +177,43 @@ def create_pattern_for_unknowns(strings: FrozenSet[str]) -> re.Pattern[str]: # 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. - - Args: - original_name: The name of the module to copy. - new_name: The name of the new module. - - Returns: - new_module: The new module, or None if original_name cannot be imported. - """ - try: - original_module = importlib.import_module(original_name) - except ImportError: - logging.debug("Failed to import module %s", original_name) - return None - except IOError: - if original_name in sys.modules: - logging.debug( - "FileError during import of %s, but module is in sys.modules", - original_name, - ) - original_module = sys.modules[original_name] - else: - logging.debug( - "Failed to import module %s due to FileError and module not in" - " sys.modules", - original_name, - ) - return None + """Copies a module and creates a new module with the same attributes. - new_module = types.ModuleType(new_name) - new_module.__dict__.update(original_module.__dict__) - new_module.__name__ = new_name - if hasattr(original_module, "__file__"): - new_module.__file__ = f"{new_name}.py" - sys.modules[new_name] = new_module + Args: + original_name: The name of the module to copy. + new_name: The name of the new module. - return new_module + Returns: + new_module: The new module, or None if original_name cannot be imported. + """ + try: + original_module = importlib.import_module(original_name) + except ImportError: + logging.debug("Failed to import module %s", original_name) + return None + except IOError: + if original_name in sys.modules: + logging.debug( + "FileError during import of %s, but module is in sys.modules", + original_name, + ) + original_module = sys.modules[original_name] + else: + logging.debug( + "Failed to import module %s due to FileError and module not in" + " sys.modules", + original_name, + ) + return None + + new_module = types.ModuleType(new_name) + new_module.__dict__.update(original_module.__dict__) + new_module.__name__ = new_name + if hasattr(original_module, "__file__"): + new_module.__file__ = f"{new_name}.py" + sys.modules[new_name] = new_module + + return new_module _COPIED_MODS_CACHE: Dict[str, types.ModuleType] = {} @@ -221,161 +221,151 @@ def copy_module(original_name: str, new_name: str) -> types.ModuleType | None: def get_copied_module(name: str) -> types.ModuleType: - """Get or create a copy of the module, caching it to avoid re-copying.""" - with _COPIED_MODS_LOCK: - if name in _COPIED_MODS_CACHE: - return _COPIED_MODS_CACHE[name] - - if name in ("pickle", "_pickle"): - # We always copy _pickle for both pickle and _pickle - if "_pickle" in _COPIED_MODS_CACHE: - mod_copied = _COPIED_MODS_CACHE["_pickle"] - else: - mod_copied = copy_module("_pickle", "pickle_copy") - if mod_copied is None: - logging.error("Failed to copy critical module _pickle") - sys.exit(1) - _COPIED_MODS_CACHE["_pickle"] = mod_copied - _COPIED_MODS_CACHE["pickle"] = mod_copied - else: - # Try to copy the module - mod_copied = copy_module(name, f"{name}_copy") - if mod_copied is None: - # Fallback to pickle copy - logging.warning( - "%s could not be imported/copied, falling back to pickle_copy", name - ) - if "_pickle" not in _COPIED_MODS_CACHE: - _ = get_copied_module("_pickle") - mod_copied = _COPIED_MODS_CACHE["_pickle"] - - _COPIED_MODS_CACHE[name] = mod_copied - - return mod_copied + """Get or create a copy of the module, caching it to avoid re-copying.""" + with _COPIED_MODS_LOCK: + if name in _COPIED_MODS_CACHE: + return _COPIED_MODS_CACHE[name] + + if name in ("pickle", "_pickle"): + # We always copy _pickle for both pickle and _pickle + if "_pickle" in _COPIED_MODS_CACHE: + mod_copied = _COPIED_MODS_CACHE["_pickle"] + else: + mod_copied = copy_module("_pickle", "pickle_copy") + if mod_copied is None: + logging.error("Failed to copy critical module _pickle") + sys.exit(1) + _COPIED_MODS_CACHE["_pickle"] = mod_copied + _COPIED_MODS_CACHE["pickle"] = mod_copied + else: + # Try to copy the module + mod_copied = copy_module(name, f"{name}_copy") + if mod_copied is None: + # Fallback to pickle copy + logging.warning( + "%s could not be imported/copied, falling back to pickle_copy", name + ) + if "_pickle" not in _COPIED_MODS_CACHE: + _ = get_copied_module("_pickle") + mod_copied = _COPIED_MODS_CACHE["_pickle"] + + _COPIED_MODS_CACHE[name] = mod_copied + + return mod_copied def _peek_bytes(file_bytes: bytes | BinaryIO, size: int) -> bytes: - """Peeks at the first `size` bytes of a bytes object or file stream.""" - if isinstance(file_bytes, bytes): - return file_bytes[:size] - - try: - is_seekable = file_bytes.seekable() - except (AttributeError, ValueError): - is_seekable = False + """Peeks at the first `size` bytes of a bytes object or file stream.""" + if isinstance(file_bytes, bytes): + return file_bytes[:size] - if is_seekable: try: - current_pos = file_bytes.tell() - peeked = file_bytes.read(size) - file_bytes.seek(current_pos) - return peeked - except (OSError, io.UnsupportedOperation): - pass - - if hasattr(file_bytes, "peek"): - try: - return file_bytes.peek(size)[:size] - except (OSError, io.UnsupportedOperation, AttributeError): - pass + is_seekable = file_bytes.seekable() + except (AttributeError, ValueError): + is_seekable = False + + if is_seekable: + try: + current_pos = file_bytes.tell() + peeked = file_bytes.read(size) + file_bytes.seek(current_pos) + return peeked + except (OSError, io.UnsupportedOperation): + pass + + if hasattr(file_bytes, "peek"): + try: + return file_bytes.peek(size)[:size] + except (OSError, io.UnsupportedOperation, AttributeError): + pass - return b"" + return b"" def is_zip_bytes(file_bytes: bytes | BinaryIO) -> bool: - """Checks if the provided bytes/stream represent a zip file. + """Checks if the provided bytes/stream represent a zip file. - Args: - file_bytes: The bytes or stream to check. + Args: + file_bytes: The bytes or stream to check. - Returns: - True if the input is a zip file, False otherwise. - """ - if not file_bytes: - return False - return _peek_bytes(file_bytes, 4).startswith( - (b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08") - ) + Returns: + True if the input is a zip file, False otherwise. + """ + if not file_bytes: + return False + return _peek_bytes(file_bytes, 4).startswith( + (b"PK\x03\x04", b"PK\x05\x06", b"PK\x07\x08") + ) def extract_zip_contents( file_bytes: bytes | BinaryIO, ) -> Generator[Tuple[str, IO[bytes]], None, None]: - """Extracts the list of files and their contents from a zip file/stream. + """Extracts the list of files and their contents from a zip file/stream. - Args: - file_bytes: The bytes or stream to check. + Args: + file_bytes: The bytes or stream to check. - Yields: - A tuple containing the file name and its stream. - """ - stream = ( - io.BytesIO(file_bytes) if isinstance(file_bytes, bytes) else file_bytes - ) - with zipfile.ZipFile(stream) as zf: - for name in zf.namelist(): - with zf.open(name) as f: - yield name, f + Yields: + A tuple containing the file name and its stream. + """ + stream = io.BytesIO(file_bytes) if isinstance(file_bytes, bytes) else file_bytes + with zipfile.ZipFile(stream) as zf: + for name in zf.namelist(): + with zf.open(name) as f: + yield name, f def is_bz2_bytes(file_bytes: bytes | BinaryIO) -> bool: - """Checks if the provided bytes represent a bz2 file.""" - return _peek_bytes(file_bytes, 3).startswith(b"\x42\x5a\x68") + """Checks if the provided bytes represent a bz2 file.""" + return _peek_bytes(file_bytes, 3).startswith(b"\x42\x5a\x68") def extract_bz2_contents(file_bytes: bytes | BinaryIO) -> IO[bytes]: - """Extracts contents from bz2 bytes/stream as a stream.""" - stream = ( - io.BytesIO(file_bytes) if isinstance(file_bytes, bytes) else file_bytes - ) - return bz2.open(stream) + """Extracts contents from bz2 bytes/stream as a stream.""" + stream = io.BytesIO(file_bytes) if isinstance(file_bytes, bytes) else file_bytes + return bz2.open(stream) def is_lzma_bytes(file_bytes: bytes | BinaryIO) -> bool: - """Checks if the provided bytes represent an lzma file.""" - return _peek_bytes(file_bytes, 6).startswith(b"\xfd\x37\x7a\x58\x5a\x00") + """Checks if the provided bytes represent an lzma file.""" + return _peek_bytes(file_bytes, 6).startswith(b"\xfd\x37\x7a\x58\x5a\x00") def extract_lzma_contents(file_bytes: bytes | BinaryIO) -> IO[bytes]: - """Extracts contents from lzma bytes/stream as a stream.""" - stream = ( - io.BytesIO(file_bytes) if isinstance(file_bytes, bytes) else file_bytes - ) - return lzma.open(stream) + """Extracts contents from lzma bytes/stream as a stream.""" + stream = io.BytesIO(file_bytes) if isinstance(file_bytes, bytes) else file_bytes + return lzma.open(stream) def is_gzip_bytes(file_bytes: bytes | BinaryIO) -> bool: - """Checks if the provided bytes represent a gzip file.""" - return _peek_bytes(file_bytes, 2).startswith(b"\x1f\x8b") + """Checks if the provided bytes represent a gzip file.""" + return _peek_bytes(file_bytes, 2).startswith(b"\x1f\x8b") def extract_gzip_contents(file_bytes: bytes | IO[bytes]) -> IO[bytes]: - """Extracts contents from gzip bytes/stream as a stream.""" - stream = ( - io.BytesIO(file_bytes) if isinstance(file_bytes, bytes) else file_bytes - ) - return cast(IO[bytes], gzip.open(stream)) + """Extracts contents from gzip bytes/stream as a stream.""" + stream = io.BytesIO(file_bytes) if isinstance(file_bytes, bytes) else file_bytes + return cast(IO[bytes], gzip.open(stream)) def is_tar_bytes(file_bytes: bytes | BinaryIO) -> bool: - """Checks if the provided bytes represent a tar file.""" - peeked = _peek_bytes(file_bytes, 262) - return len(peeked) >= 262 and peeked[257:262] == b"ustar" + """Checks if the provided bytes represent a tar file.""" + peeked = _peek_bytes(file_bytes, 262) + return len(peeked) >= 262 and peeked[257:262] == b"ustar" def extract_tar_contents( file_bytes: bytes | BinaryIO, ) -> Generator[Tuple[str, IO[bytes]], None, None]: - """Extracts contents from tar bytes/stream.""" - stream = ( - io.BytesIO(file_bytes) if isinstance(file_bytes, bytes) else file_bytes - ) - with tarfile.open(fileobj=stream, mode="r:*") as tf: - for member in tf.getmembers(): - if member.isfile(): - f = tf.extractfile(member) - if f is not None: - yield member.name, cast(IO[bytes], f) + """Extracts contents from tar bytes/stream.""" + stream = io.BytesIO(file_bytes) if isinstance(file_bytes, bytes) else file_bytes + with tarfile.open(fileobj=stream, mode="r:*") as tf: + for member in tf.getmembers(): + if member.isfile(): + f = tf.extractfile(member) + if f is not None: + yield member.name, cast(IO[bytes], f) def is_pickle_file( @@ -383,487 +373,482 @@ def is_pickle_file( return_num_bytes_read: bool = False, check_magic_bytes: bool = True, ) -> bool | tuple[bool, int]: - """Checks if the provided bytes represent a valid pickle file. - - This function reads the beginning of the input byte stream, looking for - valid pickle opcodes. It stops after reading a maximum number of bytes, - defined by `_MAX_BYTES_TO_CHECK`, to avoid processing very large inputs. - Do note that this is not foolproof and false positives are possible. - - Args: - pickle_bytes: The bytes or stream to check. - return_num_bytes_read: If True, returns a tuple containing a boolean - indicating if the file is a valid pickle file and the number of bytes - read. Otherwise, it returns only the boolean. - check_magic_bytes: If True, checks for text-based or archive magic bytes - to fast-path reject files that are not pickles. - - Returns: - If `return_num_bytes_read` is True: - - A tuple `(True, number_of_bytes_read)` if the input is likely a valid - pickle file, and the number of bytes read. - - A tuple `(False, number_of_bytes_read)` if the input is not a valid - pickle file, and the number of bytes read. - If `return_num_bytes_read` is False: - - True if the input is likely a valid pickle file. - - False if the input is not a valid pickle file. - """ - original_pos = None - is_seekable = True - if isinstance(pickle_bytes, bytes): - raw_bytes = pickle_bytes - pickle_stream = io.BytesIO(pickle_bytes) - else: - try: - original_pos = pickle_bytes.tell() - except (OSError, io.UnsupportedOperation, AttributeError): - is_seekable = False - - if is_seekable: - # Read a chunk to check text-based and non-pickle-magic prefixes safely - raw_bytes = pickle_bytes.read(1024) - pickle_bytes.seek(original_pos) # pyrefly: ignore[bad-argument-type] - pickle_stream = pickle_bytes + """Checks if the provided bytes represent a valid pickle file. + + This function reads the beginning of the input byte stream, looking for + valid pickle opcodes. It stops after reading a maximum number of bytes, + defined by `_MAX_BYTES_TO_CHECK`, to avoid processing very large inputs. + Do note that this is not foolproof and false positives are possible. + + Args: + pickle_bytes: The bytes or stream to check. + return_num_bytes_read: If True, returns a tuple containing a boolean + indicating if the file is a valid pickle file and the number of bytes + read. Otherwise, it returns only the boolean. + check_magic_bytes: If True, checks for text-based or archive magic bytes + to fast-path reject files that are not pickles. + + Returns: + If `return_num_bytes_read` is True: + - A tuple `(True, number_of_bytes_read)` if the input is likely a valid + pickle file, and the number of bytes read. + - A tuple `(False, number_of_bytes_read)` if the input is not a valid + pickle file, and the number of bytes read. + If `return_num_bytes_read` is False: + - True if the input is likely a valid pickle file. + - False if the input is not a valid pickle file. + """ + original_pos = None + is_seekable = True + if isinstance(pickle_bytes, bytes): + raw_bytes = pickle_bytes + pickle_stream = io.BytesIO(pickle_bytes) else: - # Fallback for non-seekable stream: try to peek without advancing pointer - if hasattr(pickle_bytes, "peek"): try: - raw_bytes = pickle_bytes.peek(1024) - pickle_stream = pickle_bytes - except (OSError, io.UnsupportedOperation): - raw_bytes = pickle_bytes.read(1024) - pickle_stream = io.BytesIO(raw_bytes) - else: - raw_bytes = pickle_bytes.read(1024) - pickle_stream = io.BytesIO(raw_bytes) - - if raw_bytes: - first_byte = raw_bytes[0] - if first_byte not in constants.OPCODES_INFO_INT: - if return_num_bytes_read: - return (False, 0) - return False - - if check_magic_bytes: - pickle_file_is_ascii = raw_bytes.isascii() - - if pickle_file_is_ascii: - stripped_bytes = raw_bytes.lstrip() - if stripped_bytes.startswith( - constants.TEXT_BASED_PREFIXES - ) or stripped_bytes.startswith(constants.CODE_KEYWORDS): - if return_num_bytes_read: - return (False, 0) - return False - - if raw_bytes.startswith(constants.NON_PICKLE_MAGIC_BYTES): - if raw_bytes[0] not in constants.OPCODES_INFO_INT: - if return_num_bytes_read: - return (False, 0) - return False - - valid_opcodes_count = 0 - try: - while True: - charcode = pickle_stream.read(1) - if not charcode: # EOF reached without STOP - is_suspected_pickle = valid_opcodes_count >= 3 - if return_num_bytes_read: - return (is_suspected_pickle, valid_opcodes_count) - return is_suspected_pickle - - decoded_char = charcode.decode("latin-1") - if decoded_char == ".": # STOP opcode found - valid_opcodes_count += 1 - if return_num_bytes_read: - return (True, valid_opcodes_count) - return True - - opcode = constants.OPCODES_INFO.get(decoded_char) - if opcode is None: # Invalid opcode before STOP - is_suspected_pickle = valid_opcodes_count >= 3 - if return_num_bytes_read: - return (is_suspected_pickle, valid_opcodes_count) - return is_suspected_pickle - - valid_opcodes_count += 1 - if valid_opcodes_count > constants.MAX_BYTES_TO_CHECK: - if return_num_bytes_read: - return (True, valid_opcodes_count) - return True - - if opcode.arg is None: - continue - try: - _ = opcode.arg.reader(pickle_stream) - except ValueError: - is_suspected_pickle = valid_opcodes_count >= 3 - if return_num_bytes_read: - return (is_suspected_pickle, valid_opcodes_count) - return is_suspected_pickle - finally: - if original_pos is not None: - pickle_stream.seek(original_pos) + original_pos = pickle_bytes.tell() + except (OSError, io.UnsupportedOperation, AttributeError): + is_seekable = False + + if is_seekable: + # Read a chunk to check text-based and non-pickle-magic prefixes safely + raw_bytes = pickle_bytes.read(1024) + pickle_bytes.seek(original_pos) # pyrefly: ignore[bad-argument-type] + pickle_stream = pickle_bytes + else: + # Fallback for non-seekable stream: try to peek without advancing pointer + if hasattr(pickle_bytes, "peek"): + try: + raw_bytes = pickle_bytes.peek(1024) + pickle_stream = pickle_bytes + except (OSError, io.UnsupportedOperation): + raw_bytes = pickle_bytes.read(1024) + pickle_stream = io.BytesIO(raw_bytes) + else: + raw_bytes = pickle_bytes.read(1024) + pickle_stream = io.BytesIO(raw_bytes) + + if raw_bytes: + first_byte = raw_bytes[0] + if first_byte not in constants.OPCODES_INFO_INT: + if return_num_bytes_read: + return (False, 0) + return False + + if check_magic_bytes: + pickle_file_is_ascii = raw_bytes.isascii() + + if pickle_file_is_ascii: + stripped_bytes = raw_bytes.lstrip() + if stripped_bytes.startswith( + constants.TEXT_BASED_PREFIXES + ) or stripped_bytes.startswith(constants.CODE_KEYWORDS): + if return_num_bytes_read: + return (False, 0) + return False + + if raw_bytes.startswith(constants.NON_PICKLE_MAGIC_BYTES): + if raw_bytes[0] not in constants.OPCODES_INFO_INT: + if return_num_bytes_read: + return (False, 0) + return False + + valid_opcodes_count = 0 + try: + while True: + charcode = pickle_stream.read(1) + if not charcode: # EOF reached without STOP + is_suspected_pickle = valid_opcodes_count >= 3 + if return_num_bytes_read: + return (is_suspected_pickle, valid_opcodes_count) + return is_suspected_pickle + + decoded_char = charcode.decode("latin-1") + if decoded_char == ".": # STOP opcode found + valid_opcodes_count += 1 + if return_num_bytes_read: + return (True, valid_opcodes_count) + return True + + opcode = constants.OPCODES_INFO.get(decoded_char) + if opcode is None: # Invalid opcode before STOP + is_suspected_pickle = valid_opcodes_count >= 3 + if return_num_bytes_read: + return (is_suspected_pickle, valid_opcodes_count) + return is_suspected_pickle + + valid_opcodes_count += 1 + if valid_opcodes_count > constants.MAX_BYTES_TO_CHECK: + if return_num_bytes_read: + return (True, valid_opcodes_count) + return True + + if opcode.arg is None: + continue + try: + _ = opcode.arg.reader(pickle_stream) + except ValueError: + is_suspected_pickle = valid_opcodes_count >= 3 + if return_num_bytes_read: + return (is_suspected_pickle, valid_opcodes_count) + return is_suspected_pickle + finally: + if original_pos is not None: + pickle_stream.seek(original_pos) def find_pickle_start_offset(pickle_bytes: bytes | IO[bytes]) -> int: - """Finds the start offset of a valid pickle payload in the bytes.""" - if isinstance(pickle_bytes, bytes): - max_search_len = min(len(pickle_bytes), 1024) - for offset in range(max_search_len): - char = pickle_bytes[offset : offset + 1] - if not char: - break - try: - decoded_char = char.decode("latin-1") - except UnicodeDecodeError: - continue - if decoded_char not in constants.OPCODES_INFO: - continue - if is_pickle_file(pickle_bytes[offset:]): - return offset - return 0 - - # It is a stream - stream = pickle_bytes - try: - original_pos = stream.tell() - # Read 1024 bytes to find candidates - header_bytes = stream.read(1024) - stream.seek(original_pos) - except (OSError, AttributeError, io.UnsupportedOperation): - return 0 - - max_search_len = len(header_bytes) - for offset in range(max_search_len): - char = header_bytes[offset : offset + 1] - if not char: - break - try: - decoded_char = char.decode("latin-1") - except UnicodeDecodeError: - continue - if decoded_char not in constants.OPCODES_INFO: - continue - - # Verify candidate offset using the stream + """Finds the start offset of a valid pickle payload in the bytes.""" + if isinstance(pickle_bytes, bytes): + max_search_len = min(len(pickle_bytes), 1024) + for offset in range(max_search_len): + char = pickle_bytes[offset : offset + 1] + if not char: + break + try: + decoded_char = char.decode("latin-1") + except UnicodeDecodeError: + continue + if decoded_char not in constants.OPCODES_INFO: + continue + if is_pickle_file(pickle_bytes[offset:]): + return offset + return 0 + + # It is a stream + stream = pickle_bytes try: - stream.seek(original_pos + offset) - if is_pickle_file(stream): - return offset - except (OSError, AttributeError, io.UnsupportedOperation): - pass - finally: - try: + original_pos = stream.tell() + # Read 1024 bytes to find candidates + header_bytes = stream.read(1024) stream.seek(original_pos) - except (OSError, AttributeError, io.UnsupportedOperation): - pass # If we can't seek back, it poses an issue + except (OSError, AttributeError, io.UnsupportedOperation): + return 0 - return 0 + max_search_len = len(header_bytes) + for offset in range(max_search_len): + char = header_bytes[offset : offset + 1] + if not char: + break + try: + decoded_char = char.decode("latin-1") + except UnicodeDecodeError: + continue + if decoded_char not in constants.OPCODES_INFO: + continue + + # Verify candidate offset using the stream + try: + stream.seek(original_pos + offset) + if is_pickle_file(stream): + return offset + except (OSError, AttributeError, io.UnsupportedOperation): + pass + finally: + try: + stream.seek(original_pos) + except (OSError, AttributeError, io.UnsupportedOperation): + pass # If we can't seek back, it poses an issue + + return 0 @functools.lru_cache(maxsize=None) def get_module_members(module_name: str) -> Set[str] | None: - """Tries to get module members by parsing the source file without execution of __init__.py. - - Args: - module_name: The name of the module to get members from. - - Returns: - A set of module members, or None for the following cases: - 1. if the module could not be parsed. - 2. If an ImportError is raised. - 3. If AST parsing fails, and SyntaxError or ValueError is raised. - 4. If spec is not present and its origin does not exist. - """ - - # If module is already imported, __init__.py will not be called while - # importing the module again - if module_name in sys.modules: - try: - imported_module = importlib.import_module(module_name) - except ImportError: - return None - return {member for member, _ in inspect.getmembers(imported_module)} - - if (spec := importlib.util.find_spec(module_name)) is None: - return None - if (origin := spec.origin) is None: - return None + """Tries to get module members by parsing the source file without execution of __init__.py. + + Args: + module_name: The name of the module to get members from. + + Returns: + A set of module members, or None for the following cases: + 1. if the module could not be parsed. + 2. If an ImportError is raised. + 3. If AST parsing fails, and SyntaxError or ValueError is raised. + 4. If spec is not present and its origin does not exist. + """ + + # If module is already imported, __init__.py will not be called while + # importing the module again + if module_name in sys.modules: + try: + imported_module = importlib.import_module(module_name) + except ImportError: + return None + return {member for member, _ in inspect.getmembers(imported_module)} - try: - with open(origin, "r") as f: - tree = ast.parse(f.read(), filename=origin) - except (IOError, ValueError, SyntaxError, RecursionError): - return None + if (spec := importlib.util.find_spec(module_name)) is None: + return None + if (origin := spec.origin) is None: + return None - members = set() - # This generates the list of methods and classes in the module from - # the AST tree of the module source file. - for node in ast.walk(tree): - if isinstance(node, ast.FunctionDef): - members.add(node.name) - elif isinstance(node, ast.ClassDef): - members.add(node.name) - return members + try: + with open(origin, "r") as f: + tree = ast.parse(f.read(), filename=origin) + except (IOError, ValueError, SyntaxError, RecursionError): + return None + + members = set() + # This generates the list of methods and classes in the module from + # the AST tree of the module source file. + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef): + members.add(node.name) + elif isinstance(node, ast.ClassDef): + members.add(node.name) + return members def get_optimal_workers(file_size: int) -> int: - """Calculates the optimal number of workers based on the file size using tiers. + """Calculates the optimal number of workers based on the file size using tiers. - Args: - file_size: The size of the file in bytes. + Args: + file_size: The size of the file in bytes. - Returns: - The optimal number of workers to use. - """ - if file_size is None: - return 1 - for threshold, workers in constants.WORKER_TIERS: - if file_size < threshold: - return min(constants.MAX_NUM_CHUNKS or 1, workers) + Returns: + The optimal number of workers to use. + """ + if file_size is None: + return 1 + for threshold, workers in constants.WORKER_TIERS: + if file_size < threshold: + return min(constants.MAX_NUM_CHUNKS or 1, workers) - # If file size is larger than or equal to the largest threshold, - # use logarithmic scaling. - largest_threshold, largest_workers = constants.WORKER_TIERS[-1] - # Ensure largest_workers is capped at MAX_NUM_CHUNKS before scaling up. - largest_workers = min(largest_workers, constants.MAX_NUM_CHUNKS or 1) - scaled_workers = largest_workers + int( - math.log(file_size / largest_threshold, 2) - ) + # If file size is larger than or equal to the largest threshold, + # use logarithmic scaling. + largest_threshold, largest_workers = constants.WORKER_TIERS[-1] + # Ensure largest_workers is capped at MAX_NUM_CHUNKS before scaling up. + largest_workers = min(largest_workers, constants.MAX_NUM_CHUNKS or 1) + scaled_workers = largest_workers + int(math.log(file_size / largest_threshold, 2)) - # Cap at around half the number of available CPU cores. - return min(constants.MAX_NUM_CHUNKS or 1, scaled_workers) + # Cap at around half the number of available CPU cores. + return min(constants.MAX_NUM_CHUNKS or 1, scaled_workers) @functools.lru_cache(maxsize=None) def classify_class_name(class_name: str) -> Classification | None: - """Classifies a class name based on the safe, unsafe, and suspicious patterns.""" - if re.search(safe_pattern, class_name): - return Classification.SAFE - if re.search(unsafe_pattern, class_name): - return Classification.UNSAFE - if re.search(suspicious_pattern, class_name): - return Classification.SUSPICIOUS - if re.search(unknown_pattern, class_name): - return Classification.UNKNOWN - return None + """Classifies a class name based on the safe, unsafe, and suspicious patterns.""" + if re.search(safe_pattern, class_name): + return Classification.SAFE + if re.search(unsafe_pattern, class_name): + return Classification.UNSAFE + if re.search(suspicious_pattern, class_name): + return Classification.SUSPICIOUS + if re.search(unknown_pattern, class_name): + return Classification.UNKNOWN + return None def is_unsafe_or_suspicious(class_name: str) -> bool: - allow_list = config.get_allow_list() - if any( - allowed_item.startswith(class_name) - or class_name.startswith(allowed_item) - or allowed_item.endswith(f".{class_name}") - for allowed_item in allow_list - ): - return False - classification = classify_class_name(class_name) - return classification == Classification.UNSAFE + allow_list = config.get_allow_list() + if any( + allowed_item.startswith(class_name) + or class_name.startswith(allowed_item) + or allowed_item.endswith(f".{class_name}") + for allowed_item in allow_list + ): + return False + classification = classify_class_name(class_name) + return classification == Classification.UNSAFE def resolve_library_modules_from_results( set_of_results: Set[str], ) -> Set[str]: - """Processes a set of strings to combine Python libraries and their members. - - For example, given {"os", "system", "pickle.loads"}, this function will - return {"os.system", "pickle.loads"}. - - Args: - set_of_results: A set of strings, some of which may be Python library names - and others may be their members. - - Returns: - final_results: A set of strings with libraries and their corresponding - members joined by a dot, along with any other strings from the original set. - """ - # Items with a dot are assumed to be fully qualified already. - # In case we run into cases such as os.path and join, this is a non-issue - # since this is moreso for better readability than precise library-to-member - # connections (ideal final result being os.path.join). - # Without explicit runtime introspection, the above case is not - # possible to resolve without risks of accidentally importing a - # module we don't want to. - final_results = {s for s in set_of_results if "." in s} - candidates = set_of_results - final_results - - # Identify which of the remaining candidates are actual importable modules. - importable_modules = {} - for name in candidates: - module_name = name.split(".")[0] - if module_name == "__main__": - continue - - # Find spec to avoid importing risky modules from a loose python file - # like import.py or similar. - module_spec = importlib.util.find_spec(name) - if module_spec is not None: - module_members = get_module_members(name) - if module_members is not None: - importable_modules[name] = module_members - - resolved_candidates = set() - # Combine modules with any members found in the candidates list. - for module_name, members in importable_modules.items(): - found_member_in_candidates = False - for member_name in candidates: - if member_name in members: - final_results.add(f"{module_name}.{member_name}") - resolved_candidates.add(module_name) - resolved_candidates.add(member_name) - found_member_in_candidates = True - - # If an importable module was not combined with any member, and it hasn't - # been used as a member itself, add it as a standalone item. - if ( - not found_member_in_candidates - and module_name not in resolved_candidates - ): - final_results.add(module_name) - resolved_candidates.add(module_name) - - # Add any remaining items that were not resolved as modules or members. - final_results.update(candidates - resolved_candidates) - - # Filter out base libraries if a qualified member from that module is present. - # Eg. Remove os if os.system is present. - libraries_to_remove = set() - modules_to_remove = set() - for result in final_results: - if "." in result: - base_library = result.split(".", 1)[0] - base_module = result.split(".", 1)[1] - if base_library in final_results: - libraries_to_remove.add(base_library) - if base_module in final_results: - modules_to_remove.add(base_module) - - final_results.difference_update(libraries_to_remove) - final_results.difference_update(modules_to_remove) - - # Remove less specific versions of qualified modules. - # Eg. Remove requests.api if requests.api.post is present. - underspecific_qualified_modules_to_remove = set() - for res in final_results: - attribute_parts = res.split(".") - # Check for less specific versions of the current item - for i in range(1, len(attribute_parts)): - parent = ".".join(attribute_parts[:i]) - if parent in final_results: - underspecific_qualified_modules_to_remove.add(parent) - - final_results.difference_update(underspecific_qualified_modules_to_remove) - - return final_results + """Processes a set of strings to combine Python libraries and their members. + + For example, given {"os", "system", "pickle.loads"}, this function will + return {"os.system", "pickle.loads"}. + + Args: + set_of_results: A set of strings, some of which may be Python library names + and others may be their members. + + Returns: + final_results: A set of strings with libraries and their corresponding + members joined by a dot, along with any other strings from the original set. + """ + # Items with a dot are assumed to be fully qualified already. + # In case we run into cases such as os.path and join, this is a non-issue + # since this is moreso for better readability than precise library-to-member + # connections (ideal final result being os.path.join). + # Without explicit runtime introspection, the above case is not + # possible to resolve without risks of accidentally importing a + # module we don't want to. + final_results = {s for s in set_of_results if "." in s} + candidates = set_of_results - final_results + + # Identify which of the remaining candidates are actual importable modules. + importable_modules = {} + for name in candidates: + module_name = name.split(".")[0] + if module_name == "__main__": + continue + + # Find spec to avoid importing risky modules from a loose python file + # like import.py or similar. + module_spec = importlib.util.find_spec(name) + if module_spec is not None: + module_members = get_module_members(name) + if module_members is not None: + importable_modules[name] = module_members + + resolved_candidates = set() + # Combine modules with any members found in the candidates list. + for module_name, members in importable_modules.items(): + found_member_in_candidates = False + for member_name in candidates: + if member_name in members: + final_results.add(f"{module_name}.{member_name}") + resolved_candidates.add(module_name) + resolved_candidates.add(member_name) + found_member_in_candidates = True + + # If an importable module was not combined with any member, and it hasn't + # been used as a member itself, add it as a standalone item. + if not found_member_in_candidates and module_name not in resolved_candidates: + final_results.add(module_name) + resolved_candidates.add(module_name) + + # Add any remaining items that were not resolved as modules or members. + final_results.update(candidates - resolved_candidates) + + # Filter out base libraries if a qualified member from that module is present. + # Eg. Remove os if os.system is present. + libraries_to_remove = set() + modules_to_remove = set() + for result in final_results: + if "." in result: + base_library = result.split(".", 1)[0] + base_module = result.split(".", 1)[1] + if base_library in final_results: + libraries_to_remove.add(base_library) + if base_module in final_results: + modules_to_remove.add(base_module) + + final_results.difference_update(libraries_to_remove) + final_results.difference_update(modules_to_remove) + + # Remove less specific versions of qualified modules. + # Eg. Remove requests.api if requests.api.post is present. + underspecific_qualified_modules_to_remove = set() + for res in final_results: + attribute_parts = res.split(".") + # Check for less specific versions of the current item + for i in range(1, len(attribute_parts)): + parent = ".".join(attribute_parts[:i]) + if parent in final_results: + underspecific_qualified_modules_to_remove.add(parent) + + final_results.difference_update(underspecific_qualified_modules_to_remove) + + return final_results def is_valid_python_interpreter(path: str) -> bool: - """Checks if a given path points to a valid Python interpreter. + """Checks if a given path points to a valid Python interpreter. - Args: - path: The path to the potential Python interpreter. + Args: + path: The path to the potential Python interpreter. - Returns: - True if the path is a valid and executable Python interpreter, False - otherwise. - """ - if not path or not os.path.exists(path): - return False + Returns: + True if the path is a valid and executable Python interpreter, False + otherwise. + """ + if not path or not os.path.exists(path): + return False - if "python" not in path: - return False + if "python" not in path: + return False - try: - # pass is a valid python keyword to test in python -c - subprocess.check_call( - [path, "-c", "pass"], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - return True - # If "python -c pass" doesn't work, we assume the interpreter is not valid. - except (OSError, subprocess.CalledProcessError): - return False + try: + # pass is a valid python keyword to test in python -c + subprocess.check_call( + [path, "-c", "pass"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return True + # If "python -c pass" doesn't work, we assume the interpreter is not valid. + except (OSError, subprocess.CalledProcessError): + return False def get_interpreter_path_to_patch() -> str | None: - """Returns a valid Python interpreter path if one can be found.""" - py_intrp_candidate = sys.argv[0] - python_path = "/usr/bin/python3" + """Returns a valid Python interpreter path if one can be found.""" + py_intrp_candidate = sys.argv[0] + python_path = "/usr/bin/python3" - if is_valid_python_interpreter(py_intrp_candidate): - return py_intrp_candidate + if is_valid_python_interpreter(py_intrp_candidate): + return py_intrp_candidate - logging.warning( - "Warning: %s (sys.argv[0]) is not a valid interpreter.", - py_intrp_candidate, - ) - if os.path.exists(python_path): - return python_path + logging.warning( + "Warning: %s (sys.argv[0]) is not a valid interpreter.", + py_intrp_candidate, + ) + if os.path.exists(python_path): + return python_path - return None + return None def is_sys_executable_to_path_set(path: str | None) -> bool: - """Checks if sys.executable is set to the given path.""" - if sys.executable == path or (sys.executable and "python" in sys.executable): - return True - elif path: - logging.info("Patching sys.executable to %s", path) - sys.executable = path - return True + """Checks if sys.executable is set to the given path.""" + if sys.executable == path or (sys.executable and "python" in sys.executable): + return True + elif path: + logging.info("Patching sys.executable to %s", path) + sys.executable = path + return True - logging.warning("sys.executable is not set to a valid interpreter.") - sys.executable = None # pyrefly: ignore[bad-assignment] - return False + logging.warning("sys.executable is not set to a valid interpreter.") + sys.executable = None # pyrefly: ignore[bad-assignment] + return False def is_sys_executable_patched() -> bool: - """Patches `sys.executable` if it's not set. - - This function attempts to set `sys.executable` to a valid Python interpreter - path. It first tries `sys.argv[0]`. If that's not a valid interpreter, it - falls back to "/usr/bin/python3". If neither is valid, `sys.executable` is - set to None. - - Returns: - True if `sys.executable` was successfully patched to a valid path, False - otherwise. - """ - py_interpreter_candidate = sys.argv[0] - - # If sys.executable is not set, we patch it with the interpreter path that - # was passed to the subprocess. If the path is not valid, we patch it with an - # empty string to indicate that it's invalid. - if not sys.executable: - valid_py_interpreter_path = get_interpreter_path_to_patch() - - for interpreter_path in [ - py_interpreter_candidate, - valid_py_interpreter_path, - ]: - if ( - interpreter_path and os.path.exists(interpreter_path) - ) and is_sys_executable_to_path_set(interpreter_path): - return True - - logging.warning("Warning: sys.executable is not set to a valid interpreter.") - return False + """Patches `sys.executable` if it's not set. + + This function attempts to set `sys.executable` to a valid Python interpreter + path. It first tries `sys.argv[0]`. If that's not a valid interpreter, it + falls back to "/usr/bin/python3". If neither is valid, `sys.executable` is + set to None. + + Returns: + True if `sys.executable` was successfully patched to a valid path, False + otherwise. + """ + py_interpreter_candidate = sys.argv[0] + + # If sys.executable is not set, we patch it with the interpreter path that + # was passed to the subprocess. If the path is not valid, we patch it with an + # empty string to indicate that it's invalid. + if not sys.executable: + valid_py_interpreter_path = get_interpreter_path_to_patch() + + for interpreter_path in [ + py_interpreter_candidate, + valid_py_interpreter_path, + ]: + if ( + interpreter_path and os.path.exists(interpreter_path) + ) and is_sys_executable_to_path_set(interpreter_path): + return True + + 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) + """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( @@ -872,121 +857,117 @@ def _parse_and_process_pattern( 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 + """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) + 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 + 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], - ) + """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/pyproject.toml b/pyproject.toml index 00ce422..793b75c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,55 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + [project] name = "saferpickle" version = "0.1.0" -description = "Add your description here" +description = "A safer alternative to Python's pickle module." readme = "README.md" requires-python = ">=3.13" +license = { text = "Apache-2.0" } +authors = [ + { name = "The SaferPickle Authors", email = "saferpickle-dev@google.com" }, +] dependencies = [ "absl-py>=2.5.0", "immutabledict>=4.3.1", + "pytest>=9.1.1", +] +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", + "Operating System :: OS Independent", +] + +[project.urls] +Homepage = "https://github.com/google/saferpickle" + +[project.scripts] +safer_pickle_cli = "cli:run" + +[dependency-groups] +dev = [ + "black>=26.5.1", + "isort>=8.0.1", + "pytest>=8.0.0", ] + +[tool.setuptools] +py-modules = ["saferpickle", "cli"] + +[tool.setuptools.packages.find] +include = ["lib*", "third_party*"] + +[tool.black] +target-version = ["py313"] +extend-exclude = "third_party" + +[tool.isort] +profile = "black" +skip = ["third_party"] + +[tool.pytest.ini_options] +pythonpath = ["."] +testpaths = ["tests"] diff --git a/saferpickle.py b/saferpickle.py index 974f181..97f5539 100644 --- a/saferpickle.py +++ b/saferpickle.py @@ -23,7 +23,7 @@ import logging as std_logging import lzma import math -from multiprocessing import shared_memory +import multiprocessing import os import pickle import pickletools @@ -33,17 +33,14 @@ import tarfile import tempfile import threading -from typing import Any, BinaryIO, Callable, Dict, IO, Iterator, Optional, Set, Tuple import zipfile +from multiprocessing import shared_memory +from typing import IO, Any, BinaryIO, Callable, Dict, Iterator, Optional, Set, Tuple from absl import logging -from third_party.corrupy import picklemagic -from lib import config -from lib import constants -from lib import exceptions -from lib import utils -import multiprocessing +from lib import config, constants, exceptions, utils +from third_party.corrupy import picklemagic IllegalArgumentCombinationError = exceptions.IllegalArgumentCombinationError StrictCheckError = exceptions.StrictCheckError @@ -63,200 +60,206 @@ @dataclasses.dataclass class ScanResults: - """Results from a pickle security scan.""" + """Results from a pickle security scan.""" - safe_results: Set[str] = dataclasses.field(default_factory=set) - unsafe_results: Set[str] = dataclasses.field(default_factory=set) - suspicious_results: Set[str] = dataclasses.field(default_factory=set) - unknown_results: Set[str] = dataclasses.field(default_factory=set) - is_denylisted: bool = False + safe_results: Set[str] = dataclasses.field(default_factory=set) + unsafe_results: Set[str] = dataclasses.field(default_factory=set) + suspicious_results: Set[str] = dataclasses.field(default_factory=set) + unknown_results: Set[str] = dataclasses.field(default_factory=set) + is_denylisted: bool = False def _custom_genops( pickle_bytes: bytes, ) -> Iterator[tuple[pickletools.OpcodeInfo, Any | None]]: - """Generates string-declaring opcodes and their arguments from pickle data. - - Args: - pickle_bytes: The pickle data to generate opcodes from. + """Generates string-declaring opcodes and their arguments from pickle data. - Yields: - A tuple of (opcode, opcode_argument) for each string-declaring opcode. - """ - - if isinstance(pickle_bytes, bytes): - pickle_file = io.BytesIO(pickle_bytes) - else: - pickle_file = pickle_bytes + Args: + pickle_bytes: The pickle data to generate opcodes from. - while True: - charcode = pickle_file.read(1) - if not charcode: # Indicates exhaustion of the data stream - break + Yields: + A tuple of (opcode, opcode_argument) for each string-declaring opcode. + """ - try: - opcode = constants.OPCODES_INFO_INT.get(charcode[0]) - except IndexError: - continue # Skip invalid opcode bytes + if isinstance(pickle_bytes, bytes): + pickle_file = io.BytesIO(pickle_bytes) + else: + pickle_file = pickle_bytes - if opcode is None: - # We skip processing unknown opcodes - continue + while True: + charcode = pickle_file.read(1) + if not charcode: # Indicates exhaustion of the data stream + break - opcode_argument = None - if opcode.arg is not None: - if charcode in constants.LENGTH_PREFIXED_OPCODES: try: - match charcode: - case c if c in constants.OPCODES_4BYTE_LEN: - length = int.from_bytes(pickle_file.read(4), byteorder="little") - case c if c in constants.OPCODES_8BYTE_LEN: - length = int.from_bytes(pickle_file.read(8), byteorder="little") - case c if c in constants.OPCODES_1BYTE_LEN: - length = int.from_bytes(pickle_file.read(1), byteorder="little") - case _: - length = 0 - pickle_file.seek(length, os.SEEK_CUR) - except (AttributeError, io.UnsupportedOperation, OSError): - try: - opcode_argument = opcode.arg.reader(pickle_file) - except (ValueError, pickle.UnpicklingError) as e: - raise UnsafePickleDetectedError( - f"Parser error during security scan: {e}" - ) from e - except ( - IndexError, - AttributeError, - EOFError, - TypeError, - ImportError, - ): - continue - else: - try: - opcode_argument = opcode.arg.reader(pickle_file) - except (ValueError, pickle.UnpicklingError) as e: - raise UnsafePickleDetectedError( - f"Parser error during security scan: {e}" - ) from e - except ( - IndexError, - AttributeError, - EOFError, - TypeError, - ImportError, - ): - continue + opcode = constants.OPCODES_INFO_INT.get(charcode[0]) + except IndexError: + continue # Skip invalid opcode bytes - # We only yield opcodes that declare strings and have arguments - should_yield = False - for relevant_opcode_substr in constants.OPCODE_SUBSTRS_THAT_DECLARE_STRINGS: - if relevant_opcode_substr in opcode.name: - should_yield = True - break + if opcode is None: + # We skip processing unknown opcodes + continue - if ( - should_yield - and opcode_argument is not None # Exclude opcodes without arguments - ): - # This is to be careful while processing opcode arguments. This was - # borrowed from what works in the chunked version. - if isinstance(opcode_argument, (str, bytes)) and len(opcode_argument) > 1: - yield opcode, opcode_argument - elif isinstance(opcode_argument, tuple): - yield opcode, opcode_argument + opcode_argument = None + if opcode.arg is not None: + if charcode in constants.LENGTH_PREFIXED_OPCODES: + try: + match charcode: + case c if c in constants.OPCODES_4BYTE_LEN: + length = int.from_bytes( + pickle_file.read(4), byteorder="little" + ) + case c if c in constants.OPCODES_8BYTE_LEN: + length = int.from_bytes( + pickle_file.read(8), byteorder="little" + ) + case c if c in constants.OPCODES_1BYTE_LEN: + length = int.from_bytes( + pickle_file.read(1), byteorder="little" + ) + case _: + length = 0 + pickle_file.seek(length, os.SEEK_CUR) + except (AttributeError, io.UnsupportedOperation, OSError): + try: + opcode_argument = opcode.arg.reader(pickle_file) + except (ValueError, pickle.UnpicklingError) as e: + raise UnsafePickleDetectedError( + f"Parser error during security scan: {e}" + ) from e + except ( + IndexError, + AttributeError, + EOFError, + TypeError, + ImportError, + ): + continue + else: + try: + opcode_argument = opcode.arg.reader(pickle_file) + except (ValueError, pickle.UnpicklingError) as e: + raise UnsafePickleDetectedError( + f"Parser error during security scan: {e}" + ) from e + except ( + IndexError, + AttributeError, + EOFError, + TypeError, + ImportError, + ): + continue + + # We only yield opcodes that declare strings and have arguments + should_yield = False + for relevant_opcode_substr in constants.OPCODE_SUBSTRS_THAT_DECLARE_STRINGS: + if relevant_opcode_substr in opcode.name: + should_yield = True + break - if charcode == b".": - break + if ( + should_yield + and opcode_argument is not None # Exclude opcodes without arguments + ): + # This is to be careful while processing opcode arguments. This was + # borrowed from what works in the chunked version. + if isinstance(opcode_argument, (str, bytes)) and len(opcode_argument) > 1: + yield opcode, opcode_argument + elif isinstance(opcode_argument, tuple): + yield opcode, opcode_argument + + if charcode == b".": + break def _custom_chunked_genops( pickle_file: IO[bytes], chunk_range: Tuple[int, int], ) -> Iterator[tuple[pickletools.OpcodeInfo, Any | None]]: - """Generates string-declaring opcodes and arguments from a chunk. + """Generates string-declaring opcodes and arguments from a chunk. - This function reads a specific byte range (chunk) of the pickle bytecode - and yields opcodes that are known to declare strings, along with their - arguments. It's designed to be used in parallel for large pickle files. + This function reads a specific byte range (chunk) of the pickle bytecode + and yields opcodes that are known to declare strings, along with their + arguments. It's designed to be used in parallel for large pickle files. - Args: - pickle_file: The pickle data stream to generate opcodes from. - chunk_range: A tuple (start, end) defining the byte range to process. + Args: + pickle_file: The pickle data stream to generate opcodes from. + chunk_range: A tuple (start, end) defining the byte range to process. - Yields: - A tuple of (opcode, opcode_argument) for each string-declaring opcode. - """ - pickle_file.seek(chunk_range[0]) + Yields: + A tuple of (opcode, opcode_argument) for each string-declaring opcode. + """ + pickle_file.seek(chunk_range[0]) - while True: - current_file_position = pickle_file.tell() - if not (chunk_range[0] <= current_file_position < chunk_range[1]): - break + while True: + current_file_position = pickle_file.tell() + if not (chunk_range[0] <= current_file_position < chunk_range[1]): + break - charcode = pickle_file.read(1) - if not charcode: # Indicates exhaustion of the data stream - break + charcode = pickle_file.read(1) + if not charcode: # Indicates exhaustion of the data stream + break - try: - opcode = constants.OPCODES_INFO_INT.get(charcode[0]) - except IndexError: - continue # Skip invalid opcode bytes - - if opcode is None: - # We skip processing unknown opcodes - if not charcode: - break - continue - - opcode_argument = None - if opcode.arg is not None: - pos_before_arg_read = pickle_file.tell() - try: - opcode_argument = opcode.arg.reader(pickle_file) - new_pos = pickle_file.tell() - - # Ensure we don't read past the chunk boundary accidentally - if new_pos > chunk_range[1]: - pickle_file.seek(pos_before_arg_read) - continue - - except (ValueError, pickle.UnpicklingError) as e: - raise UnsafePickleDetectedError( - f"Parser error during security scan: {e}" - ) from e - except ( - IndexError, - AttributeError, - EOFError, - TypeError, - ImportError, - ): - # Continue if we can't read the argument within the chunk - pickle_file.seek(pos_before_arg_read) - continue - - # We only yield opcodes that declare strings and have arguments - should_yield = False - for relevant_opcode_substr in constants.OPCODE_SUBSTRS_THAT_DECLARE_STRINGS: - if relevant_opcode_substr in opcode.name: - should_yield = True - break + try: + opcode = constants.OPCODES_INFO_INT.get(charcode[0]) + except IndexError: + continue # Skip invalid opcode bytes - if ( - should_yield - and opcode_argument is not None # Exclude opcodes without arguments - ): - # Filter to ensure the argument is string-like if needed - if isinstance(opcode_argument, (str, bytes)) and len(opcode_argument) > 1: - yield opcode, opcode_argument - elif isinstance( - opcode_argument, tuple - ): # Sometimes these arguments are memoized tuples - yield opcode, opcode_argument + if opcode is None: + # We skip processing unknown opcodes + if not charcode: + break + continue + + opcode_argument = None + if opcode.arg is not None: + pos_before_arg_read = pickle_file.tell() + try: + opcode_argument = opcode.arg.reader(pickle_file) + new_pos = pickle_file.tell() + + # Ensure we don't read past the chunk boundary accidentally + if new_pos > chunk_range[1]: + pickle_file.seek(pos_before_arg_read) + continue + + except (ValueError, pickle.UnpicklingError) as e: + raise UnsafePickleDetectedError( + f"Parser error during security scan: {e}" + ) from e + except ( + IndexError, + AttributeError, + EOFError, + TypeError, + ImportError, + ): + # Continue if we can't read the argument within the chunk + pickle_file.seek(pos_before_arg_read) + continue + + # We only yield opcodes that declare strings and have arguments + should_yield = False + for relevant_opcode_substr in constants.OPCODE_SUBSTRS_THAT_DECLARE_STRINGS: + if relevant_opcode_substr in opcode.name: + should_yield = True + break - if charcode == b".": - break + if ( + should_yield + and opcode_argument is not None # Exclude opcodes without arguments + ): + # Filter to ensure the argument is string-like if needed + if isinstance(opcode_argument, (str, bytes)) and len(opcode_argument) > 1: + yield opcode, opcode_argument + elif isinstance( + opcode_argument, tuple + ): # Sometimes these arguments are memoized tuples + yield opcode, opcode_argument + + if charcode == b".": + break def _process_chunk_for_generate_ops( @@ -265,50 +268,52 @@ def _process_chunk_for_generate_ops( is_shared_memory: bool = False, abort_event: Optional[Any] = None, ) -> Set[str]: - """Helper function to process a chunk of pickle data.""" - chunked_operands = set() - try: - if is_shared_memory: - shm = shared_memory.SharedMemory(name=pickle_data_source) # pyrefly: ignore[bad-argument-type] - try: - # Use BytesIO on the memoryview for compatibility with - # _custom_chunked_genops - data_view = shm.buf - with io.BytesIO(data_view) as f: # pyrefly: ignore[bad-argument-type] - for _, operand in _custom_chunked_genops(f, chunk_range): - if abort_event and abort_event.is_set(): - break - if operand is None: - continue - operand_str = str(operand) - chunked_operands.add(operand_str) - if abort_event: - if utils.is_unsafe_or_suspicious(operand_str): - abort_event.set() - break - finally: - shm.close() - else: - with open(pickle_data_source, "rb") as f: - f.seek(chunk_range[0]) - chunk_data = f.read(chunk_range[1] - chunk_range[0]) - with io.BytesIO(chunk_data) as memory_f: - for _, operand in _custom_chunked_genops( - memory_f, (0, len(chunk_data)) - ): - if abort_event and abort_event.is_set(): - break - if operand is None: - continue - operand_str = str(operand) - chunked_operands.add(operand_str) - if abort_event: - if utils.is_unsafe_or_suspicious(operand_str): - abort_event.set() - break - except StopIteration: - pass - return chunked_operands + """Helper function to process a chunk of pickle data.""" + chunked_operands = set() + try: + if is_shared_memory: + shm = shared_memory.SharedMemory( + name=pickle_data_source + ) # pyrefly: ignore[bad-argument-type] + try: + # Use BytesIO on the memoryview for compatibility with + # _custom_chunked_genops + data_view = shm.buf + with io.BytesIO(data_view) as f: # pyrefly: ignore[bad-argument-type] + for _, operand in _custom_chunked_genops(f, chunk_range): + if abort_event and abort_event.is_set(): + break + if operand is None: + continue + operand_str = str(operand) + chunked_operands.add(operand_str) + if abort_event: + if utils.is_unsafe_or_suspicious(operand_str): + abort_event.set() + break + finally: + shm.close() + else: + with open(pickle_data_source, "rb") as f: + f.seek(chunk_range[0]) + chunk_data = f.read(chunk_range[1] - chunk_range[0]) + with io.BytesIO(chunk_data) as memory_f: + for _, operand in _custom_chunked_genops( + memory_f, (0, len(chunk_data)) + ): + if abort_event and abort_event.is_set(): + break + if operand is None: + continue + operand_str = str(operand) + chunked_operands.add(operand_str) + if abort_event: + if utils.is_unsafe_or_suspicious(operand_str): + abort_event.set() + break + except StopIteration: + pass + return chunked_operands def generate_ops_from_file( @@ -317,285 +322,297 @@ def generate_ops_from_file( pickle_length: Optional[int] = None, fail_fast: Optional[bool] = DEFAULT_FAIL_FAST, ) -> Set[str]: - """Returns opcodes that declare strings from a path or shared memory. - - Args: - pickle_file_path: The path to the pickle file. - shm_name: Optional name of the shared memory block. - pickle_length: Optional length of the pickle data. - fail_fast: Whether to fail fast on first unsafe or suspicious match. - - Returns: - genops_output: The operands associated with the opcodes that declare - strings. - """ - filtered_operands = set() - num_workers = utils.get_optimal_workers(pickle_length) # pyrefly: ignore[bad-argument-type] - - if ( - pickle_length < constants.MIN_SIZE_FOR_CHUNKING # pyrefly: ignore[unsupported-operation] - or not utils.is_sys_executable_patched() - ): - if shm_name: - shm = shared_memory.SharedMemory(name=shm_name) - pickle_bytes = bytes(shm.buf[:pickle_length]) # pyrefly: ignore[unsupported-operation] - else: - with open(pickle_file_path, "rb") as f: - pickle_bytes = f.read() - try: - for _, operand in _custom_genops(pickle_bytes): - if operand is None: - continue - operand_str = str(operand) - filtered_operands.add(operand_str) - if fail_fast: - if utils.is_unsafe_or_suspicious(operand_str): - break - except StopIteration: - pass - return filtered_operands - else: - # Divide into constants.MAX_NUM_CHUNKS for larger files - chunk_size = math.ceil(pickle_length / num_workers) # pyrefly: ignore[unsupported-operation] - ranges = [] - for chunk_index in range(num_workers): - chunk_start_size = chunk_index * chunk_size - # Extend the chunk end by CHUNK_OVERLAP, but don't exceed pickle_length - chunk_end = min( # pyrefly: ignore[bad-specialization] - chunk_start_size + chunk_size + constants.CHUNK_OVERLAP, pickle_length - ) - if chunk_start_size < pickle_length: # pyrefly: ignore[unsupported-operation] - ranges.append((chunk_start_size, chunk_end)) - if chunk_end == pickle_length: - break # Last chunk reaches the end - - ctx = multiprocessing.get_context("spawn") - manager = ctx.Manager() if fail_fast else None - abort_event = manager.Event() if manager else None + """Returns opcodes that declare strings from a path or shared memory. - try: - with concurrent.futures.ProcessPoolExecutor( - max_workers=num_workers, mp_context=ctx - ) as executor: - future_to_range_tuple = { - executor.submit( - _process_chunk_for_generate_ops, - shm_name if shm_name else pickle_file_path, - range_tuple, - is_shared_memory=bool(shm_name), - abort_event=abort_event, - ): range_tuple - for range_tuple in ranges - } - for future in concurrent.futures.as_completed(future_to_range_tuple): - try: - chunk_results = future.result() - filtered_operands.update(chunk_results) - if fail_fast: - has_blocked_operand = False - for op in chunk_results: - if utils.is_unsafe_or_suspicious(op): - has_blocked_operand = True - break - if has_blocked_operand: - if abort_event: - abort_event.set() - for f in future_to_range_tuple: - f.cancel() - break - except ( - EOFError, - ValueError, - IndexError, - TypeError, - ) as exc: - logging.exception( - "Error processing chunk %s: %s", - future_to_range_tuple[future], - exc, + Args: + pickle_file_path: The path to the pickle file. + shm_name: Optional name of the shared memory block. + pickle_length: Optional length of the pickle data. + fail_fast: Whether to fail fast on first unsafe or suspicious match. + + Returns: + genops_output: The operands associated with the opcodes that declare + strings. + """ + filtered_operands = set() + num_workers = utils.get_optimal_workers( + pickle_length + ) # pyrefly: ignore[bad-argument-type] + + if ( + pickle_length + < constants.MIN_SIZE_FOR_CHUNKING # pyrefly: ignore[unsupported-operation] + or not utils.is_sys_executable_patched() + ): + if shm_name: + shm = shared_memory.SharedMemory(name=shm_name) + pickle_bytes = bytes( + shm.buf[:pickle_length] + ) # pyrefly: ignore[unsupported-operation] + else: + with open(pickle_file_path, "rb") as f: + pickle_bytes = f.read() + try: + for _, operand in _custom_genops(pickle_bytes): + if operand is None: + continue + operand_str = str(operand) + filtered_operands.add(operand_str) + if fail_fast: + if utils.is_unsafe_or_suspicious(operand_str): + break + except StopIteration: + pass + return filtered_operands + else: + # Divide into constants.MAX_NUM_CHUNKS for larger files + chunk_size = math.ceil( + pickle_length / num_workers + ) # pyrefly: ignore[unsupported-operation] + ranges = [] + for chunk_index in range(num_workers): + chunk_start_size = chunk_index * chunk_size + # Extend the chunk end by CHUNK_OVERLAP, but don't exceed pickle_length + chunk_end = min( # pyrefly: ignore[bad-specialization] + chunk_start_size + chunk_size + constants.CHUNK_OVERLAP, pickle_length ) - finally: - if manager: - manager.shutdown() + if ( + chunk_start_size < pickle_length + ): # pyrefly: ignore[unsupported-operation] + ranges.append((chunk_start_size, chunk_end)) + if chunk_end == pickle_length: + break # Last chunk reaches the end + + ctx = multiprocessing.get_context("spawn") + manager = ctx.Manager() if fail_fast else None + abort_event = manager.Event() if manager else None + + try: + with concurrent.futures.ProcessPoolExecutor( + max_workers=num_workers, mp_context=ctx + ) as executor: + future_to_range_tuple = { + executor.submit( + _process_chunk_for_generate_ops, + shm_name if shm_name else pickle_file_path, + range_tuple, + is_shared_memory=bool(shm_name), + abort_event=abort_event, + ): range_tuple + for range_tuple in ranges + } + for future in concurrent.futures.as_completed(future_to_range_tuple): + try: + chunk_results = future.result() + filtered_operands.update(chunk_results) + if fail_fast: + has_blocked_operand = False + for op in chunk_results: + if utils.is_unsafe_or_suspicious(op): + has_blocked_operand = True + break + if has_blocked_operand: + if abort_event: + abort_event.set() + for f in future_to_range_tuple: + f.cancel() + break + except ( + EOFError, + ValueError, + IndexError, + TypeError, + ) as exc: + logging.exception( + "Error processing chunk %s: %s", + future_to_range_tuple[future], + exc, + ) + finally: + if manager: + manager.shutdown() - return filtered_operands + return filtered_operands def generate_ops( pickle_bytes: bytes | IO[bytes], fail_fast: Optional[bool] = DEFAULT_FAIL_FAST, ) -> Set[str]: - """Returns string-declaring opcodes. + """Returns string-declaring opcodes. - Args: - pickle_bytes: The pickle bytecode or stream to yield opcode information for. - fail_fast: Whether to fail fast on first unsafe or suspicious match. + Args: + pickle_bytes: The pickle bytecode or stream to yield opcode information for. + fail_fast: Whether to fail fast on first unsafe or suspicious match. - Returns: - genops_output: The operands associated with the opcodes that declare - strings. - """ + Returns: + genops_output: The operands associated with the opcodes that declare + strings. + """ - filtered_operands = set() - original_pos = None - if not isinstance(pickle_bytes, bytes): - original_pos = pickle_bytes.tell() + filtered_operands = set() + original_pos = None + if not isinstance(pickle_bytes, bytes): + original_pos = pickle_bytes.tell() - try: try: - for _, operand in _custom_genops(pickle_bytes): # pyrefly: ignore[bad-argument-type] - if operand is None: - continue - operand_str = str(operand) - filtered_operands.add(operand_str) - if fail_fast: - if utils.is_unsafe_or_suspicious(operand_str): - break - except StopIteration: - pass - return filtered_operands - finally: - if original_pos is not None: - pickle_bytes.seek(original_pos) # pyrefly: ignore[missing-attribute] + try: + for _, operand in _custom_genops( + pickle_bytes + ): # pyrefly: ignore[bad-argument-type] + if operand is None: + continue + operand_str = str(operand) + filtered_operands.add(operand_str) + if fail_fast: + if utils.is_unsafe_or_suspicious(operand_str): + break + except StopIteration: + pass + return filtered_operands + finally: + if original_pos is not None: + pickle_bytes.seek(original_pos) # pyrefly: ignore[missing-attribute] def get_class_instantiations( pickle_bytes: bytes | BinaryIO, ) -> tuple[io.StringIO, bool, bool]: - """Gets the class instantiations from a pickle file/stream. - - Args: - pickle_bytes: The pickle bytecode or stream to disassemble. - - Returns: - A tuple containing: - - 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 - if isinstance(pickle_bytes, bytes): - pickle_stream = io.BytesIO(pickle_bytes) - else: - original_pos = pickle_bytes.tell() - pickle_stream = pickle_bytes - - try: + """Gets the class instantiations from a pickle file/stream. + + Args: + pickle_bytes: The pickle bytecode or stream to disassemble. + + Returns: + A tuple containing: + - 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 + if isinstance(pickle_bytes, bytes): + pickle_stream = io.BytesIO(pickle_bytes) + else: + original_pos = pickle_bytes.tell() + pickle_stream = pickle_bytes + try: - factory = picklemagic.FakeClassFactory([], picklemagic.FakeWarning) - - # Instead of using safe_loads, we do this to get the - # has_blocked_unsafe_build_instr boolean properly. - unpickler = picklemagic.SafeUnpickler( - pickle_stream, - class_factory=factory, - safe_modules=constants.SAFE_STRINGS, - unsafe_modules=constants.UNSAFE_STRINGS, - ) - factory.default.unpickler = unpickler - unpickler.load() + try: + factory = picklemagic.FakeClassFactory([], picklemagic.FakeWarning) + + # Instead of using safe_loads, we do this to get the + # has_blocked_unsafe_build_instr boolean properly. + unpickler = picklemagic.SafeUnpickler( + pickle_stream, + class_factory=factory, + safe_modules=constants.SAFE_STRINGS, + unsafe_modules=constants.UNSAFE_STRINGS, + ) + factory.default.unpickler = unpickler + unpickler.load() - except ( - ValueError, - AttributeError, - TypeError, - picklemagic.FakeUnpicklingError, - pickle.UnpicklingError, - IndexError, - EOFError, - KeyError, - struct.error, - ) 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) # pyrefly: ignore[missing-attribute] - - was_unsafe_build_blocked = False - if unpickler: - was_unsafe_build_blocked = getattr( - unpickler, "has_blocked_unsafe_build_instr", False - ) + except ( + ValueError, + AttributeError, + TypeError, + picklemagic.FakeUnpicklingError, + pickle.UnpicklingError, + IndexError, + EOFError, + KeyError, + struct.error, + ) 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) # pyrefly: ignore[missing-attribute] + + was_unsafe_build_blocked = False + if unpickler: + was_unsafe_build_blocked = getattr( + unpickler, "has_blocked_unsafe_build_instr", False + ) - return picklemagic_output, was_unsafe_build_blocked, has_scan_error + 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.""" - 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) # pyrefly: ignore[bad-argument-type] - return _reclassify_with_resolution(safe, unsafe, suspicious, unknown) + """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 + ) # pyrefly: ignore[bad-argument-type] + return _reclassify_with_resolution(safe, unsafe, suspicious, unknown) 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() - - 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) + """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() + + 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) - return safe_results, unsafe_results, suspicious_results, unknown_results + return safe_results, unsafe_results, suspicious_results, unknown_results def _reclassify_with_resolution( @@ -604,125 +621,127 @@ def _reclassify_with_resolution( 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( - unsafe_results, suspicious_results, unknown_results - ) - resolved_results = utils.resolve_library_modules_from_results(all_results) - - # Re-categorize the resolved results - new_safe_results = set() - new_unsafe_results = set() - new_suspicious_results = set() - new_unknown_results = set() - is_denylisted = False - - for result in resolved_results: - if any(result.startswith(denied_item) for denied_item in deny_list): - new_unsafe_results.add(result) - is_denylisted = True - continue - - if any(result.startswith(allowed_item) for allowed_item in allow_list): - new_safe_results.add(result) - continue - - if result == "builtins": - new_unknown_results.add(result) - continue - - # Classify the resolved result - classification = utils.classify_class_name(result) - - if classification == utils.Classification.SAFE: - new_safe_results.add(result) - elif classification == utils.Classification.UNSAFE: - new_unsafe_results.add(result) - elif classification == utils.Classification.SUSPICIOUS: - new_suspicious_results.add(result) - elif classification == utils.Classification.UNKNOWN: - # Fallback: Check against original categories if - # classify_class_name returns UNKNOWN. - if result in unsafe_results: - new_unsafe_results.add(result) - elif result in suspicious_results: - new_suspicious_results.add(result) - elif result in safe_results: - new_safe_results.add(result) - else: - new_unknown_results.add(result) - - return ScanResults( - safe_results=new_safe_results, - unsafe_results=new_unsafe_results, - suspicious_results=new_suspicious_results, - unknown_results=new_unknown_results, - is_denylisted=is_denylisted, - ) + """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( + unsafe_results, suspicious_results, unknown_results + ) + resolved_results = utils.resolve_library_modules_from_results(all_results) + + # Re-categorize the resolved results + new_safe_results = set() + new_unsafe_results = set() + new_suspicious_results = set() + new_unknown_results = set() + is_denylisted = False + + for result in resolved_results: + if any(result.startswith(denied_item) for denied_item in deny_list): + new_unsafe_results.add(result) + is_denylisted = True + continue -def strict_security_scan(pickle_bytes: bytes | BinaryIO) -> bool: - """Strict security scan for malicious content in pickle files. + if any(result.startswith(allowed_item) for allowed_item in allow_list): + new_safe_results.add(result) + continue - Args: - pickle_bytes: Pickle bytecode or stream to scan. + if result == "builtins": + new_unknown_results.add(result) + continue - Returns: - True if the pickle file is dangerous, False otherwise. - """ + # Classify the resolved result + classification = utils.classify_class_name(result) + + if classification == utils.Classification.SAFE: + new_safe_results.add(result) + elif classification == utils.Classification.UNSAFE: + new_unsafe_results.add(result) + elif classification == utils.Classification.SUSPICIOUS: + new_suspicious_results.add(result) + elif classification == utils.Classification.UNKNOWN: + # Fallback: Check against original categories if + # classify_class_name returns UNKNOWN. + if result in unsafe_results: + new_unsafe_results.add(result) + elif result in suspicious_results: + new_suspicious_results.add(result) + elif result in safe_results: + new_safe_results.add(result) + else: + new_unknown_results.add(result) + + return ScanResults( + safe_results=new_safe_results, + unsafe_results=new_unsafe_results, + suspicious_results=new_suspicious_results, + unknown_results=new_unknown_results, + is_denylisted=is_denylisted, + ) - original_pos = None - if not isinstance(pickle_bytes, bytes): - original_pos = pickle_bytes.tell() - try: - unsafe_and_suspicious_strings = constants.UNSAFE_STRINGS.union( - constants.SUSPICIOUS_STRINGS - ) - for _, operand in _custom_genops(pickle_bytes): # pyrefly: ignore[bad-argument-type] - if operand is None: - continue - stmt = str(operand) - for pattern in unsafe_and_suspicious_strings: - if re.search(pattern, stmt): - return True - - # Seek back the stream before running picklemagic - if original_pos is not None and hasattr(pickle_bytes, "seek"): - try: - pickle_bytes.seek(original_pos) - except (OSError, AttributeError, io.UnsupportedOperation): - 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, has_scan_error = ( - get_class_instantiations(pickle_bytes) - ) +def strict_security_scan(pickle_bytes: bytes | BinaryIO) -> bool: + """Strict security scan for malicious content in pickle files. - if was_unsafe_build_blocked or has_scan_error: - return True + Args: + pickle_bytes: Pickle bytecode or stream to scan. - instantiations = instantiations_output.getvalue().split("\n") - for instantiation in instantiations: - if re.search(utils.unknown_pattern, instantiation): - return True - # This is a noisy but necessary check for a small number of cases where - # a library is not explicitly imported but is used in a - # class instantiation in a suspicious manner. - if re.search(utils.suspicious_pattern, instantiation): - return True - finally: - if original_pos is not None and hasattr(pickle_bytes, "seek"): - try: - pickle_bytes.seek(original_pos) - except (OSError, AttributeError, io.UnsupportedOperation): - logging.debug("Failed to reset file pointer after strict scan.") + Returns: + True if the pickle file is dangerous, False otherwise. + """ + + original_pos = None + if not isinstance(pickle_bytes, bytes): + original_pos = pickle_bytes.tell() + + try: + unsafe_and_suspicious_strings = constants.UNSAFE_STRINGS.union( + constants.SUSPICIOUS_STRINGS + ) + for _, operand in _custom_genops( + pickle_bytes + ): # pyrefly: ignore[bad-argument-type] + if operand is None: + continue + stmt = str(operand) + for pattern in unsafe_and_suspicious_strings: + if re.search(pattern, stmt): + return True + + # Seek back the stream before running picklemagic + if original_pos is not None and hasattr(pickle_bytes, "seek"): + try: + pickle_bytes.seek(original_pos) + except (OSError, AttributeError, io.UnsupportedOperation): + 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, has_scan_error = ( + get_class_instantiations(pickle_bytes) + ) + + if was_unsafe_build_blocked or has_scan_error: + return True + + instantiations = instantiations_output.getvalue().split("\n") + for instantiation in instantiations: + if re.search(utils.unknown_pattern, instantiation): + return True + # This is a noisy but necessary check for a small number of cases where + # a library is not explicitly imported but is used in a + # class instantiation in a suspicious manner. + if re.search(utils.suspicious_pattern, instantiation): + return True + finally: + if original_pos is not None and hasattr(pickle_bytes, "seek"): + try: + pickle_bytes.seek(original_pos) + except (OSError, AttributeError, io.UnsupportedOperation): + logging.debug("Failed to reset file pointer after strict scan.") - return False + return False def is_unsafe( @@ -730,64 +749,64 @@ def is_unsafe( number_of_unsafe_results: int, number_of_suspicious_results: int, ) -> bool: - """Conditional check for safeness. - - Args: - number_of_safe_results: Number of safe results from the security scan. - number_of_unsafe_results: Number of unsafe results from the security scan. - number_of_suspicious_results: Number of suspicious results from the security - scan. - - Returns: - True if the pickle file is dangerous, False otherwise. - """ - if number_of_unsafe_results == 0 and number_of_suspicious_results == 0: - return False + """Conditional check for safeness. + + Args: + number_of_safe_results: Number of safe results from the security scan. + number_of_unsafe_results: Number of unsafe results from the security scan. + number_of_suspicious_results: Number of suspicious results from the security + scan. - # We halve the weight of suspicious results to lower false positives - # caused by greedy matches of unknown method-like strings (Ex. "google.com") - if ( - number_of_suspicious_results + number_of_unsafe_results - >= number_of_safe_results - ): - return True + Returns: + True if the pickle file is dangerous, False otherwise. + """ + if number_of_unsafe_results == 0 and number_of_suspicious_results == 0: + return False + + # We halve the weight of suspicious results to lower false positives + # caused by greedy matches of unknown method-like strings (Ex. "google.com") + if ( + number_of_suspicious_results + number_of_unsafe_results + >= number_of_safe_results + ): + return True - sum_of_unsafe_and_suspicious_results = ( - number_of_unsafe_results + 0.5 * number_of_suspicious_results - ) + sum_of_unsafe_and_suspicious_results = ( + number_of_unsafe_results + 0.5 * number_of_suspicious_results + ) - unsafe = (sum_of_unsafe_and_suspicious_results > number_of_safe_results) or ( - number_of_safe_results == 0 and sum_of_unsafe_and_suspicious_results >= 1 - ) + unsafe = (sum_of_unsafe_and_suspicious_results > number_of_safe_results) or ( + number_of_safe_results == 0 and sum_of_unsafe_and_suspicious_results >= 1 + ) - return unsafe + return unsafe def picklemagic_scan( pickle_bytes: bytes, ) -> ScanResults: - """Picklemagic scan for malicious content in pickle files. + """Picklemagic scan for malicious content in pickle files. - Args: - pickle_bytes: Pickle bytecode to scan. + Args: + pickle_bytes: Pickle bytecode to scan. - Returns: - A ScanResults object. - """ - picklemagic_output, was_unsafe_build_blocked, has_scan_error = ( - get_class_instantiations(pickle_bytes) - ) + Returns: + A ScanResults object. + """ + picklemagic_output, was_unsafe_build_blocked, has_scan_error = ( + get_class_instantiations(pickle_bytes) + ) - results = categorize_strings(picklemagic_output, use_picklemagic=True) + results = categorize_strings(picklemagic_output, use_picklemagic=True) - if was_unsafe_build_blocked: - # 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") + if was_unsafe_build_blocked: + # 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 + return results def genops_scan( @@ -797,38 +816,40 @@ def genops_scan( fail_fast: Optional[bool] = DEFAULT_FAIL_FAST, pickle_length: Optional[int] = None, ) -> ScanResults: - """Genops scan for malicious content in pickle files. - - Args: - pickle_bytes: Pickle bytecode to scan. - pickle_file_path: Optional path to the pickle file for streaming scan. - shm_name: Optional name of the shared memory block. - fail_fast: Whether to fail fast on first unsafe or suspicious match. - pickle_length: Optional length of the pickle data. - - Returns: - A ScanResults object. - """ - resolved_pickle_length = ( - pickle_length if pickle_length is not None else len(pickle_bytes) # pyrefly: ignore[bad-argument-type] - ) - if shm_name: - genops_output = generate_ops_from_file( - "", - shm_name=shm_name, - pickle_length=resolved_pickle_length, - fail_fast=fail_fast, - ) - elif pickle_file_path: - genops_output = generate_ops_from_file( - pickle_file_path, - pickle_length=resolved_pickle_length, - fail_fast=fail_fast, + """Genops scan for malicious content in pickle files. + + Args: + pickle_bytes: Pickle bytecode to scan. + pickle_file_path: Optional path to the pickle file for streaming scan. + shm_name: Optional name of the shared memory block. + fail_fast: Whether to fail fast on first unsafe or suspicious match. + pickle_length: Optional length of the pickle data. + + Returns: + A ScanResults object. + """ + resolved_pickle_length = ( + pickle_length + if pickle_length is not None + else len(pickle_bytes) # pyrefly: ignore[bad-argument-type] ) - else: - genops_output = generate_ops(pickle_bytes, fail_fast=fail_fast) - results = categorize_strings(genops_output) - return results + if shm_name: + genops_output = generate_ops_from_file( + "", + shm_name=shm_name, + pickle_length=resolved_pickle_length, + fail_fast=fail_fast, + ) + elif pickle_file_path: + genops_output = generate_ops_from_file( + pickle_file_path, + pickle_length=resolved_pickle_length, + fail_fast=fail_fast, + ) + else: + genops_output = generate_ops(pickle_bytes, fail_fast=fail_fast) + results = categorize_strings(genops_output) + return results def score_results( @@ -837,34 +858,34 @@ def score_results( suspicious_results: Set[str], unknown_results: Set[str], ) -> Tuple[int, int, int, int]: - """Count the results from the security scan. - - Args: - safe_results: List of safe strings. - unsafe_results: List of unsafe strings. - suspicious_results: List of suspicious strings. - unknown_results: List of unknown strings. + """Count the results from the security scan. - Returns: - A tuple of safe, unsafe, suspicious, and unknown scores. - """ - - number_of_safe_results = len(safe_results) - number_of_unsafe_results = len(unsafe_results) - number_of_suspicious_results = len(suspicious_results) - number_of_unknown_results = len(unknown_results) + Args: + safe_results: List of safe strings. + unsafe_results: List of unsafe strings. + suspicious_results: List of suspicious strings. + unknown_results: List of unknown strings. - safe_score = math.log(number_of_safe_results + 1) * 2 - unsafe_score = math.log(number_of_unsafe_results + 1) * 4 - suspicious_score = math.log(number_of_suspicious_results + 1) * 3 - unknown_score = math.log(number_of_unknown_results + 1) * 1 + Returns: + A tuple of safe, unsafe, suspicious, and unknown scores. + """ - return ( - round(safe_score), - round(unsafe_score), - round(suspicious_score), - round(unknown_score), - ) + number_of_safe_results = len(safe_results) + number_of_unsafe_results = len(unsafe_results) + number_of_suspicious_results = len(suspicious_results) + number_of_unknown_results = len(unknown_results) + + safe_score = math.log(number_of_safe_results + 1) * 2 + unsafe_score = math.log(number_of_unsafe_results + 1) * 4 + suspicious_score = math.log(number_of_suspicious_results + 1) * 3 + unknown_score = math.log(number_of_unknown_results + 1) * 1 + + return ( + round(safe_score), + round(unsafe_score), + round(suspicious_score), + round(unknown_score), + ) def apply_approach( @@ -875,70 +896,70 @@ def apply_approach( fail_fast: Optional[bool] = DEFAULT_FAIL_FAST, pickle_length: Optional[int] = None, ) -> Dict[str, int]: - """Applies the given scan approach to the data. - - Args: - scan_approach: The scan approach to apply to the data. - pickle_bytes: The data to scan. - pickle_file_path: Optional path to the pickle file for streaming scan. - shm_name: Optional name of the shared memory block. - fail_fast: Whether to fail fast on first unsafe or suspicious match. - pickle_length: Optional length of the pickle data. - - Returns: - A dictionary of the resulting scores. - """ - if scan_approach is genops_scan: - results = scan_approach( - pickle_bytes, - pickle_file_path=pickle_file_path, - shm_name=shm_name, - fail_fast=fail_fast, - pickle_length=pickle_length, + """Applies the given scan approach to the data. + + Args: + scan_approach: The scan approach to apply to the data. + pickle_bytes: The data to scan. + pickle_file_path: Optional path to the pickle file for streaming scan. + shm_name: Optional name of the shared memory block. + fail_fast: Whether to fail fast on first unsafe or suspicious match. + pickle_length: Optional length of the pickle data. + + Returns: + A dictionary of the resulting scores. + """ + if scan_approach is genops_scan: + results = scan_approach( + pickle_bytes, + pickle_file_path=pickle_file_path, + shm_name=shm_name, + fail_fast=fail_fast, + pickle_length=pickle_length, + ) + else: + results = scan_approach(pickle_bytes) + + if DEBUG_MODE: + logging.info("Scan approach: %s", scan_approach.__name__) + logging.info(" Safe results: %s", results.safe_results) + logging.info(" Unsafe results: %s", results.unsafe_results) + logging.info(" Suspicious results: %s", results.suspicious_results) + logging.info(" Unknown results: %s\n", results.unknown_results) + + ( + number_of_safe_results, + number_of_unsafe_results, + number_of_suspicious_results, + number_of_unknown_results, + ) = score_results( + results.safe_results, + results.unsafe_results, + results.suspicious_results, + results.unknown_results, ) - else: - results = scan_approach(pickle_bytes) - - if DEBUG_MODE: - logging.info("Scan approach: %s", scan_approach.__name__) - logging.info(" Safe results: %s", results.safe_results) - logging.info(" Unsafe results: %s", results.unsafe_results) - logging.info(" Suspicious results: %s", results.suspicious_results) - logging.info(" Unknown results: %s\n", results.unknown_results) - - ( - number_of_safe_results, - number_of_unsafe_results, - number_of_suspicious_results, - number_of_unknown_results, - ) = score_results( - results.safe_results, - results.unsafe_results, - results.suspicious_results, - results.unknown_results, - ) - scores = { - "unsafe": number_of_unsafe_results, - "suspicious": number_of_suspicious_results, - "unknown": number_of_unknown_results, - } - should_fail_fast = fail_fast and ( - number_of_unsafe_results > 0 and number_of_suspicious_results == 0 - ) - if ( - results.is_denylisted - or should_fail_fast - or is_unsafe( - number_of_safe_results, - number_of_unsafe_results, - number_of_suspicious_results, - ) - ): - return scores + scores = { + "unsafe": number_of_unsafe_results, + "suspicious": number_of_suspicious_results, + "unknown": number_of_unknown_results, + } + should_fail_fast = fail_fast and ( + number_of_unsafe_results > 0 and number_of_suspicious_results == 0 + ) + if ( + results.is_denylisted + or should_fail_fast + or is_unsafe( + number_of_safe_results, + number_of_unsafe_results, + number_of_suspicious_results, + ) + ): + return scores - scores["unsafe"] = 0 - scores["suspicious"] = 0 - return scores + scores["unsafe"] = 0 + scores["suspicious"] = 0 + return scores def security_scan( @@ -948,125 +969,149 @@ def security_scan( fail_fast: Optional[bool] = DEFAULT_FAIL_FAST, check_magic_bytes: bool = True, ) -> Dict[str, int]: - """Security scan to detect malicious content in pickle files. - - Args: - pickle_bytes: Pickle bytecode or stream to scan. - force_scan: If True, force scan even if the file is not a pickle file. - recursion_depth: Current recursion depth for nested archives. - fail_fast: Whether to fail fast on first unsafe or suspicious match. - check_magic_bytes: Whether to perform magic byte checks. - - Returns: - A dictionary containing the scores for unsafe, suspicious, and unknown - results. - """ - if recursion_depth > 10: - raise MaxRecursionDepthExceededError("Max recursion depth of 10 exceeded.") - if recursion_depth > 3: - logging.warning("Suspiciously deep recursion depth of %d", recursion_depth) - - original_pos = None - if not isinstance(pickle_bytes, bytes): - original_pos = pickle_bytes.tell() - - try: - is_archive = False + """Security scan to detect malicious content in pickle files. + + Args: + pickle_bytes: Pickle bytecode or stream to scan. + force_scan: If True, force scan even if the file is not a pickle file. + recursion_depth: Current recursion depth for nested archives. + fail_fast: Whether to fail fast on first unsafe or suspicious match. + check_magic_bytes: Whether to perform magic byte checks. + + Returns: + A dictionary containing the scores for unsafe, suspicious, and unknown + results. + """ + if recursion_depth > 10: + raise MaxRecursionDepthExceededError("Max recursion depth of 10 exceeded.") + if recursion_depth > 3: + logging.warning("Suspiciously deep recursion depth of %d", recursion_depth) + + original_pos = None if not isinstance(pickle_bytes, bytes): - # Peek first 262 bytes to identify archive streams - current_pos = pickle_bytes.tell() - header = pickle_bytes.read(262) - pickle_bytes.seek(current_pos) - if header.startswith( - (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") - ) or (len(header) >= 262 and header[257:262] == b"ustar"): - is_archive = True - - if is_archive: - # Temporarily archive streams fully to bytes - archive_bytes = pickle_bytes.read() # pyrefly: ignore[missing-attribute] - if archive_bytes.startswith(b"PK\x03\x04"): - archive_type = "zip" - elif archive_bytes.startswith(b"BZh"): - archive_type = "bz2" - elif archive_bytes.startswith(b"\xfd7zXZ\x00"): - archive_type = "lzma" - elif archive_bytes.startswith(b"\x1f\x8b"): - archive_type = "gzip" - else: - archive_type = "tar" - return _extract_and_scan_archive( - archive_bytes, - archive_type, - recursion_depth, - force_scan, - fail_fast=fail_fast, - check_magic_bytes=check_magic_bytes, - ) - - # Check for compression signatures if input was raw bytes - if isinstance(pickle_bytes, bytes): - if pickle_bytes.startswith(b"PK\x03\x04"): - return _extract_and_scan_archive( - pickle_bytes, - "zip", - recursion_depth, - force_scan, - fail_fast=fail_fast, - check_magic_bytes=check_magic_bytes, - ) - elif pickle_bytes.startswith(b"BZh"): - return _extract_and_scan_archive( - pickle_bytes, - "bz2", - recursion_depth, - force_scan, - fail_fast=fail_fast, - check_magic_bytes=check_magic_bytes, - ) - elif pickle_bytes.startswith(b"\xfd7zXZ\x00"): - return _extract_and_scan_archive( - pickle_bytes, - "lzma", - recursion_depth, - force_scan, - fail_fast=fail_fast, - check_magic_bytes=check_magic_bytes, - ) - elif pickle_bytes.startswith(b"\x1f\x8b"): - return _extract_and_scan_archive( - pickle_bytes, - "gzip", - recursion_depth, - force_scan, - fail_fast=fail_fast, - check_magic_bytes=check_magic_bytes, - ) - elif len(pickle_bytes) >= 262 and pickle_bytes[257:262] == b"ustar": - return _extract_and_scan_archive( + original_pos = pickle_bytes.tell() + + try: + is_archive = False + if not isinstance(pickle_bytes, bytes): + # Peek first 262 bytes to identify archive streams + current_pos = pickle_bytes.tell() + header = pickle_bytes.read(262) + pickle_bytes.seek(current_pos) + if header.startswith( + (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") + ) or (len(header) >= 262 and header[257:262] == b"ustar"): + is_archive = True + + if is_archive: + # Temporarily archive streams fully to bytes + archive_bytes = pickle_bytes.read() # pyrefly: ignore[missing-attribute] + if archive_bytes.startswith(b"PK\x03\x04"): + archive_type = "zip" + elif archive_bytes.startswith(b"BZh"): + archive_type = "bz2" + elif archive_bytes.startswith(b"\xfd7zXZ\x00"): + archive_type = "lzma" + elif archive_bytes.startswith(b"\x1f\x8b"): + archive_type = "gzip" + else: + archive_type = "tar" + return _extract_and_scan_archive( + archive_bytes, + archive_type, + recursion_depth, + force_scan, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + + # Check for compression signatures if input was raw bytes + if isinstance(pickle_bytes, bytes): + if pickle_bytes.startswith(b"PK\x03\x04"): + return _extract_and_scan_archive( + pickle_bytes, + "zip", + recursion_depth, + force_scan, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + elif pickle_bytes.startswith(b"BZh"): + return _extract_and_scan_archive( + pickle_bytes, + "bz2", + recursion_depth, + force_scan, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + elif pickle_bytes.startswith(b"\xfd7zXZ\x00"): + return _extract_and_scan_archive( + pickle_bytes, + "lzma", + recursion_depth, + force_scan, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + elif pickle_bytes.startswith(b"\x1f\x8b"): + return _extract_and_scan_archive( + pickle_bytes, + "gzip", + recursion_depth, + force_scan, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + elif len(pickle_bytes) >= 262 and pickle_bytes[257:262] == b"ustar": + return _extract_and_scan_archive( + pickle_bytes, + "tar", + recursion_depth, + force_scan, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + + return _security_scan_internal( pickle_bytes, - "tar", - recursion_depth, force_scan, fail_fast=fail_fast, check_magic_bytes=check_magic_bytes, ) - - return _security_scan_internal( - pickle_bytes, - force_scan, - fail_fast=fail_fast, - check_magic_bytes=check_magic_bytes, - ) - finally: - if original_pos is not None: - pickle_bytes.seek(original_pos) # pyrefly: ignore[missing-attribute] + finally: + if original_pos is not None: + pickle_bytes.seek(original_pos) # pyrefly: ignore[missing-attribute] def _merge_scores(total: Dict[str, int], new: Dict[str, int]): - total["unsafe"] += new.get("unsafe", 0) - total["suspicious"] += new.get("suspicious", 0) - total["unknown"] += new.get("unknown", 0) + total["unsafe"] += new.get("unsafe", 0) + total["suspicious"] += new.get("suspicious", 0) + total["unknown"] += new.get("unknown", 0) + + +def _is_unsafe_archive_member(name: str) -> bool: + """Detects path-traversal and absolute-path archive member names. + + This catches classic zip-slip and tar-slip payloads, including Windows-style + separators and drive letters that a naive ``..`` / leading-slash check misses. + + Args: + name: The archive member name to inspect. + + Returns: + True if the member name could escape the extraction directory. + """ + if not name: + return False + if name.startswith(("/", "\\")): + return True + # Windows drive-letter absolute paths (e.g. "C:\\...") and UNC-style paths. + if re.match(r"^[a-zA-Z]:", name): + return True + # Normalize separators and check for a traversal component. + parts = name.replace("\\", "/").split("/") + return any(part == ".." for part in parts) def _extract_and_scan_archive( @@ -1077,119 +1122,122 @@ def _extract_and_scan_archive( fail_fast: Optional[bool] = DEFAULT_FAIL_FAST, check_magic_bytes: bool = True, ) -> Dict[str, int]: - """Extracts and scans contents of an archive.""" - if not isinstance(data, bytes): - data = data.read() - - all_scores = {"unsafe": 0, "suspicious": 0, "unknown": 0} - - try: - if archive_type == "zip": - try: - with zipfile.ZipFile(io.BytesIO(data)) as zf: - for name in zf.namelist(): - if ".." in name or name.startswith("/"): - # Zip slip detection - logging.warning("Zip slip detected: %s", name) - return { - "unsafe": constants.HIGH_SEVERITY_ZIPSLIP, - "suspicious": 0, - "unknown": 0, - } # Return early - - with zf.open(name) as f: - content = f.read() - scores = security_scan( - content, - force_scan=force_scan, - recursion_depth=recursion_depth + 1, - fail_fast=fail_fast, - check_magic_bytes=check_magic_bytes, - ) - _merge_scores(all_scores, scores) - except zipfile.BadZipFile as e: - logging.warning("Error processing zip archive: %s", e) + """Extracts and scans contents of an archive.""" + if not isinstance(data, bytes): + data = data.read() + + all_scores = {"unsafe": 0, "suspicious": 0, "unknown": 0} + + try: + if archive_type == "zip": + try: + with zipfile.ZipFile(io.BytesIO(data)) as zf: + for name in zf.namelist(): + if _is_unsafe_archive_member(name): + # Zip slip detection + logging.warning("Zip slip detected: %s", name) + return { + "unsafe": constants.HIGH_SEVERITY_ZIPSLIP, + "suspicious": 0, + "unknown": 0, + } # Return early + + with zf.open(name) as f: + content = f.read() + scores = security_scan( + content, + force_scan=force_scan, + recursion_depth=recursion_depth + 1, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + _merge_scores(all_scores, scores) + except zipfile.BadZipFile as e: + logging.warning("Error processing zip archive: %s", e) + return { + "unsafe": constants.HIGH_SEVERITY_ZIPSLIP, + "suspicious": 0, + "unknown": 0, + } + + elif archive_type == "bz2": + content = utils.extract_bz2_contents(data) + scores = security_scan( + content, + force_scan=force_scan, + recursion_depth=recursion_depth + 1, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + _merge_scores(all_scores, scores) + + elif archive_type == "lzma": + content = utils.extract_lzma_contents(data) + scores = security_scan( + content, + force_scan=force_scan, + recursion_depth=recursion_depth + 1, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + _merge_scores(all_scores, scores) + + elif archive_type == "gzip": + content = utils.extract_gzip_contents(data) + scores = security_scan( + content, + force_scan=force_scan, + recursion_depth=recursion_depth + 1, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + _merge_scores(all_scores, scores) + + elif archive_type == "tar": + for name, content in utils.extract_tar_contents(data): + if _is_unsafe_archive_member(name): + logging.warning("Tar slip detected: %s", name) + return { + "unsafe": constants.HIGH_SEVERITY_ZIPSLIP, + "suspicious": 0, + "unknown": 0, + } # Return early + scores = security_scan( + content, + force_scan=force_scan, + recursion_depth=recursion_depth + 1, + fail_fast=fail_fast, + check_magic_bytes=check_magic_bytes, + ) + _merge_scores(all_scores, scores) + + else: + logging.warning("Unsupported archive type: %s", archive_type) + return { + "unsafe": constants.HIGH_SEVERITY_ARCHIVE_ERROR, + "suspicious": 0, + "unknown": 0, + } + + except MaxRecursionDepthExceededError: + raise + except ( + zipfile.BadZipFile, + tarfile.TarError, + lzma.LZMAError, + OSError, + EOFError, + ValueError, + ) as e: + logging.warning("Error processing %s archive: %s", archive_type, e) + # Block file if extraction fails to prevent security bypass return { "unsafe": constants.HIGH_SEVERITY_ZIPSLIP, "suspicious": 0, "unknown": 0, } - elif archive_type == "bz2": - content = utils.extract_bz2_contents(data) - scores = security_scan( - content, - force_scan=force_scan, - recursion_depth=recursion_depth + 1, - fail_fast=fail_fast, - ) - _merge_scores(all_scores, scores) - - elif archive_type == "lzma": - content = utils.extract_lzma_contents(data) - scores = security_scan( - content, - force_scan=force_scan, - recursion_depth=recursion_depth + 1, - fail_fast=fail_fast, - ) - _merge_scores(all_scores, scores) - - elif archive_type == "gzip": - content = utils.extract_gzip_contents(data) - scores = security_scan( - content, - force_scan=force_scan, - recursion_depth=recursion_depth + 1, - fail_fast=fail_fast, - ) - _merge_scores(all_scores, scores) - - elif archive_type == "tar": - for name, content in utils.extract_tar_contents(data): - if ".." in name or name.startswith("/"): - logging.warning("Tar slip detected: %s", name) - return { - "unsafe": constants.HIGH_SEVERITY_ZIPSLIP, - "suspicious": 0, - "unknown": 0, - } # Return early - scores = security_scan( - content, - force_scan=force_scan, - recursion_depth=recursion_depth + 1, - fail_fast=fail_fast, - check_magic_bytes=check_magic_bytes, - ) - _merge_scores(all_scores, scores) - - else: - logging.warning("Unsupported archive type: %s", archive_type) - return { - "unsafe": 0, - "suspicious": 0, - "unknown": constants.HIGH_SEVERITY_ZIPSLIP, - } - - except MaxRecursionDepthExceededError: - raise - except ( - zipfile.BadZipFile, - tarfile.TarError, - lzma.LZMAError, - OSError, - EOFError, - ValueError, - ) as e: - logging.warning("Error processing %s archive: %s", archive_type, e) - # Block file if extraction fails to prevent security bypass - return { - "unsafe": constants.HIGH_SEVERITY_ZIPSLIP, - "suspicious": 0, - "unknown": 0, - } - - return all_scores + return all_scores def _security_scan_internal( @@ -1198,162 +1246,164 @@ def _security_scan_internal( fail_fast: Optional[bool] = DEFAULT_FAIL_FAST, check_magic_bytes: bool = True, ) -> Dict[str, int]: - """Security scan to detect malicious content in pickle files. - - Args: - pickle_bytes: Pickle bytecode or stream to scan. - force_scan: If True, force scan even if the file is not a pickle file. - fail_fast: Whether to fail fast on first unsafe or suspicious match. - check_magic_bytes: Whether to perform magic byte checks. - - Returns: - A dictionary containing the scores for unsafe, suspicious, and unknown - finds. - """ - # Normalize to seekable stream and get length - if isinstance(pickle_bytes, bytes): - stream = io.BytesIO(pickle_bytes) - pickle_length = len(pickle_bytes) - else: - stream = pickle_bytes - try: - is_seekable = stream.seekable() - except (AttributeError, ValueError): - is_seekable = False - - if not is_seekable: - data = stream.read() - stream = io.BytesIO(data) - pickle_length = len(data) - else: - current_pos = stream.tell() - stream.seek(0, io.SEEK_END) - pickle_length = stream.tell() - stream.seek(current_pos) - - # Check if pickle (always stream now) - try: - is_pickle = utils.is_pickle_file( - stream, check_magic_bytes=check_magic_bytes - ) - except (OSError, AttributeError, io.UnsupportedOperation, ValueError): - is_pickle = False - - if not is_pickle and not force_scan: - return {"unsafe": 0, "suspicious": 0, "unknown": 0} - - # Find actual start of valid pickle payload if it has leading garbage - start_offset = 0 - try: - current_pos = stream.tell() - stream.seek(0) - header_bytes = stream.read(1024) - stream.seek(current_pos) - is_archive = header_bytes.startswith( - (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") - ) or (len(header_bytes) >= 262 and header_bytes[257:262] == b"ustar") - if not is_archive: - start_offset = utils.find_pickle_start_offset(stream) - else: - start_offset = 0 + """Security scan to detect malicious content in pickle files. - if start_offset > 0: - stream.seek(start_offset) - pickle_length -= start_offset - else: - stream.seek(current_pos) - except (OSError, AttributeError, io.UnsupportedOperation): - pass + Args: + pickle_bytes: Pickle bytecode or stream to scan. + force_scan: If True, force scan even if the file is not a pickle file. + fail_fast: Whether to fail fast on first unsafe or suspicious match. + check_magic_bytes: Whether to perform magic byte checks. - pickle_file_path = None - shm = None - shm_name = None + Returns: + A dictionary containing the scores for unsafe, suspicious, and unknown + finds. + """ + # Normalize to seekable stream and get length + if isinstance(pickle_bytes, bytes): + stream = io.BytesIO(pickle_bytes) + pickle_length = len(pickle_bytes) + else: + stream = pickle_bytes + try: + is_seekable = stream.seekable() + except (AttributeError, ValueError): + is_seekable = False + + if not is_seekable: + data = stream.read() + stream = io.BytesIO(data) + pickle_length = len(data) + else: + current_pos = stream.tell() + stream.seek(0, io.SEEK_END) + pickle_length = stream.tell() + stream.seek(current_pos) - if pickle_length >= constants.MIN_SIZE_FOR_CHUNKING: + # Check if pickle (always stream now) try: - shm = shared_memory.SharedMemory(create=True, size=pickle_length) - shm_name = shm.name + is_pickle = utils.is_pickle_file(stream, check_magic_bytes=check_magic_bytes) + except (OSError, AttributeError, io.UnsupportedOperation, ValueError): + is_pickle = False - # Fast path for BytesIO, chunked fallback for other streams - if isinstance(stream, io.BytesIO): - shm.buf[:pickle_length] = stream.getbuffer()[ # pyrefly: ignore[unsupported-operation] - start_offset : start_offset + pickle_length - ] - else: - offset = 0 + if not is_pickle and not force_scan: + return {"unsafe": 0, "suspicious": 0, "unknown": 0} + + # Find actual start of valid pickle payload if it has leading garbage + start_offset = 0 + try: current_pos = stream.tell() - stream.seek(start_offset) - try: - while True: - chunk = stream.read(1024 * 1024) - if not chunk: - break - shm.buf[offset : offset + len(chunk)] = chunk # pyrefly: ignore[unsupported-operation] - offset += len(chunk) - finally: - stream.seek(current_pos) - except OSError: - # Fallback to tempfile with chunked buffering - with tempfile.NamedTemporaryFile(delete=False) as temp_file: - pickle_file_path = temp_file.name - if isinstance(stream, io.BytesIO): - temp_file.write( - stream.getbuffer()[start_offset : start_offset + pickle_length] - ) + stream.seek(0) + header_bytes = stream.read(1024) + stream.seek(current_pos) + is_archive = header_bytes.startswith( + (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") + ) or (len(header_bytes) >= 262 and header_bytes[257:262] == b"ustar") + if not is_archive: + start_offset = utils.find_pickle_start_offset(stream) + else: + start_offset = 0 + + if start_offset > 0: + stream.seek(start_offset) + pickle_length -= start_offset else: - current_pos = stream.tell() - stream.seek(start_offset) - try: - while True: - chunk = stream.read(1024 * 1024) - if not chunk: - break - temp_file.write(chunk) - finally: stream.seek(current_pos) + except (OSError, AttributeError, io.UnsupportedOperation): + pass - original_stream_pos = None - try: - original_stream_pos = stream.tell() - except (OSError, AttributeError, io.UnsupportedOperation): - pass - - try: - final_scores = {"unsafe": 0, "suspicious": 0, "unknown": 0} - for scan_approach in [picklemagic_scan, genops_scan]: - scores = apply_approach( - scan_approach, - stream, - pickle_file_path, - shm_name, - fail_fast=fail_fast, - pickle_length=pickle_length, - ) - # Restore stream pointer after each scan approach to prevent EOF errors - if original_stream_pos is not None and hasattr(stream, "seek"): + pickle_file_path = None + shm = None + shm_name = None + + if pickle_length >= constants.MIN_SIZE_FOR_CHUNKING: try: - stream.seek(original_stream_pos) - except (OSError, AttributeError, io.UnsupportedOperation): - logging.debug("Failed to restore stream pointer inside scan loop.") + shm = shared_memory.SharedMemory(create=True, size=pickle_length) + shm_name = shm.name + + # Fast path for BytesIO, chunked fallback for other streams + if isinstance(stream, io.BytesIO): + shm.buf[ + :pickle_length + ] = stream.getbuffer()[ # pyrefly: ignore[unsupported-operation] + start_offset : start_offset + pickle_length + ] + else: + offset = 0 + current_pos = stream.tell() + stream.seek(start_offset) + try: + while True: + chunk = stream.read(1024 * 1024) + if not chunk: + break + shm.buf[offset : offset + len(chunk)] = ( + chunk # pyrefly: ignore[unsupported-operation] + ) + offset += len(chunk) + finally: + stream.seek(current_pos) + except OSError: + # Fallback to tempfile with chunked buffering + with tempfile.NamedTemporaryFile(delete=False) as temp_file: + pickle_file_path = temp_file.name + if isinstance(stream, io.BytesIO): + temp_file.write( + stream.getbuffer()[start_offset : start_offset + pickle_length] + ) + else: + current_pos = stream.tell() + stream.seek(start_offset) + try: + while True: + chunk = stream.read(1024 * 1024) + if not chunk: + break + temp_file.write(chunk) + finally: + stream.seek(current_pos) + + original_stream_pos = None + try: + original_stream_pos = stream.tell() + except (OSError, AttributeError, io.UnsupportedOperation): + pass - if scores["unsafe"] > 0 or scores["suspicious"] > 0: - return scores - final_scores["unknown"] += scores["unknown"] - return final_scores - finally: - # Ensure stream is seeked back before exiting so load_func gets - # a clean stream - if original_stream_pos is not None and hasattr(stream, "seek"): - try: - stream.seek(original_stream_pos) - except (OSError, AttributeError, io.UnsupportedOperation): - logging.debug("Failed to restore stream pointer before exiting scan.") - - if shm: - shm.close() - shm.unlink() - if pickle_file_path: - os.remove(pickle_file_path) + try: + final_scores = {"unsafe": 0, "suspicious": 0, "unknown": 0} + for scan_approach in [picklemagic_scan, genops_scan]: + scores = apply_approach( + scan_approach, + stream, + pickle_file_path, + shm_name, + fail_fast=fail_fast, + pickle_length=pickle_length, + ) + # Restore stream pointer after each scan approach to prevent EOF errors + if original_stream_pos is not None and hasattr(stream, "seek"): + try: + stream.seek(original_stream_pos) + except (OSError, AttributeError, io.UnsupportedOperation): + logging.debug("Failed to restore stream pointer inside scan loop.") + + if scores["unsafe"] > 0 or scores["suspicious"] > 0: + return scores + final_scores["unknown"] += scores["unknown"] + return final_scores + finally: + # Ensure stream is seeked back before exiting so load_func gets + # a clean stream + if original_stream_pos is not None and hasattr(stream, "seek"): + try: + stream.seek(original_stream_pos) + except (OSError, AttributeError, io.UnsupportedOperation): + logging.debug("Failed to restore stream pointer before exiting scan.") + + if shm: + shm.close() + shm.unlink() + if pickle_file_path: + os.remove(pickle_file_path) _ORIG_METHODS_BEFORE_HOOKING = {} @@ -1363,17 +1413,17 @@ def _security_scan_internal( def _report_or_raise( classification: utils.Classification, report_only: bool, log_info=False ): - """Reports or raises an error based on classification and report_only flag.""" + """Reports or raises an error based on classification and report_only flag.""" - if report_only: - logging_function = logging.info if log_info else logging.error - logging_function( + if report_only: + logging_function = logging.info if log_info else logging.error + logging_function( + constants.ERROR_STRING.substitute(classification=classification.value) + ) + return + raise UnsafePickleDetectedError( constants.ERROR_STRING.substitute(classification=classification.value) ) - return - raise UnsafePickleDetectedError( - constants.ERROR_STRING.substitute(classification=classification.value) - ) def _scan_and_load( @@ -1389,128 +1439,130 @@ def _scan_and_load( *args: Any, **kwargs: Any, ): - """Internal helper to scan and load pickle data.""" + """Internal helper to scan and load pickle data.""" - if is_load: - if not isinstance(pickle_file_or_bytes, io.IOBase): - raise TypeError("pickle_file_or_bytes must be IOBase when is_load=True") + if is_load: + if not isinstance(pickle_file_or_bytes, io.IOBase): + raise TypeError("pickle_file_or_bytes must be IOBase when is_load=True") - pickle_file = pickle_file_or_bytes + pickle_file = pickle_file_or_bytes - # Dynamically handle non-seekable streams - try: - is_seekable = pickle_file.seekable() - except (AttributeError, ValueError): - is_seekable = False + # Dynamically handle non-seekable streams + try: + is_seekable = pickle_file.seekable() + except (AttributeError, ValueError): + is_seekable = False - if not is_seekable: - # Fallback: read non-seekable stream into a seekable BytesIO - # We only read it fully when seek is not supported - data_bytes = pickle_file.read() - pickle_file = io.BytesIO(data_bytes) + if not is_seekable: + # Fallback: read non-seekable stream into a seekable BytesIO + # We only read it fully when seek is not supported + data_bytes = pickle_file.read() + pickle_file = io.BytesIO(data_bytes) - try: - current_pos = pickle_file.tell() - pickle_file.seek(0) - header_bytes = pickle_file.read(1024) - is_archive = header_bytes.startswith( - (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") - ) or (len(header_bytes) >= 262 and header_bytes[257:262] == b"ustar") - if not is_archive: - start_offset = utils.find_pickle_start_offset(header_bytes) - else: - start_offset = 0 - - if start_offset > 0: - pickle_file.seek(start_offset) - else: - pickle_file.seek(current_pos) - except (OSError, AttributeError, io.UnsupportedOperation): - pass - - scan_source = pickle_file - else: - if not isinstance(pickle_file_or_bytes, bytes): - raise TypeError("pickle_file_or_bytes must be bytes when is_load=False") - data_bytes = pickle_file_or_bytes - is_archive = data_bytes.startswith( - (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") - ) or (len(data_bytes) >= 262 and data_bytes[257:262] == b"ustar") - if not is_archive: - start_offset = utils.find_pickle_start_offset(data_bytes) + try: + current_pos = pickle_file.tell() + pickle_file.seek(0) + header_bytes = pickle_file.read(1024) + is_archive = header_bytes.startswith( + (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") + ) or (len(header_bytes) >= 262 and header_bytes[257:262] == b"ustar") + if not is_archive: + start_offset = utils.find_pickle_start_offset(header_bytes) + else: + start_offset = 0 + + if start_offset > 0: + pickle_file.seek(start_offset) + else: + pickle_file.seek(current_pos) + except (OSError, AttributeError, io.UnsupportedOperation): + pass + + scan_source = pickle_file else: - start_offset = 0 + if not isinstance(pickle_file_or_bytes, bytes): + raise TypeError("pickle_file_or_bytes must be bytes when is_load=False") + data_bytes = pickle_file_or_bytes + is_archive = data_bytes.startswith( + (b"PK\x03\x04", b"BZh", b"\xfd7zXZ\x00", b"\x1f\x8b") + ) or (len(data_bytes) >= 262 and data_bytes[257:262] == b"ustar") + if not is_archive: + start_offset = utils.find_pickle_start_offset(data_bytes) + else: + start_offset = 0 - if start_offset > 0: - data_bytes = data_bytes[start_offset:] - scan_source = data_bytes - pickle_file = None + if start_offset > 0: + data_bytes = data_bytes[start_offset:] + scan_source = data_bytes + pickle_file = None - loader_mod = utils.get_copied_module(hooked_mod_name or "_pickle") + loader_mod = utils.get_copied_module(hooked_mod_name or "_pickle") - if is_load: - load_func = loader_mod.load # pyrefly: ignore[missing-attribute] - load_args = (pickle_file,) - else: - load_func = loader_mod.loads # pyrefly: ignore[missing-attribute] - load_args = (data_bytes,) # pyrefly: ignore[unbound-name] + if is_load: + load_func = loader_mod.load # pyrefly: ignore[missing-attribute] + load_args = (pickle_file,) + else: + load_func = loader_mod.loads # pyrefly: ignore[missing-attribute] + load_args = (data_bytes,) # pyrefly: ignore[unbound-name] - if strict_check and allow_unsafe: - error_string_illegal_combination = ( - "Strict scanning and allow_unsafe cannot be used together." - ) - if report_only: - logging.error(error_string_illegal_combination) - return - raise IllegalArgumentCombinationError(error_string_illegal_combination) - elif allow_unsafe: - if report_only: - logging.info("Loading pickle file with allow_unsafe set to True.") - elif strict_check: - if strict_security_scan(scan_source): # pyrefly: ignore[bad-argument-type] - error_string_strict_check = "Pickle file failed strict security check." - if report_only: - logging.error(error_string_strict_check) - return - raise StrictCheckError(error_string_strict_check) - else: - # Default scanning routines - scan_scores = security_scan( - scan_source, force_scan=force_scan, check_magic_bytes=check_magic_bytes # pyrefly: ignore[bad-argument-type] - ) - number_of_unsafe_results = scan_scores["unsafe"] - number_of_suspicious_results = scan_scores["suspicious"] - number_of_unknown_results = scan_scores["unknown"] - - if number_of_suspicious_results == 0 and number_of_unsafe_results == 0: - if report_only: - logging.info("Loading safe pickle file") - if number_of_unknown_results > 0: - logging.warning( - "SaferPickle: File contains %d unknown items that were ignored.", - number_of_unknown_results, - ) - elif number_of_unsafe_results > number_of_suspicious_results: - _report_or_raise(utils.Classification.UNSAFE, report_only, log_info) + if strict_check and allow_unsafe: + error_string_illegal_combination = ( + "Strict scanning and allow_unsafe cannot be used together." + ) + if report_only: + logging.error(error_string_illegal_combination) + return + raise IllegalArgumentCombinationError(error_string_illegal_combination) + elif allow_unsafe: + if report_only: + logging.info("Loading pickle file with allow_unsafe set to True.") + elif strict_check: + if strict_security_scan(scan_source): # pyrefly: ignore[bad-argument-type] + error_string_strict_check = "Pickle file failed strict security check." + if report_only: + logging.error(error_string_strict_check) + return + raise StrictCheckError(error_string_strict_check) else: - _report_or_raise(utils.Classification.SUSPICIOUS, report_only, log_info) - - # Load the pickle if report_only is True and no exceptions were raised earlier - try: - return load_func(*load_args, *args, **kwargs) - except ( - AttributeError, - pickle.UnpicklingError, - ModuleNotFoundError, - EOFError, - ImportError, - ) as exc: - logging.debug( - "Safe pickle failed to load due to environmental constraints: %s", - exc, - exc_info=True, - ) - raise + # Default scanning routines + scan_scores = security_scan( + scan_source, + force_scan=force_scan, + check_magic_bytes=check_magic_bytes, # pyrefly: ignore[bad-argument-type] + ) + number_of_unsafe_results = scan_scores["unsafe"] + number_of_suspicious_results = scan_scores["suspicious"] + number_of_unknown_results = scan_scores["unknown"] + + if number_of_suspicious_results == 0 and number_of_unsafe_results == 0: + if report_only: + logging.info("Loading safe pickle file") + if number_of_unknown_results > 0: + logging.warning( + "SaferPickle: File contains %d unknown items that were ignored.", + number_of_unknown_results, + ) + elif number_of_unsafe_results > number_of_suspicious_results: + _report_or_raise(utils.Classification.UNSAFE, report_only, log_info) + else: + _report_or_raise(utils.Classification.SUSPICIOUS, report_only, log_info) + + # Load the pickle if report_only is True and no exceptions were raised earlier + try: + return load_func(*load_args, *args, **kwargs) + except ( + AttributeError, + pickle.UnpicklingError, + ModuleNotFoundError, + EOFError, + ImportError, + ) as exc: + logging.debug( + "Safe pickle failed to load due to environmental constraints: %s", + exc, + exc_info=True, + ) + raise def hook_pickle( @@ -1518,77 +1570,222 @@ def hook_pickle( log_info: bool = False, config_path: Optional[str] = None, ) -> None: - """This implements the hooking of pickle-like libraries.""" - config.set_config_path(config_path) - - def custom_loads( - pickle_bytes: bytes, - allow_unsafe: bool = False, - strict_check: bool = False, - report_only: bool = False, - force_scan: bool = False, - hooked_mod_name: str = "", - check_magic_bytes: bool = True, - *args: Any, - **kwargs: Any, - ) -> Any: - """Custom loads function for pickle to security scan before loading pickle files. + """This implements the hooking of pickle-like libraries.""" + config.set_config_path(config_path) + + def custom_loads( + pickle_bytes: bytes, + allow_unsafe: bool = False, + strict_check: bool = False, + report_only: bool = False, + force_scan: bool = False, + hooked_mod_name: str = "", + check_magic_bytes: bool = True, + *args: Any, + **kwargs: Any, + ) -> Any: + """Custom loads function for pickle to security scan before loading pickle files. + + Args: + pickle_bytes: The pickle file bytes to load. + allow_unsafe: If True, allow unsafe pickle files to be loaded. + strict_check: If True, perform a strict security check on the pickle file. + report_only: If True, only report errors and do not raise them. + force_scan: If True, force scan even if the file is not a pickle file. + hooked_mod_name: The name of the hooked module that called this function. + check_magic_bytes: Whether to perform magic byte checks to fast path + reject files. + *args: Additional arguments to pass to pickle.loads. + **kwargs: Additional keyword arguments to pass to pickle.loads. + + Returns: + None if we are in report_only mode and the pickle file is unsafe. + Result of loader_mod.loads if pickle file is safe. + + Raises: + IllegalArgumentCombinationError: If both allow_unsafe and strict_check are + set to True. + StrictCheckError: If the pickle file fails the strict security check. + UnsafePickleDetectedError: If the pickle file is unsafe. + + Logs: + If report_only is True, logs the above raised exceptions and unknown + results. + Logs if an absent class is encountered. We return even if benign. + """ + if force_report_only: + report_only = True + return _scan_and_load( + pickle_bytes, + allow_unsafe, + strict_check, + report_only, + force_scan, + hooked_mod_name, + False, + log_info, + check_magic_bytes, + *args, + **kwargs, + ) - Args: - pickle_bytes: The pickle file bytes to load. - allow_unsafe: If True, allow unsafe pickle files to be loaded. - strict_check: If True, perform a strict security check on the pickle file. - report_only: If True, only report errors and do not raise them. - force_scan: If True, force scan even if the file is not a pickle file. - hooked_mod_name: The name of the hooked module that called this function. - check_magic_bytes: Whether to perform magic byte checks to fast path - reject files. - *args: Additional arguments to pass to pickle.loads. - **kwargs: Additional keyword arguments to pass to pickle.loads. + def custom_load( + pickle_file: Any, + allow_unsafe: bool = False, + strict_check: bool = False, + report_only: bool = False, + force_scan: bool = False, + hooked_mod_name: str = "", + check_magic_bytes: bool = True, + *args: Any, + **kwargs: Any, + ) -> Any: + """Custom load function for pickle to security scan before loading pickle files. + + Args: + pickle_file: The pickle file to load. + allow_unsafe: If True, allow unsafe pickle files to be loaded. + strict_check: If True, perform a strict security check on the pickle file. + report_only: If True, only report errors and do not raise them. + force_scan: If True, force scan even if the file is not a pickle file. + hooked_mod_name: The name of the hooked module that called this function. + check_magic_bytes: Whether to perform magic byte checks to fast path + reject files. + *args: Additional arguments to pass to pickle.load. + **kwargs: Additional keyword arguments to pass to pickle.load. + + Returns: + None if we are in report_only mode and the pickle file is unsafe. + result of loader_mod.load if pickle file is safe. + + Raises: + IllegalArgumentCombinationError: If both allow_unsafe and strict_check are + set to True. + StrictCheckError: If the pickle file fails the strict security check. + UnsafePickleDetectedError: If the pickle file is unsafe. + + Logs: + If report_only is True, logs the above raised exceptions and unknown + results. + Logs if an absent class is encountered. We return even if benign. + """ + if force_report_only: + report_only = True + return _scan_and_load( + pickle_file, + allow_unsafe, + strict_check, + report_only, + force_scan, + hooked_mod_name, + True, + log_info, + check_magic_bytes, + *args, + **kwargs, + ) - Returns: - None if we are in report_only mode and the pickle file is unsafe. - Result of loader_mod.loads if pickle file is safe. + # The main hooking routine + hookable_mods: Set[str] = set( + [ + "_pickle", + "joblib", + "cloudpickle", + "torch", + "pickle", + "dill", + ] + ) - Raises: - IllegalArgumentCombinationError: If both allow_unsafe and strict_check are - set to True. - StrictCheckError: If the pickle file fails the strict security check. - UnsafePickleDetectedError: If the pickle file is unsafe. + for hookable_mod in hookable_mods: + if sys.modules.get(hookable_mod): + module = sys.modules[hookable_mod] + else: + logging.debug("%s DOES NOT exist in sys.modules", hookable_mod) + logging.debug("Importing %s now", hookable_mod) + try: + # Imports are necessary for hooking to work + module = importlib.import_module(hookable_mod) + except (ImportError, ModuleNotFoundError): + logging.debug("Failed to import %s", hookable_mod) + continue + + # Force copy before patching to ensure we copy the unhooked version + _ = utils.get_copied_module(hookable_mod) + + with _HOOKING_LOCK: + if hookable_mod not in _ORIG_METHODS_BEFORE_HOOKING: + _ORIG_METHODS_BEFORE_HOOKING[hookable_mod] = {} + + methods_to_patch = { + "load": functools.partial(custom_load, hooked_mod_name=hookable_mod), + "_load": functools.partial(custom_load, hooked_mod_name=hookable_mod), + "loads": functools.partial(custom_loads, hooked_mod_name=hookable_mod), + "_loads": functools.partial(custom_loads, hooked_mod_name=hookable_mod), + } + for method_name, custom_func in methods_to_patch.items(): + if hasattr(module, method_name): + if method_name not in _ORIG_METHODS_BEFORE_HOOKING[hookable_mod]: + _ORIG_METHODS_BEFORE_HOOKING[hookable_mod][method_name] = ( + getattr(module, method_name) + ) + setattr(module, method_name, custom_func) - Logs: - If report_only is True, logs the above raised exceptions and unknown - results. - Logs if an absent class is encountered. We return even if benign. + +@contextlib.contextmanager +def hook_pickle_libs( + report_only: bool = True, + log_info: bool = False, + config_path: Optional[str] = None, +) -> Iterator[None]: + """Context manager that hooks pickle on entry and unhooks on exit. + + Args: + report_only: If True, hooks will only log errors instead of raising them. + log_info: If True, use logging.info instead of logging.error for + reporting. + config_path: Optional path to a JSON config file for the allow-list. """ - if force_report_only: - report_only = True - return _scan_and_load( - pickle_bytes, - allow_unsafe, - strict_check, - report_only, - force_scan, - hooked_mod_name, - False, - log_info, - check_magic_bytes, - *args, - **kwargs, + hook_pickle( + force_report_only=report_only, log_info=log_info, config_path=config_path ) + try: + yield + finally: + unhook_pickle() + + +def unhook_pickle() -> None: + """Unhooks the pickle-like libraries.""" + with _HOOKING_LOCK: + for module_name, methods in _ORIG_METHODS_BEFORE_HOOKING.items(): + try: + module = importlib.import_module(module_name) + for method_name, original_method in methods.items(): + if hasattr(module, method_name): + setattr(module, method_name, original_method) + except (ImportError, ModuleNotFoundError): + logging.debug("Failed to import %s for unhooking", module_name) + continue + # Empty stored methods to avoid re-unhooking on a second unhook call + _ORIG_METHODS_BEFORE_HOOKING.clear() + + +def load( + pickle_file: Any, + allow_unsafe: bool = False, + strict_check: bool = False, + report_only: bool = False, + force_scan: bool = False, + log_info: bool = False, + check_magic_bytes: bool = True, + *args: Any, + **kwargs: Any, +) -> Any: + """Custom load function to security scan before loading pickle files. - def custom_load( - pickle_file: Any, - allow_unsafe: bool = False, - strict_check: bool = False, - report_only: bool = False, - force_scan: bool = False, - hooked_mod_name: str = "", - check_magic_bytes: bool = True, - *args: Any, - **kwargs: Any, - ) -> Any: - """Custom load function for pickle to security scan before loading pickle files. + This function can be used as a replacement for pickle.load or torch.load, + providing security scan features. Args: pickle_file: The pickle file to load. @@ -1596,36 +1793,26 @@ def custom_load( strict_check: If True, perform a strict security check on the pickle file. report_only: If True, only report errors and do not raise them. force_scan: If True, force scan even if the file is not a pickle file. - hooked_mod_name: The name of the hooked module that called this function. - check_magic_bytes: Whether to perform magic byte checks to fast path - reject files. - *args: Additional arguments to pass to pickle.load. - **kwargs: Additional keyword arguments to pass to pickle.load. + log_info: If True, use logging.info instead of logging.error for reporting. + check_magic_bytes: Whether to perform magic byte checks to fast path reject + files. + *args: Additional arguments to pass to torch.load. + **kwargs: Additional keyword arguments to pass to torch.load. Returns: - None if we are in report_only mode and the pickle file is unsafe. - result of loader_mod.load if pickle file is safe. + The unpickled object or None if the pickle file is unsafe and report_only is + True. Raises: - IllegalArgumentCombinationError: If both allow_unsafe and strict_check are - set to True. - StrictCheckError: If the pickle file fails the strict security check. UnsafePickleDetectedError: If the pickle file is unsafe. - - Logs: - If report_only is True, logs the above raised exceptions and unknown - results. - Logs if an absent class is encountered. We return even if benign. """ - if force_report_only: - report_only = True return _scan_and_load( pickle_file, allow_unsafe, strict_check, report_only, force_scan, - hooked_mod_name, + "torch", True, log_info, check_magic_bytes, @@ -1633,96 +1820,9 @@ def custom_load( **kwargs, ) - # The main hooking routine - hookable_mods: Set[str] = set([ - "_pickle", - "joblib", - "cloudpickle", - "torch", - "pickle", - "dill", - ]) - - for hookable_mod in hookable_mods: - if sys.modules.get(hookable_mod): - module = sys.modules[hookable_mod] - else: - logging.debug("%s DOES NOT exist in sys.modules", hookable_mod) - logging.debug("Importing %s now", hookable_mod) - try: - # Imports are necessary for hooking to work - module = importlib.import_module(hookable_mod) - except (ImportError, ModuleNotFoundError): - logging.debug("Failed to import %s", hookable_mod) - continue - - # Force copy before patching to ensure we copy the unhooked version - _ = utils.get_copied_module(hookable_mod) - - with _HOOKING_LOCK: - if hookable_mod not in _ORIG_METHODS_BEFORE_HOOKING: - _ORIG_METHODS_BEFORE_HOOKING[hookable_mod] = {} - - methods_to_patch = { - "load": functools.partial(custom_load, hooked_mod_name=hookable_mod), - "_load": functools.partial(custom_load, hooked_mod_name=hookable_mod), - "loads": functools.partial( - custom_loads, hooked_mod_name=hookable_mod - ), - "_loads": functools.partial( - custom_loads, hooked_mod_name=hookable_mod - ), - } - for method_name, custom_func in methods_to_patch.items(): - if hasattr(module, method_name): - if method_name not in _ORIG_METHODS_BEFORE_HOOKING[hookable_mod]: - _ORIG_METHODS_BEFORE_HOOKING[hookable_mod][method_name] = getattr( - module, method_name - ) - setattr(module, method_name, custom_func) - - -@contextlib.contextmanager -def hook_pickle_libs( - report_only: bool = True, - log_info: bool = False, - config_path: Optional[str] = None, -) -> Iterator[None]: - """Context manager that hooks pickle on entry and unhooks on exit. - - Args: - report_only: If True, hooks will only log errors instead of raising them. - log_info: If True, use logging.info instead of logging.error for - reporting. - config_path: Optional path to a JSON config file for the allow-list. - """ - hook_pickle( - force_report_only=report_only, log_info=log_info, config_path=config_path - ) - try: - yield - finally: - unhook_pickle() - -def unhook_pickle() -> None: - """Unhooks the pickle-like libraries.""" - with _HOOKING_LOCK: - for module_name, methods in _ORIG_METHODS_BEFORE_HOOKING.items(): - try: - module = importlib.import_module(module_name) - for method_name, original_method in methods.items(): - if hasattr(module, method_name): - setattr(module, method_name, original_method) - except (ImportError, ModuleNotFoundError): - logging.debug("Failed to import %s for unhooking", module_name) - continue - # Empty stored methods to avoid re-unhooking on a second unhook call - _ORIG_METHODS_BEFORE_HOOKING.clear() - - -def load( - pickle_file: Any, +def loads( + pickle_data: bytes, allow_unsafe: bool = False, strict_check: bool = False, report_only: bool = False, @@ -1732,88 +1832,88 @@ def load( *args: Any, **kwargs: Any, ) -> Any: - """Custom load function to security scan before loading pickle files. - - This function can be used as a replacement for pickle.load or torch.load, - providing security scan features. - - Args: - pickle_file: The pickle file to load. - allow_unsafe: If True, allow unsafe pickle files to be loaded. - strict_check: If True, perform a strict security check on the pickle file. - report_only: If True, only report errors and do not raise them. - force_scan: If True, force scan even if the file is not a pickle file. - log_info: If True, use logging.info instead of logging.error for reporting. - check_magic_bytes: Whether to perform magic byte checks to fast path reject - files. - *args: Additional arguments to pass to torch.load. - **kwargs: Additional keyword arguments to pass to torch.load. - - Returns: - The unpickled object or None if the pickle file is unsafe and report_only is - True. - - Raises: - UnsafePickleDetectedError: If the pickle file is unsafe. - """ - return _scan_and_load( - pickle_file, - allow_unsafe, - strict_check, - report_only, - force_scan, - "torch", - True, - log_info, - check_magic_bytes, - *args, - **kwargs, - ) + """Custom loads function to security scan before loading pickle data. + This function can be used as a replacement for pickle.loads, providing + security scan features on serialized pickle bytes. -class Unpickler(pickle.Unpickler): - """Custom unpickler class to security scan before unpickling.""" - - def __init__( - self, - file: Any, - allow_unsafe: bool = False, - strict_check: bool = False, - report_only: bool = False, - force_scan: bool = False, - log_info: bool = False, - check_magic_bytes: bool = True, - *args: Any, - **kwargs: Any, - ): - super().__init__(file, *args, **kwargs) - self.file = file - self.args = args - self.kwargs = kwargs - self._allow_unsafe = allow_unsafe - self._strict_check = strict_check - self._report_only = report_only - self._force_scan = force_scan - self._log_info = log_info - self._check_magic_bytes = check_magic_bytes - - def load(self) -> Any: - """Security scan before loading pickle files.""" + Args: + pickle_data: The serialized pickle data to load. + allow_unsafe: If True, allow unsafe pickle files to be loaded. + strict_check: If True, perform a strict security check on the pickle file. + report_only: If True, only report errors and do not raise them. + force_scan: If True, force scan even if the file is not a pickle file. + log_info: If True, use logging.info instead of logging.error for reporting. + check_magic_bytes: Whether to perform magic byte checks to fast path reject + files. + *args: Additional positional arguments to pass to pickle.loads. + **kwargs: Additional keyword arguments to pass to pickle.loads. + + Returns: + The unpickled object or None if the pickle data is unsafe and report_only + is True. + + Raises: + UnsafePickleDetectedError: If the pickle data is unsafe. + """ return _scan_and_load( - self.file, - self._allow_unsafe, - self._strict_check, - self._report_only, - self._force_scan, + pickle_data, + allow_unsafe, + strict_check, + report_only, + force_scan, "pickle", - True, - self._log_info, - self._check_magic_bytes, - *self.args, - **self.kwargs, + False, + log_info, + check_magic_bytes, + *args, + **kwargs, ) +class Unpickler(pickle.Unpickler): + """Custom unpickler class to security scan before unpickling.""" + + def __init__( + self, + file: Any, + allow_unsafe: bool = False, + strict_check: bool = False, + report_only: bool = False, + force_scan: bool = False, + log_info: bool = False, + check_magic_bytes: bool = True, + *args: Any, + **kwargs: Any, + ): + super().__init__(file, *args, **kwargs) + self.file = file + self.args = args + self.kwargs = kwargs + self._allow_unsafe = allow_unsafe + self._strict_check = strict_check + self._report_only = report_only + self._force_scan = force_scan + self._log_info = log_info + self._check_magic_bytes = check_magic_bytes + + def load(self) -> Any: + """Security scan before loading pickle files.""" + return _scan_and_load( + self.file, + self._allow_unsafe, + self._strict_check, + self._report_only, + self._force_scan, + "pickle", + True, + self._log_info, + self._check_magic_bytes, + *self.args, + **self.kwargs, + ) + + if __name__ == "__main__": - if IS_COLAB_ENABLED: - hook_pickle() + if IS_COLAB_ENABLED: + hook_pickle() diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..efa3d28 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,36 @@ +"""Shared fixtures for the SaferPickle test suite.""" + +import os +import pickle + +import pytest + + +class _SystemExploit: + def __reduce__(self): + return (os.system, ("echo saferpickle_test",)) + + +class _EvalExploit: + def __reduce__(self): + return (eval, ("40 + 2",)) + + +@pytest.fixture(scope="session") +def benign_bytes(): + return pickle.dumps({"a": [1, 2, 3], "b": "hello"}) + + +@pytest.fixture(scope="session") +def benign_obj(): + return {"a": [1, 2, 3], "b": "hello"} + + +@pytest.fixture(scope="session") +def malicious_system_bytes(): + return pickle.dumps(_SystemExploit()) + + +@pytest.fixture(scope="session") +def malicious_eval_bytes(): + return pickle.dumps(_EvalExploit()) diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..cf365ba --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,25 @@ +"""Tests for the SaferPickle CLI helpers.""" + +import io +import zipfile + +import cli + + +def test_benign_pickle_is_benign(benign_bytes): + result = cli.security_scan_with_justifications(benign_bytes) + assert result["classification"] == "benign" + + +def test_malicious_pickle_is_unsafe(malicious_system_bytes): + result = cli.security_scan_with_justifications(malicious_system_bytes) + assert result["classification"] == "unsafe" + + +def test_zip_scans_all_members_not_just_first(benign_bytes, malicious_system_bytes): + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr("a.pkl", benign_bytes) + zf.writestr("b.pkl", malicious_system_bytes) + result = cli.security_scan_with_justifications(buf.getvalue()) + assert result["classification"] == "unsafe" diff --git a/tests/test_load.py b/tests/test_load.py new file mode 100644 index 0000000..1a64832 --- /dev/null +++ b/tests/test_load.py @@ -0,0 +1,67 @@ +"""Tests for load(), loads(), Unpickler, and pickle hooking.""" + +import io +import pickle + +import pytest + +import saferpickle + + +def test_loads_roundtrip(benign_bytes, benign_obj): + assert saferpickle.loads(benign_bytes) == benign_obj + + +def test_load_roundtrip(benign_bytes, benign_obj): + assert saferpickle.load(io.BytesIO(benign_bytes)) == benign_obj + + +def test_loads_raises_on_malicious(malicious_system_bytes): + with pytest.raises(saferpickle.UnsafePickleDetectedError): + saferpickle.loads(malicious_system_bytes) + + +def test_load_raises_on_malicious(malicious_system_bytes): + with pytest.raises(saferpickle.UnsafePickleDetectedError): + saferpickle.load(io.BytesIO(malicious_system_bytes)) + + +def test_loads_allow_unsafe(malicious_eval_bytes): + assert saferpickle.loads(malicious_eval_bytes, allow_unsafe=True) == 42 + + +def test_loads_strict_check(malicious_system_bytes): + with pytest.raises(saferpickle.StrictCheckError): + saferpickle.loads(malicious_system_bytes, strict_check=True) + + +def test_loads_illegal_combo(benign_bytes): + with pytest.raises(saferpickle.IllegalArgumentCombinationError): + saferpickle.loads(benign_bytes, allow_unsafe=True, strict_check=True) + + +def test_loads_report_only_returns_object(benign_bytes, benign_obj): + assert saferpickle.loads(benign_bytes, report_only=True) == benign_obj + + +def test_unpickler(benign_bytes, benign_obj, malicious_system_bytes): + assert saferpickle.Unpickler(io.BytesIO(benign_bytes)).load() == benign_obj + with pytest.raises(saferpickle.UnsafePickleDetectedError): + saferpickle.Unpickler(io.BytesIO(malicious_system_bytes)).load() + + +def test_hook_and_unhook(benign_bytes, benign_obj, malicious_system_bytes): + saferpickle.hook_pickle() + try: + assert pickle.loads(benign_bytes) == benign_obj + with pytest.raises(saferpickle.UnsafePickleDetectedError): + pickle.loads(malicious_system_bytes) + finally: + saferpickle.unhook_pickle() + + +def test_unhook_restores_pickle(benign_bytes, benign_obj, malicious_system_bytes): + saferpickle.hook_pickle() + saferpickle.unhook_pickle() + # After unhooking, plain pickle no longer raises. + assert pickle.loads(benign_bytes) == benign_obj diff --git a/tests/test_saferpickle.py b/tests/test_saferpickle.py new file mode 100644 index 0000000..c0cd7fa --- /dev/null +++ b/tests/test_saferpickle.py @@ -0,0 +1,58 @@ +"""Tests for classification and archive scanning in saferpickle.""" + +import bz2 +import gzip +import io +import lzma +import zipfile + +import saferpickle + + +def test_benign_is_clean(benign_bytes): + result = saferpickle.security_scan(benign_bytes) + assert result["unsafe"] == 0 + + +def test_malicious_system_is_unsafe(malicious_system_bytes): + assert saferpickle.security_scan(malicious_system_bytes)["unsafe"] > 0 + + +def test_malicious_eval_is_unsafe(malicious_eval_bytes): + assert saferpickle.security_scan(malicious_eval_bytes)["unsafe"] > 0 + + +def test_strict_security_scan(benign_bytes, malicious_system_bytes): + assert saferpickle.strict_security_scan(benign_bytes) is False + assert saferpickle.strict_security_scan(malicious_system_bytes) is True + + +def test_zip_slip_detection(): + assert saferpickle._is_unsafe_archive_member("../evil") + assert saferpickle._is_unsafe_archive_member("..\\evil") + assert saferpickle._is_unsafe_archive_member("/absolute/path") + assert saferpickle._is_unsafe_archive_member("C:\\evil") + assert saferpickle._is_unsafe_archive_member("\\\\unc\\path") + assert not saferpickle._is_unsafe_archive_member("normal.txt") + assert not saferpickle._is_unsafe_archive_member("foo..bar") + + +def test_compressed_archives_scan_inside(malicious_system_bytes): + for compress in (gzip.compress, bz2.compress, lzma.compress): + result = saferpickle.security_scan(compress(malicious_system_bytes)) + assert result["unsafe"] > 0 + + +def test_zip_archive_detects_malicious_member(benign_bytes, malicious_system_bytes): + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr("safe.pkl", benign_bytes) + zf.writestr("evil.pkl", malicious_system_bytes) + assert saferpickle.security_scan(buf.getvalue())["unsafe"] > 0 + + +def test_zip_slip_member_detected(benign_bytes): + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr("../evil.pkl", benign_bytes) + assert saferpickle.security_scan(buf.getvalue())["unsafe"] > 0 diff --git a/third_party/__init__.py b/third_party/__init__.py new file mode 100644 index 0000000..428bdb6 --- /dev/null +++ b/third_party/__init__.py @@ -0,0 +1 @@ +"""Vendored third-party code used by saferpickle.""" diff --git a/uv.lock b/uv.lock index 62a53d7..d6e38c5 100644 --- a/uv.lock +++ b/uv.lock @@ -11,6 +11,54 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl", hash = "sha256:0f17b89f2a4eaaedc4f28c622998aa690564b3012a396a4ffad0821007fe03ba", size = 137410, upload-time = "2026-07-03T10:57:46.735Z" }, ] +[[package]] +name = "black" +version = "26.5.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "mypy-extensions" }, + { name = "packaging" }, + { name = "pathspec" }, + { name = "platformdirs" }, + { name = "pytokens" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/37/5628dd55bf2b34257fc7603f0fe97c40e3aaf24265f416a9c85c95ca1436/black-26.5.1.tar.gz", hash = "sha256:dd321f668053961824bcc1be1cc1df748b2d7e4fa28086b08331e577b0100a73", size = 679439, upload-time = "2026-05-18T16:53:36.107Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/5c/c384363980e11e25ca6b93205949bb331fbf35f4e0dbec376dfa6326cec8/black-26.5.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:2b36cf2ddf5566e205f6535f782a62194a184d33e175b64ae8c40b1737522be3", size = 2009020, upload-time = "2026-05-18T17:05:28.132Z" }, + { url = "https://files.pythonhosted.org/packages/0b/df/9f31c5e0babbfed77d505fc5d120beb98b21b33feaeded3924ea941fe360/black-26.5.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1f7ea64ebfa01b50f693508fc39f875e264446d3b097088f84f203b9d09618a0", size = 1813335, upload-time = "2026-05-18T17:05:31.266Z" }, + { url = "https://files.pythonhosted.org/packages/fb/24/8e7b9a2fa61b0afd82209efe937557d180a1fa055bd7f6161eb9defc3719/black-26.5.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ecb3e624844c798144e9bd986954e0adc81d8911a1f30f375e1252fe26e8c294", size = 1881614, upload-time = "2026-05-18T17:05:32.718Z" }, + { url = "https://files.pythonhosted.org/packages/49/ad/b4e0d9365ba8ac34f6bbab62a4b1b2dd5d618fac3fa1b8db968c844201b5/black-26.5.1-cp313-cp313-win_amd64.whl", hash = "sha256:e1a26503279b6b310669fb0b219c39e4820b77e8189fe80f522bb511f247db0a", size = 1488925, upload-time = "2026-05-18T17:05:34.259Z" }, + { url = "https://files.pythonhosted.org/packages/a1/4b/652b859bf5df88a751c30451b09338f7fd26a77d1271c666992f836b7711/black-26.5.1-cp313-cp313-win_arm64.whl", hash = "sha256:5c34b25da232ead53a6f335b76dbea124f4d152ad568b9080d6f944bc2b34b52", size = 1289883, upload-time = "2026-05-18T17:05:36.019Z" }, + { url = "https://files.pythonhosted.org/packages/a6/16/a8da8eb208c51c7f4ce74609a45d0dcc6d8a2141e45e81ee5289d1bb0d59/black-26.5.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:e88976690a64b0af98312ca958415849cb42423423c5f2ee74af4b49a97a2168", size = 2004800, upload-time = "2026-05-18T17:05:38.182Z" }, + { url = "https://files.pythonhosted.org/packages/11/8a/a479296a19e383b70a725882a6cf3d786540601ff03cabbaaf1cce864c5a/black-26.5.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:32d5ea7f6c8bdfa6e648326ebca1f02b0764e2a029edc6f8dce2627e19d468c3", size = 1815576, upload-time = "2026-05-18T17:05:40.309Z" }, + { url = "https://files.pythonhosted.org/packages/81/6b/cfaf3d39f25132c156a068f6b805576c9103a84086019507c70e1911ee7d/black-26.5.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ea8d16dc41655aa113cd64665e7219446cd7e4ff2248d7178eaa905190c86b18", size = 1877927, upload-time = "2026-05-18T17:05:42.463Z" }, + { url = "https://files.pythonhosted.org/packages/66/76/302e313964bcff7e28df329d39f84f5270095730d85ff0acc260610a0d82/black-26.5.1-cp314-cp314-win_amd64.whl", hash = "sha256:577f21094ea469ef92ec1adaf2c9441a226d2144d01a5be2fa823cecf6543e50", size = 1511860, upload-time = "2026-05-18T17:05:43.943Z" }, + { url = "https://files.pythonhosted.org/packages/27/4e/a3827e35e0e567f9f9ee59e2a0ab979267dca98718f25547ca8c6733afd4/black-26.5.1-cp314-cp314-win_arm64.whl", hash = "sha256:ed1a20af114c301a0269bf01163d51dbef72737fd65f850001e7cbe7f3c7abae", size = 1316632, upload-time = "2026-05-18T17:05:45.521Z" }, + { url = "https://files.pythonhosted.org/packages/94/51/f975cae76d44274cc2868dc9040ac5d58d464784610234455b4e7b19c6ef/black-26.5.1-py3-none-any.whl", hash = "sha256:4ed7f7da04046d2e488437170797d3b4a4ad83906683bcb7dfc68b673bbce5e2", size = 213693, upload-time = "2026-05-18T16:53:33.964Z" }, +] + +[[package]] +name = "click" +version = "8.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/76/d4/81420972a676e8ffea40450d8c8c92943e7218a78fe9b64359836cc9876b/click-8.4.2.tar.gz", hash = "sha256:9a6cea6e60b17ebe0a44c5cc636d94f09bd66142c1cd7d8b4cd731c4917a15f6", size = 338000, upload-time = "2026-06-24T17:45:15.148Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + [[package]] name = "immutabledict" version = "4.3.1" @@ -20,17 +68,145 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a3/ce/f9018bf69ae91b273b6391a095e7c93fa5e1617f25b6ba81ad4b20c9df10/immutabledict-4.3.1-py3-none-any.whl", hash = "sha256:c9facdc0ff30fdb8e35bd16532026cac472a549e182c94fa201b51b25e4bf7bf", size = 5000, upload-time = "2026-02-15T10:32:33.672Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "isort" +version = "8.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ef/7c/ec4ab396d31b3b395e2e999c8f46dec78c5e29209fac49d1f4dace04041d/isort-8.0.1.tar.gz", hash = "sha256:171ac4ff559cdc060bcfff550bc8404a486fee0caab245679c2abe7cb253c78d", size = 769592, upload-time = "2026-02-28T10:08:20.685Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3e/95/c7c34aa53c16353c56d0b802fba48d5f5caa2cdee7958acbcb795c830416/isort-8.0.1-py3-none-any.whl", hash = "sha256:28b89bc70f751b559aeca209e6120393d43fbe2490de0559662be7a9787e3d75", size = 89733, upload-time = "2026-02-28T10:08:19.466Z" }, +] + +[[package]] +name = "mypy-extensions" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, +] + +[[package]] +name = "packaging" +version = "26.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, +] + +[[package]] +name = "pathspec" +version = "1.1.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5a/82/42f767fc1c1143d6fd36efb827202a2d997a375e160a71eb2888a925aac1/pathspec-1.1.1.tar.gz", hash = "sha256:17db5ecd524104a120e173814c90367a96a98d07c45b2e10c2f3919fff91bf5a", size = 135180, upload-time = "2026-04-27T01:46:08.907Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/d9/7fb5aa316bc299258e68c73ba3bddbc499654a07f151cba08f6153988714/pathspec-1.1.1-py3-none-any.whl", hash = "sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189", size = 57328, upload-time = "2026-04-27T01:46:07.06Z" }, +] + +[[package]] +name = "platformdirs" +version = "4.11.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/50/bb/ebc6636e1ae41314f796ebb7215fd28febb45f9aac72f2b04cb74b5071dc/platformdirs-4.11.4.tar.gz", hash = "sha256:f3373be828247211d0febabea97e238c3dfde8a60b3c90c32756fb52cb21556d", size = 34079, upload-time = "2026-08-24T14:53:49.676Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/28/be/0ff05fcd2938fb58ad9219bd54135968342d214737e012d62d43f06a2dd6/platformdirs-4.11.4-py3-none-any.whl", hash = "sha256:e34ff91a24bcddc6d939b878bdf3f5c437c9c46fe9e212b1bf455fdf1ee57586", size = 23741, upload-time = "2026-08-24T14:53:48.406Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "pygments" +version = "2.21.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/2e/ced460408999b33da6b31b0021b0f37d329e202d4169aeb164493778f25b/pygments-2.21.0.tar.gz", hash = "sha256:610ca751c9bc2492b38eb9a38a7fbc93edbbb2d7182edaf34e66ae493dee5c8c", size = 5005329, upload-time = "2026-08-17T08:02:48.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/46/17f022dd3e953bf20a04a028a21ec746d942f8d2af30fa0f124fa0e6a684/pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9", size = 1250147, upload-time = "2026-08-17T08:02:44.912Z" }, +] + +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + +[[package]] +name = "pytokens" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b6/34/b4e015b99031667a7b960f888889c5bd34ef585c85e1cb56a594b92836ac/pytokens-0.4.1.tar.gz", hash = "sha256:292052fe80923aae2260c073f822ceba21f3872ced9a68bb7953b348e561179a", size = 23015, upload-time = "2026-01-30T01:03:45.924Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/dc/08b1a080372afda3cceb4f3c0a7ba2bde9d6a5241f1edb02a22a019ee147/pytokens-0.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8bdb9d0ce90cbf99c525e75a2fa415144fd570a1ba987380190e8b786bc6ef9b", size = 160720, upload-time = "2026-01-30T01:03:13.843Z" }, + { url = "https://files.pythonhosted.org/packages/64/0c/41ea22205da480837a700e395507e6a24425151dfb7ead73343d6e2d7ffe/pytokens-0.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5502408cab1cb18e128570f8d598981c68a50d0cbd7c61312a90507cd3a1276f", size = 254204, upload-time = "2026-01-30T01:03:14.886Z" }, + { url = "https://files.pythonhosted.org/packages/e0/d2/afe5c7f8607018beb99971489dbb846508f1b8f351fcefc225fcf4b2adc0/pytokens-0.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:29d1d8fb1030af4d231789959f21821ab6325e463f0503a61d204343c9b355d1", size = 268423, upload-time = "2026-01-30T01:03:15.936Z" }, + { url = "https://files.pythonhosted.org/packages/68/d4/00ffdbd370410c04e9591da9220a68dc1693ef7499173eb3e30d06e05ed1/pytokens-0.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:970b08dd6b86058b6dc07efe9e98414f5102974716232d10f32ff39701e841c4", size = 266859, upload-time = "2026-01-30T01:03:17.458Z" }, + { url = "https://files.pythonhosted.org/packages/a7/c9/c3161313b4ca0c601eeefabd3d3b576edaa9afdefd32da97210700e47652/pytokens-0.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:9bd7d7f544d362576be74f9d5901a22f317efc20046efe2034dced238cbbfe78", size = 103520, upload-time = "2026-01-30T01:03:18.652Z" }, + { url = "https://files.pythonhosted.org/packages/8f/a7/b470f672e6fc5fee0a01d9e75005a0e617e162381974213a945fcd274843/pytokens-0.4.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4a14d5f5fc78ce85e426aa159489e2d5961acf0e47575e08f35584009178e321", size = 160821, upload-time = "2026-01-30T01:03:19.684Z" }, + { url = "https://files.pythonhosted.org/packages/80/98/e83a36fe8d170c911f864bfded690d2542bfcfacb9c649d11a9e6eb9dc41/pytokens-0.4.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97f50fd18543be72da51dd505e2ed20d2228c74e0464e4262e4899797803d7fa", size = 254263, upload-time = "2026-01-30T01:03:20.834Z" }, + { url = "https://files.pythonhosted.org/packages/0f/95/70d7041273890f9f97a24234c00b746e8da86df462620194cef1d411ddeb/pytokens-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dc74c035f9bfca0255c1af77ddd2d6ae8419012805453e4b0e7513e17904545d", size = 268071, upload-time = "2026-01-30T01:03:21.888Z" }, + { url = "https://files.pythonhosted.org/packages/da/79/76e6d09ae19c99404656d7db9c35dfd20f2086f3eb6ecb496b5b31163bad/pytokens-0.4.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f66a6bbe741bd431f6d741e617e0f39ec7257ca1f89089593479347cc4d13324", size = 271716, upload-time = "2026-01-30T01:03:23.633Z" }, + { url = "https://files.pythonhosted.org/packages/79/37/482e55fa1602e0a7ff012661d8c946bafdc05e480ea5a32f4f7e336d4aa9/pytokens-0.4.1-cp314-cp314-win_amd64.whl", hash = "sha256:b35d7e5ad269804f6697727702da3c517bb8a5228afa450ab0fa787732055fc9", size = 104539, upload-time = "2026-01-30T01:03:24.788Z" }, + { url = "https://files.pythonhosted.org/packages/30/e8/20e7db907c23f3d63b0be3b8a4fd1927f6da2395f5bcc7f72242bb963dfe/pytokens-0.4.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:8fcb9ba3709ff77e77f1c7022ff11d13553f3c30299a9fe246a166903e9091eb", size = 168474, upload-time = "2026-01-30T01:03:26.428Z" }, + { url = "https://files.pythonhosted.org/packages/d6/81/88a95ee9fafdd8f5f3452107748fd04c24930d500b9aba9738f3ade642cc/pytokens-0.4.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:79fc6b8699564e1f9b521582c35435f1bd32dd06822322ec44afdeba666d8cb3", size = 290473, upload-time = "2026-01-30T01:03:27.415Z" }, + { url = "https://files.pythonhosted.org/packages/cf/35/3aa899645e29b6375b4aed9f8d21df219e7c958c4c186b465e42ee0a06bf/pytokens-0.4.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d31b97b3de0f61571a124a00ffe9a81fb9939146c122c11060725bd5aea79975", size = 303485, upload-time = "2026-01-30T01:03:28.558Z" }, + { url = "https://files.pythonhosted.org/packages/52/a0/07907b6ff512674d9b201859f7d212298c44933633c946703a20c25e9d81/pytokens-0.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:967cf6e3fd4adf7de8fc73cd3043754ae79c36475c1c11d514fc72cf5490094a", size = 306698, upload-time = "2026-01-30T01:03:29.653Z" }, + { url = "https://files.pythonhosted.org/packages/39/2a/cbbf9250020a4a8dd53ba83a46c097b69e5eb49dd14e708f496f548c6612/pytokens-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:584c80c24b078eec1e227079d56dc22ff755e0ba8654d8383b2c549107528918", size = 116287, upload-time = "2026-01-30T01:03:30.912Z" }, + { url = "https://files.pythonhosted.org/packages/c6/78/397db326746f0a342855b81216ae1f0a32965deccfd7c830a2dbc66d2483/pytokens-0.4.1-py3-none-any.whl", hash = "sha256:26cef14744a8385f35d0e095dc8b3a7583f6c953c2e3d269c7f82484bf5ad2de", size = 13729, upload-time = "2026-01-30T01:03:45.029Z" }, +] + [[package]] name = "saferpickle" version = "0.1.0" -source = { virtual = "." } +source = { editable = "." } dependencies = [ { name = "absl-py" }, { name = "immutabledict" }, + { name = "pytest" }, +] + +[package.dev-dependencies] +dev = [ + { name = "black" }, + { name = "isort" }, + { name = "pytest" }, ] [package.metadata] requires-dist = [ { name = "absl-py", specifier = ">=2.5.0" }, { name = "immutabledict", specifier = ">=4.3.1" }, + { name = "pytest", specifier = ">=9.1.1" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "black", specifier = ">=26.5.1" }, + { name = "isort", specifier = ">=8.0.1" }, + { name = "pytest", specifier = ">=8.0.0" }, ] From f83f68c4e813c21d92a4b91373b68575cb4e4ad8 Mon Sep 17 00:00:00 2001 From: Vijay Panchal Date: Wed, 26 Aug 2026 01:02:23 +0530 Subject: [PATCH 4/5] fix: remove pytest from runtime dependencies --- pyproject.toml | 1 - uv.lock | 2 -- 2 files changed, 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 793b75c..6ed9c1f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,6 @@ authors = [ dependencies = [ "absl-py>=2.5.0", "immutabledict>=4.3.1", - "pytest>=9.1.1", ] classifiers = [ "Programming Language :: Python :: 3", diff --git a/uv.lock b/uv.lock index d6e38c5..7f63d12 100644 --- a/uv.lock +++ b/uv.lock @@ -187,7 +187,6 @@ source = { editable = "." } dependencies = [ { name = "absl-py" }, { name = "immutabledict" }, - { name = "pytest" }, ] [package.dev-dependencies] @@ -201,7 +200,6 @@ dev = [ requires-dist = [ { name = "absl-py", specifier = ">=2.5.0" }, { name = "immutabledict", specifier = ">=4.3.1" }, - { name = "pytest", specifier = ">=9.1.1" }, ] [package.metadata.requires-dev] From 87cf976e4258c1ed77a1a4ba6ebc67a5710d90ff Mon Sep 17 00:00:00 2001 From: Vijay Panchal Date: Wed, 26 Aug 2026 01:30:16 +0530 Subject: [PATCH 5/5] chore: remove stray root package marker --- __init__.py | 1 - 1 file changed, 1 deletion(-) delete mode 100644 __init__.py diff --git a/__init__.py b/__init__.py deleted file mode 100644 index 0e632e1..0000000 --- a/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Package marker