Python装饰器实现MongoDB参数清洗：处理嵌套列表与关键字参数

原装饰器仅能处理顶层位置参数，无法覆盖嵌套列表（如$in查询场景）和关键字参数的清洗需求，以下是改进后的实现：

改进后的代码 #

def sanitize(function_to_decorate):
    def _sanitize_value(value):
        # 检查字符串是否包含MongoDB注入风险字符
        if isinstance(value, str):
            if any(elem in value for elem in r'${/}\\'):
                raise ValueError(f"Invalid character found in value: {repr(value)}")
        # 递归处理列表/元组等可迭代结构
        elif isinstance(value, (list, tuple)):
            for item in value:
                _sanitize_value(item)
        # 其他类型（数字、布尔、None等）无需检查
        return value

    def wrapper(*args, **kwargs):
        # 清洗位置参数
        for arg in args:
            _sanitize_value(arg)
        # 清洗关键字参数
        for key, value in kwargs.items():
            _sanitize_value(value)
        # 执行原函数
        return function_to_decorate(*args, **kwargs)
    return wrapper

# 测试示例
@sanitize
def test(arg, arg2, list_arg=None, **kwargs):
    print(arg, arg2, list_arg, kwargs)

# 合法调用：通过
test('Hi', 'Me', ['1', '2'], extra='test')
# 嵌套列表含非法字符：抛出异常
# test('Hi', 'Me', ['Me', '{1'])
# 关键字参数含非法字符：抛出异常
# test('Hi', 'Me', extra='${test}')
# 多层嵌套含非法字符：抛出异常
# test('Hi', 'Me', ['a', ['b', '$c']])

关键改进点 #

• 递归检查逻辑：通过_sanitize_value函数递归处理所有嵌套结构（列表、元组），确保深层元素也被清洗
• 覆盖关键字参数：新增对**kwargs的遍历清洗，适配带关键字参数的MongoDB查询函数
• 明确错误信息：使用ValueError替代通用Exception，并标注含非法字符的具体值，便于调试
• 类型区分处理：仅对字符串和可迭代结构做检查，跳过数字、布尔等无注入风险的类型

内容的提问来源于stack exchange，提问作者Mauro