如何使用Python检测不安全的pickle操作？

pickle

__reduce__

__reduce__

os

os.system('rm -rf /')

pickletools

pickletools.dis()

GLOBAL

REDUCE

import pickle
import pickletools
import io

# 正常情况下的pickle数据
class MySafeClass:
    def __init__(self, name):
        self.name = name

safe_obj = MySafeClass("hello")
safe_data = pickle.dumps(safe_obj)

print("--- 安全数据的pickletools反汇编 ---")
pickletools.dis(safe_data)

# 构造一个潜在的恶意pickle数据
# 警告：不要在生产环境或不可信环境中运行此代码，它展示了潜在的攻击方式。
class Malicious:
    def __reduce__(self):
        # 尝试执行一个系统命令，例如 'echo Hacked!'
        # 在真实攻击中，这里可能是更破坏性的命令
        return (eval, ("__import__('os').system('echo Hacked!')",))

mal_obj = Malicious()
mal_data = pickle.dumps(mal_obj)

print("\n--- 恶意数据的pickletools反汇编 ---")
pickletools.dis(mal_data)

# 实际检查时，你可以解析dis的输出
# 这是一个简化的检查逻辑，实际可能需要更复杂的模式匹配
def check_for_dangerous_opcodes(data):
    stream = io.BytesIO(data)
    for opcode, arg, pos in pickletools.genops(stream):
        # 关注GLOBAL (导入模块) 和 REDUCE (执行__reduce__方法)
        # 还有 BUILD (构建对象) 如果和 GLOBAL/REDUCE 组合
        if opcode.name in ('GLOBAL', 'REDUCE'):
            print(f"检测到潜在危险操作码: {opcode.name} (参数: {arg})")
            # 更复杂的逻辑可能需要检查arg的值，判断是否是危险模块或函数
            return True
    return False

print("\n--- 检查恶意数据 ---")
if check_for_dangerous_opcodes(mal_data):
    print("警告：该pickle数据可能包含危险操作！")
else:
    print("未检测到明显危险操作码。")

print("\n--- 检查安全数据 ---")
if check_for_dangerous_opcodes(safe_data):
    print("警告：该pickle数据可能包含危险操作！")
else:
    print("未检测到明显危险操作码。")

GLOBAL

__import__

os.system

pickle.Unpickler

find_class

import pickle
import sys

# 定义一个允许加载的模块和类的白名单
# 仅允许加载内置类型和我们明确定义的MySafeClass
SAFE_MODULES = {
    '__builtins__': ['dict', 'list', 'tuple', 'set', 'int', 'float', 'str', 'bool', 'NoneType'],
    '__main__': ['MySafeClass'], # 假设MySafeClass定义在当前脚本中
}

class RestrictedUnpickler(pickle.Unpickler):
    def find_class(self, module, name):
        # 检查模块是否在白名单中
        if module not in SAFE_MODULES:
            raise pickle.UnpicklingError(f"Attempted to load unauthorized module: {module}")

        # 检查类是否在白名单中
        if name not in SAFE_MODULES[module]:
            raise pickle.UnpicklingError(f"Attempted to load unauthorized class: {module}.{name}")

        # 如果都在白名单中，则调用父类的find_class方法加载
        return super().find_class(module, name)

# 定义一个安全的类
class MySafeClass:
    def __init__(self, data):
        self.data = data
    def __repr__(self):
        return f"MySafeClass(data={self.data})"

# 定义一个尝试执行危险操作的类 (用于测试限制)
class DangerousClass:
    def __reduce__(self):
        return (eval, ("__import__('os').system('echo Evil!')",))

# --- 测试安全数据 ---
safe_obj = MySafeClass("This is safe data.")
safe_pickled_data = pickle.dumps(safe_obj)

print("--- 尝试加载安全数据 ---")
try:
    loaded_safe_obj = RestrictedUnpickler(io.BytesIO(safe_pickled_data)).load()
    print(f"成功加载安全对象: {loaded_safe_obj}")
except pickle.UnpicklingError as e:
    print(f"加载安全数据失败: {e}")
except Exception as e:
    print(f"发生其他错误: {e}")

# --- 测试不安全数据 ---
dangerous_obj = DangerousClass()
dangerous_pickled_data = pickle.dumps(dangerous_obj)

print("\n--- 尝试加载不安全数据 ---")
try:
    # 假设DangerousClass在当前模块，但我们没有把它加入白名单
    # 如果它在其他模块，而那个模块不在白名单，效果也一样
    loaded_dangerous_obj = RestrictedUnpickler(io.BytesIO(dangerous_pickled_data)).load()
    print(f"成功加载不安全对象: {loaded_dangerous_obj} (这不应该发生!)")
except pickle.UnpicklingError as e:
    print(f"加载不安全数据失败 (预期): {e}")
except Exception as e:
    print(f"发生其他错误: {e}")

# 如果DangerousClass是在一个单独的模块里，并且那个模块不在SAFE_MODULES里，也会被阻止
# 比如：
# # malicious_module.py
# import os
# class MaliciousPayload:
#     def __reduce__(self):
#         return (os.system, ('echo Malicious Code Executed!',))
#
# # 在主程序中：
# # import malicious_module
# # mal_obj_from_other_module = malicious_module.MaliciousPayload()
# # mal_pickled_data_from_other_module = pickle.dumps(mal_obj_from_other_module)
# # try:
# #     RestrictedUnpickler(io.BytesIO(mal_pickled_data_from_other_module)).load()
# # except pickle.UnpicklingError as e:
# #     print(f"阻止了加载外部恶意模块: {e}")

Unpickler

os.system

os

os.system

find_class