token-vault

1import type { Register } from 'claude-code';
2
3// 当日本地日期前缀（YYYYMMDD）。编号按天分区，跨天自动重新从 001 计数。
4export const todayPrefix = (d = new Date()): string => {
5  const p = (x: number) => String(x).padStart(2, '0');
6  return `${d.getFullYear()}${p(d.getMonth() + 1)}${p(d.getDate())}`;
7};
8
9export const pad3 = (n: number): string => String(n).padStart(3, '0');
10
11// 全箱 FIFO 淘汰计划：日期与序号均定宽，编号字符串排序即时间序，
12// 超出上限时淘汰最旧的一批。抽成纯函数直接单测——事件级构造 >999 键
13// 需要真实跨日时钟，不值得为此引入时钟注入。
14export const planEvict = (nums: string[], cap: number): string[] =>
15  nums.length <= cap ? [] : [...nums].sort().slice(0, nums.length - cap);
16
17// 常见 token 形态：GitHub PAT（经典/细粒度/各前缀）、Anthropic key、通用 sk-、
18// GitLab PAT、Slack token。长度取下限贪婪匹配——超出标称长度的尾部一并截获，
19// 不假设 token 恰好是文档标称长度。命中即截获入箱，原文中的值全部替换为占位符。
20const TOKEN_RE =
21  /(ghp_[A-Za-z0-9]{20,}|gho_[A-Za-z0-9]{20,}|ghu_[A-Za-z0-9]{20,}|ghs_[A-Za-z0-9]{20,}|github_pat_[A-Za-z0-9_]{20,}|sk-ant-[A-Za-z0-9_-]{24,}|sk-[A-Za-z0-9_-]{20,}|glpat-[A-Za-z0-9_-]{20,}|xox[baprs]-[A-Za-z0-9-]{10,})/g;
22
23export const register: Register = (on) => {
24  // 模型可调用的落值工具：参数只有目标路径、变量名和入箱编号，值始终留在本地保险箱。
25  on('session.start', async ($, e, next) => {
26    await $.tool.register({
27      name: 'vault-apply',
28      description:
29        '把用户存入本地保险箱的 token 写入目标文件（一行 KEY=value）。token 的值对模型不可见，模型只提供路径、变量名和入箱编号。',
30      inputSchema: {
31        type: 'object',
32        properties: {
33          target: { type: 'string', description: '目标文件的绝对路径' },
34          key: { type: 'string', description: '变量名，默认 TOKEN' },
35          index: { type: 'string', description: '占位符中的入箱编号（YYYYMMDD-NNN，如 20260101-001）；箱内仅一个 token 时可省略' },
36        },
37        required: ['target'],
38      },
39    });
40    return next(e);
41  });
42
43  // 落值严格化：编号必须显式且存在，报错列出可用编号并禁止模型改用其他编号。
44  // 缺省回退与淘汰后猜编号是"写错值"的全部来源，这里一并堵死。
45  on('tool.call', { tool: 'mcp__token-vault__vault-apply' }, async ($, e, next) => {
46    const args = e as { target?: string; key?: string; index?: string };
47    if (!args.target) {
48      return { result: 'token-vault：缺少 target 参数' };
49    }
50    const nums = (await $.store.keys())
51      .filter((k) => k.startsWith('token:'))
52      .map((k) => k.slice('token:'.length))
53      .sort();
54    const today = todayPrefix();
55    const avail = () => {
56      const ordered = [
57        ...nums.filter((n) => n.startsWith(`${today}-`)),
58        ...nums.filter((n) => !n.startsWith(`${today}-`)),
59      ];
60      if (!ordered.length) return '（保险箱为空：请用户把 token 直接发在消息里，会自动入箱脱敏）';
61      const shown = ordered.slice(0, 20).join('、');
62      return ordered.length > 20 ? `${shown}…等 ${ordered.length} 个` : shown;
63    };
64    const apply = async (n: string) => {
65      const value = await $.store.get(`token:${n}`);
66      if (typeof value !== 'string' || !value) return undefined;
67      const key = args.key ?? 'TOKEN';
68      await $.fs.write(args.target!, `${key}=${value}\n`);
69      return `已写入 ${key}=***（${value.length} 字符）到 ${args.target}，值未经过模型`;
70    };
71    // index 违反 schema 传成非字符串（如 JSON number）时归一为字符串走格式报错，
72    // 而不是让 match 抛 TypeError 落进宿主的晦涩 HooksError、断掉恢复指引。
73    const idx = args.index === undefined || args.index === null ? '' : String(args.index);
74    if (!idx) {
75      if (nums.length === 1) {
76        const r = await apply(nums[0]);
77        if (r) return { result: r };
78        // 键在而值异常（正常流程不产生，仅手动改存储可触发）：不再重试，避免与
79        // 「提供 index」指引形成死循环，由用户删除存储文件清理。
80      }
81      return { result: `token-vault：请提供 index（占位符中的入箱编号）。当前可用：${avail()}；禁止改用其他编号落值，请让用户重发` };
82    }
83    const m = idx.match(/^(\d{8})-(\d{1,3})$/);
84    if (!m) {
85      return { result: `token-vault：编号格式应为 YYYYMMDD-NNN。当前可用：${avail()}；禁止改用其他编号落值，请让用户重发` };
86    }
87    const n = `${m[1]}-${m[2].padStart(3, '0')}`;
88    const r = await apply(n);
89    if (r) return { result: r };
90    return { result: `token-vault：编号 ${n} 不存在（可能已被 FIFO 淘汰）。当前可用：${avail()}；禁止改用其他编号落值，请让用户重发该 token` };
91  });
92
93  // 入口脱敏：带 token 的消息在进入会话前把值换成占位符，真值进 $.store，
94  // 模型与 transcript 从头到尾只见占位符。
95  // 编号 YYYYMMDD-NNN：按天分区，每日从 001 起、当天内只增不复用；
96  // 日期与序号均定宽，字符串排序即时间序。去重仅当天范围，跨天重发入新编号。
97  // 全箱至多 MAX_KEYS 个键，超出按时间序 FIFO 淘汰最旧（无墓碑，查不到即没有）。
98  // 单日新值将超过 MAX_PER_DAY 时 fail-closed 整体拦截，报错指引删存储文件恢复；
99  // 上限取 999 与 NNN 三位编号格式一致，编号恒为三位。
100  // 升级前的存量纯数字键不参与当天去重与 nextN；其字符串序与时间序不对应
101  // （'12' 排最前、'5' 排在全部日期键之后），过渡期 FIFO 顺序对它们不精确，
102  // 属已拍板的过渡取舍，必要时删除存储文件一次清掉。
103  const MAX_PER_DAY = 999;
104  const MAX_KEYS = 999;
105
106  on('prompt.submit', async ($, e, next) => {
107    if (e.command) return next(e);
108    const found = e.text.match(TOKEN_RE);
109    if (!found?.length) return next(e);
110    const today = todayPrefix();
111    const known = new Map<string, string>(); // 真值 → 编号（仅当天）
112    const nums: string[] = []; // 全箱编号，字符串升序即时间序
113    let nextN = 0;
114    for (const k of await $.store.keys()) {
115      if (!k.startsWith('token:')) continue;
116      const n = k.slice('token:'.length);
117      nums.push(n);
118      if (n.startsWith(`${today}-`)) {
119        const v = await $.store.get(k);
120        if (typeof v === 'string') known.set(v, n);
121        const m = Number(n.slice(today.length + 1));
122        if (m > nextN) nextN = m;
123      }
124    }
125    // 单日额度预检：本条消息的全部新值都要能入箱才动手，避免拦一半留一半。
126    const fresh = new Set(found.filter((t) => !known.has(t)));
127    if (known.size + fresh.size > MAX_PER_DAY) {
128      return {
129        drop: `token-vault：今日脱敏额度已满（${MAX_PER_DAY}），信息脱敏不可用；请删除 ~/.claude/plugins/store/ 下 token-vault_ 开头的 JSON 文件（全箱清空）后重试`,
130      };
131    }
132    let text = e.text;
133    let added = 0;
134    for (const t of found) {
135      let n = known.get(t);
136      if (n === undefined) {
137        added += 1;
138        nextN += 1;
139        n = `${today}-${pad3(nextN)}`;
140        await $.store.set(`token:${n}`, t);
141        nums.push(n);
142        known.set(t, n);
143      }
144      text = text.replace(t, `[token-vault:${n}已入箱脱敏;需要落值时调用vault-apply并传此编号]`);
145    }
146    const evicted = planEvict(nums, MAX_KEYS);
147    for (const n of evicted) await $.store.delete(`token:${n}`);
148    const dup = found.length - added;
149    const parts: string[] = [];
150    if (added > 0) parts.push(`新入箱 ${added}`);
151    if (dup > 0) parts.push(`去重 ${dup}`);
152    if (evicted.length > 0) parts.push(`淘汰最旧 ${evicted.length}`);
153    $.ui.toast(
154      parts.length
155        ? `token-vault：截获 ${found.length} 个 token（${parts.join('，')}）`
156        : `token-vault：截获 ${found.length} 个 token，均已入箱`,
157    );
158    return next({ ...e, text });
159  }).catch(async ($, e, next) => {
160    // fail-closed：消息确实含 token 而脱敏没完成时宁可整体拦截，不放原文过。
161    if (next.called || !e.text?.match(TOKEN_RE)) return next(e);
162    return { drop: 'token-vault：消息疑似含 token 但脱敏失败，已整体拦截' };
163  });
164};
165