Resk-Security/reskSecure

# reskSecure **基于 Bitmask 的 LLM 安全防火墙。** 这是一个 Python 包，它根据以 bitmask 编码的用户权限，限制语言模型可以生成的内容。它的工作原理是拦截模型的 token 预测，并在不允许的短语出现在输出中之前，将其拦截或施加惩罚。 与基于提示词的过滤器（可能会被越狱绕过）或生成后的内容审核（在检测之前就已经让违规内容泄露）不同，reskSecure 作用于 logits 级别——直接在模型的生成循环内部进行处理。每个 token 在输出之前都必须通过安全策略的检查。 ## 工作原理 ``` User request with bitmask (e.g. 7) | v BitmaskLogitsProcessor intercepts each token prediction | v For every candidate token, the Aho-Corasick automaton (from resklogits) checks if selecting it would start or complete a banned phrase | v Hard-mode phrases: the token's logit is set to -inf (impossible to generate) Bias-mode phrases: the token's logit is reduced by a configurable penalty | v On complete match: EOS token is forced, generation stops immediately | v Output response or tool call ``` 工具调用也会在 token 级别被拦截。如果用户的 bitmask 不包含某个工具所需的位，该工具的触发短语（如 `"send_email("` 或 `"create_ticket("`）将被自动添加到 hard 模式的拦截列表中。无论通过何种提示词工程或越狱尝试，模型都无法生成被禁工具调用的第一个 token。 后置生成检查（`verify_tool_action`）也可用作深度防御层，但主要的保护机制仍在于 logits 级别。 ## 为什么不直接使用提示词工程？ - 提示词注入可以绕过基于指令的过滤器 - 生成后的正则表达式或分类器扫描虽然能在违规内容出现后将其捕获，但违规内容此时已经被生成了 - logits 级别的过滤会在 token 被采样之前将其拦截——模型永远不会将禁用序列视为补全的候选项 ## 功能 - **两种严重性模式**：`hard` 完全拦截 token (-inf logit)，`bias` 通过可配置的惩罚值来降低其概率 - **Strict 模式**：一旦生成的生成前缀匹配到被禁短语的开头，即使完整短语尚未形成，也会强制结束序列 - **GPU 加速模式匹配**：使用 resklogits 包中的 VectorizedAhoCorasick 进行快速的 token 扫描 - **策略系统**：通过 YAML 配置，将功能 bitmask 与短语规则和工具权限关联起来 - **热重载**：PolicyWatcher 会检测文件变动，并在无需重启服务器的情况下重建自动机 - **线程安全缓存**：自动机通过 (mask, model_name) 进行缓存，并带有可配置的 TTL - **工具防护**：在生成后对工具调用执行进行 bitmask 检查 - **不处理 JWT**：该包接收原始的整数 bitmask；身份验证和 JWT 解码由调用的应用程序负责处理 ## 环境要求 - Python >= 3.13 - PyTorch >= 2.0.0 - transformers >= 4.35.0 - resklogits >= 0.1.0 ## 安装 ``` pip install resksecure ``` 从源码安装： ``` git clone https://github.com/Resk-Security/reskSecure.git cd reskSecure pip install -e . ``` ## 快速开始 定义一个策略文件 (`policy.yaml`)： ``` version: "1.0" policies: - mask: 7 name: contributor strict: false default: true rules: - phrase: "DROP TABLE" mode: hard - phrase: "DELETE FROM" mode: hard - phrase: "salaries" mode: bias penalty: -5.0 tools: read_email: required_bit: 0 send_email: required_bit: 1 read_sql: required_bit: 2 ``` 在你的生成 pipeline 中使用它： ``` from resksecure import BitmaskLogitsProcessor, load_policy, verify_tool_action policy_set = load_policy("policy.yaml") processor = BitmaskLogitsProcessor( mask=7, model_name="mistralai/Mistral-7B-v0.1", tokenizer=tokenizer, policy_set=policy_set, device="cuda", ) outputs = model.generate(**inputs, logits_processor=[processor]) # 可选：针对 bitmask 验证 tool calls if has_tool_call(response): if not verify_tool_action("send_email", user_mask=7, policy_set=policy_set): raise PermissionError("Action not authorized") ``` ## 策略参考 策略级别字段： | 字段 | 类型 | 描述 | |-------------|---------|----------------------------------------------------------| | `mask` | int | 用于标识此策略的功能 bitmask | | `name` | string | 人类可读的策略名称 | | `strict` | bool | 如果为 true，遇到第一个被禁前缀时即停止生成 | | `default` | bool | 如果为 true，当没有精确匹配的 mask 时将使用此策略 | | `rules` | list | 短语规则列表（见下文） | | `tools` | dict | 工具名称到工具规则对象的映射（见下文） | 短语规则字段： | 字段 | 类型 | 描述 | |-----------|--------|------------------------------------------------------------| | `phrase` | string | 要拦截或惩罚的文本模式 | | `mode` | string | 可以是 `hard`（完全拦截）或 `bias`（减少 logits） | | `penalty` | float | bias 模式下的 logit 惩罚值（例如 -5.0） | 工具规则字段： | 字段 | 类型 | 描述 | |--------------------|----------|----------------------------------------------------------------| | `required_bit` | int | 用户 mask 中必须设置的位，才能使用此工具 | | `trigger_phrases` | string[] | 启动此工具调用的短语（如果未设置该位则会被拦截） | ## 包结构 ``` reskSecure/ src/resksecure/ __init__.py Public exports, version policy_loader.py YAML parsing, Policy/PolicySet models trie_factory.py Builds VectorizedAhoCorasick from a Policy bitmask_processor.py BitmaskLogitsProcessor (LogitsProcessor subclass) tool_guard.py Post-generation tool action verification cache.py Thread-safe TTL cache for automata policy_watcher.py Hot-reload daemon that watches YAML mtime config/example_policy.yaml tests/ examples/ ``` ## 许可证 本软件根据 RESK 软件许可证进行授权。商业使用需要单独购买付费许可证。详情请参阅 [LICENSE](LICENSE) 文件。 如需咨询商业授权事宜，请联系：contact@resk.fr

标签：Aho-Corasick, DLL 劫持, Logits拦截, Python, 凭据扫描, 大语言模型, 无后门, 逆向工具