Updated whitelist and Python Environment checking

[ShadowPickle-Bench]

Summary

Integrated the whitelist from PyTorch's weights_only unpickler (_weights_only_unpickler.py) into the non_standard_imports check (fickle.py:1295) to reduce false positive detections in pickle files. Additionally, accounted for CWE-1395 within the updated whitelist, specifically addressing the Overwritten-Modules exploit involving pickles from the ShadowPickle paper.

Motivation:

Our work, ShadowPickle, demonstrates that the top 15 libraries detected during their experimental evaluation of Fickling against 3,000 benign pickle models are all present within PyTorch's weights_only scanner whitelist and that Fickling does not whitelist any additional libraries beyond those covered by the weights_only unpickler. Furthermore, these whitelisted libraries are not known to be leveraged in the construction of malicious payloads. Incorporating them into Fickling's existing whitelist would therefore reduce the rate of false positives while preserving true positive detection performance.

However, implicitly trusting third-party libraries to behave as expected introduces a notable security risk (CWE-1395) , as these libraries may be manipulated to execute malicious operations. The ShadowPickle Overwritten-Modules attack exemplifies this threat by overwriting entries in sys.modules, redirecting the original import path of a given module (e.g., torch.FloatStorage, collections.OrderedDict) to an attacker-controlled implementation capable of executing arbitrary payloads. Accordingly, validating the Python environment to ensure that modules are imported from their expected /usr/lib/python/ locations is equally important.

Changes

Added a list of libraries and modules considered part of the whitelist for the weights_only unpickler:

WEIGHTS_ONLY_WHITELIST: frozenset[str] = frozenset(
    [
        "collections.OrderedDict",
        "collections.Counter",
        "torch.nn.parameter.Parameter",
        "torch.serialization._get_layout",
        "torch.Size",
        "torch.Tensor",
        "torch.device",
        ...
    ]
)

To account for module member imports, helper functions create_full_import, is_in_weights_only, and is_sys_modules_tampered were introduced (fickle.py:1512) to determine whether imported modules and members belong to the weights_only unpickler's whitelist. The is_sys_modules_tampered function checks whether the sys.modules path of a built-in library isn’t modified/tampered. Once confirmed that the python environment has not been tampered with, the create_full_import function constructs the full module member path (e.g., collections.OrderedDict), which is then compared against the whitelist via is_in_weights_only.

def non_standard_imports(self) -> Iterator[ast.Import | ast.ImportFrom]:
    for node in self.properties.imports:
        if isinstance(node, ast.ImportFrom):
            # from module import x - check if module is standard
            # module can be None for relative imports (from . import x)
            # Also verify if the imported library has not been tampered with.
            if node.module is not None:
                full_imports = create_full_import(node)
                if not is_sys_modules_tampered(node.module):
                    if not is_std_module(node.module):
                        if not is_in_weights_only(full_imports):
                            yield node
                else:
                    print(f"{full_imports} may have been overwritten with a custom implementation, and may have additional functionality that may compromise the environment. Please ensure the sys.modules path hasn't been tampered with.")
                    yield node
…

Similarly, for whole-module imports, the module path in sys.modules is checked for any case of tampering. If not, each imported name is checked against both the standard module list and the weights_only unpickler whitelist:

else:
  # import x, y, z - check if any name is non-standard or not part of the weights_only whitelist
  # Also verify if the imported library has not been tampered with.
  if any(is_sys_modules_tampered(alias.name) or (not is_std_module(alias.name) and not is_in_weights_only([alias.name])) for alias in node.names):
      yield node

Tests

Two test cases were added to adequately validate this patch:

• test_weights_only_whitelist: Verifies the scanner's behavior against all libraries included in the weights_only unpickler's whitelist.
• test_sys_modules_tampering: Verifies the scanner's ability to detect an instance of sys.modules path tampering.

All tests in the test folder were executed and passed, confirming that default behavior remains intact.

[CLAassistant]

All committers have signed the CLA.