secret-vault

1import { atom, read, update } from 'claude-code'
2import type { ApiContentBlock, Register } from 'claude-code'
3
4import type { Vault, VaultEntry } from '../types'
5
6const vault = atom({ plugin: 'secret-vault', key: 'vault' } as const, { entries: [], issued: 0 } as Vault)
7
8const PLACEHOLDER = /⟨secret:(\d+)⟩/g
9const placeholder = (n: number) => `⟨secret:${n}⟩`
10
11// Masking a value shorter than this would rewrite ordinary words in output.
12const MIN_VALUE = 4
13
14// `group` names the capture that is the secret; without it the whole match is.
15type Rule = { kind: string; re: RegExp; group?: number }
16
17// Order matters: a JWT is taken before the Bearer rule sees it.
18const RULES: Rule[] = [
19  { kind: 'private-key', re: /-----BEGIN [A-Z ]*PRIVATE KEY-----[\s\S]*?-----END [A-Z ]*PRIVATE KEY-----/g },
20  { kind: 'jwt', re: /\beyJ[\w-]{4,}\.eyJ[\w-]{4,}\.[\w-]*/g },
21  { kind: 'github', re: /\b(?:gh[pousr]_[A-Za-z0-9]{30,}|github_pat_\w{40,})/g },
22  { kind: 'aws-key-id', re: /\b(?:AKIA|ASIA)[0-9A-Z]{16}\b/g },
23  { kind: 'aws-secret', re: /aws_secret_access_key["']?\s*[=:]\s*["']?([A-Za-z0-9/+=]{40})/gi, group: 1 },
24  { kind: 'slack', re: /\bxox[abprs]-[A-Za-z0-9-]{10,}/g },
25  { kind: 'api-key', re: /\bsk-(?:ant-)?[\w-]{20,}/g },
26  { kind: 'bearer', re: /\bBearer\s+([\w.~+/=-]{16,})/gi, group: 1 },
27  { kind: 'basic-auth', re: /\bAuthorization["']?\s*:\s*["']?(?:Basic|Token)\s+([\w.~+/=-]{8,})/gi, group: 1 },
28  { kind: 'url-password', re: /\b[a-z][\w+.-]*:\/\/[^\s:/@]+:([^\s@/]+)@/gi, group: 1 },
29  { kind: 'password', re: /\b(?:password|passwd|pwd)["']?\s*[=:]\s*["']?([^\s"'&;,]{4,})/gi, group: 1 },
30  { kind: 'api-key', re: /\bx-api-key["']?\s*[=:]\s*["']?([\w.~+/=-]{12,})/gi, group: 1 },
31  {
32    kind: 'token',
33    re: /\b(?:session_?id|access_token|refresh_token|auth_token|api_?key|client_secret)["']?\s*[=:]\s*["']?([\w%.~+/=-]{16,})/gi,
34    group: 1,
35  },
36]
37
38// A Cookie header (or curl -b / --cookie) carries several name=value pairs.
39const COOKIE_LINE = /^.*(?:\bcookie\s*:|--cookie\b|\s-b\s+['"]).*$/gim
40const COOKIE_VALUE = /=([^;\s'"]{16,})/g
41
42export function mask(text: string, entries: readonly VaultEntry[]): string {
43  const longestFirst = [...entries].sort((a, b) => b.value.length - a.value.length)
44  return longestFirst.reduce(
45    (s, e) => (e.value.length >= MIN_VALUE ? s.split(e.value).join(placeholder(e.n)) : s),
46    text,
47  )
48}
49
50export function reveal(text: string, entries: readonly VaultEntry[]): string {
51  return text.replace(PLACEHOLDER, (whole, n: string) => entries.find(e => e.n === Number(n))?.value ?? whole)
52}
53
54// Replaces known values and newly detected secrets; `vault` gains the new ones.
55export function redact(text: string, start: Vault): { text: string; vault: Vault; hits: number } {
56  const entries = [...start.entries]
57  let issued = start.issued
58  let hits = 0
59
60  const take = (value: string, kind: string): string => {
61    hits += 1
62    const known = entries.find(e => e.value === value)
63    if (known) {
64      return placeholder(known.n)
65    }
66    issued += 1
67    entries.push({ n: issued, kind, value })
68    return placeholder(issued)
69  }
70
71  const sub = (whole: string, value: string | undefined, kind: string): string => {
72    if (!value || value.length < MIN_VALUE || value.includes('⟨')) {
73      return whole
74    }
75    const at = whole.lastIndexOf(value)
76    return whole.slice(0, at) + take(value, kind) + whole.slice(at + value.length)
77  }
78
79  let out = text
80  const knownHits = entries.filter(e => e.value.length >= MIN_VALUE && out.includes(e.value)).length
81  out = mask(out, entries)
82  hits += knownHits
83
84  for (const rule of RULES) {
85    out = out.replace(rule.re, (...m: unknown[]) => {
86      const whole = m[0] as string
87      return sub(whole, (rule.group ? m[rule.group] : whole) as string | undefined, rule.kind)
88    })
89  }
90  out = out.replace(COOKIE_LINE, line => line.replace(COOKIE_VALUE, (whole, v: string) => sub(whole, v, 'cookie')))
91
92  return { text: out, vault: { entries, issued }, hits }
93}
94
95function mapStrings(value: unknown, f: (s: string) => string): unknown {
96  if (typeof value === 'string') {
97    return f(value)
98  }
99  if (Array.isArray(value)) {
100    return value.map(v => mapStrings(v, f))
101  }
102  if (value !== null && typeof value === 'object') {
103    return Object.fromEntries(Object.entries(value).map(([k, v]) => [k, mapStrings(v, f)]))
104  }
105  return value
106}
107
108function someString(value: unknown, test: (s: string) => boolean): boolean {
109  if (typeof value === 'string') {
110    return test(value)
111  }
112  if (Array.isArray(value)) {
113    return value.some(v => someString(v, test))
114  }
115  if (value !== null && typeof value === 'object') {
116    return Object.values(value).some(v => someString(v, test))
117  }
118  return false
119}
120
121const holdsSecret = (value: unknown, entries: readonly VaultEntry[]) =>
122  someString(value, s => entries.some(e => e.value.length >= MIN_VALUE && s.includes(e.value)))
123
124// The engine pins these keys; the rest are the tool's own arguments.
125const RESERVED = new Set(['tool', 'tool_use_id', 'agentId', 'consent'])
126
127// Inputs left literal: a file or another agent's transcript would keep the real value.
128const LITERAL_TOOLS = new Set(['Write', 'Edit', 'MultiEdit', 'NotebookEdit', 'Agent', 'Task', 'SendMessage'])
129
130function maskBlock(block: ApiContentBlock, entries: readonly VaultEntry[]): ApiContentBlock {
131  const m = (s: string) => mask(s, entries)
132  if (block.type === 'text' && typeof block.text === 'string') {
133    return { ...block, text: m(block.text) }
134  }
135  if (block.type === 'tool_result') {
136    return { ...block, content: mapStrings(block.content, m) }
137  }
138  return block
139}
140
141const NOTE =
142  'secret-vault: secrets in this prompt were replaced with ⟨secret:N⟩ placeholders. Use a placeholder verbatim ' +
143  'where its value is needed in a Bash command, WebFetch or MCP tool input; the real value is put in just before ' +
144  'the tool runs and masked again in its output. Do not ask the user to paste the value again.'
145
146const SECTION =
147  '# Secret placeholders\n' +
148  'A value written ⟨secret:N⟩ stands for a real secret the user gave in this session (a token, password or key). ' +
149  'Use the placeholder verbatim where the value is needed in Bash commands, WebFetch and MCP tool inputs: the real ' +
150  'value is put in just before the tool runs and masked back to the placeholder in the output. Inputs of Write, ' +
151  'Edit, MultiEdit, NotebookEdit, Agent and SendMessage are not substituted, so a placeholder written there stays ' +
152  'literal; to put a secret into a file, write it with a Bash command. Do not try to print or decode a placeholder.'
153
154const statusOf = (n: number) => (n > 0 ? `vault: ${n} secret${n === 1 ? '' : 's'}` : undefined)
155
156const hint = (value: string) => (value.length >= 12 ? `${value.slice(0, 4)}… (${value.length} chars)` : `… (${value.length} chars)`)
157
158export const register: Register = on => {
159  on('session.start', async ($, e, next) => {
160    await $.command.register({
161      name: 'vault',
162      description: 'List or clear the secrets replaced with ⟨secret:N⟩ placeholders this session',
163      argumentHint: '[list|clear]',
164    })
165    $.ui.status(statusOf((await read($, vault)).entries.length))
166    return next(e)
167  })
168
169  on('prompt.submit', async ($, e, next) => {
170    const scan = redact(e.text, await read($, vault))
171    if (scan.hits === 0) {
172      return next(e)
173    }
174    const saved = await update($, vault, v => redact(e.text, v).vault)
175    $.ui.status(statusOf(saved.entries.length))
176    return next({ ...e, text: redact(e.text, saved).text, context: [...(e.context ?? []), NOTE] })
177  })
178
179  on('prompt.compose', async ($, e, next) => {
180    const composed = await next(e)
181    if ((await read($, vault)).entries.length === 0) {
182      return composed
183    }
184    return { sections: [...composed.sections, { id: 'secret-vault:placeholders', text: SECTION, scope: 'session' as const }] }
185  })
186
187  on('tool.call', async ($, e, next) => {
188    const { entries } = await read($, vault)
189    if (entries.length === 0) {
190      return next(e)
191    }
192
193    const literal = LITERAL_TOOLS.has(String(e.tool))
194    const input = literal
195      ? e
196      : (Object.fromEntries(
197          Object.entries(e).map(([k, v]) => [k, RESERVED.has(k) ? v : mapStrings(v, s => reveal(s, entries))]),
198        ) as typeof e)
199    const ran = await next(input)
200
201    if (ran.deny !== undefined) {
202      return { deny: mask(ran.deny, entries) }
203    }
204    if (!holdsSecret([ran.result, ran.text, ran.context], entries)) {
205      return ran
206    }
207
208    // Answering without `ref` makes core record and map our masked copy instead of its own.
209    const m = (s: string) => mask(s, entries)
210    const context = ran.context?.map(m)
211    if (ran.isError) {
212      return { isError: true, result: mapStrings(ran.result, m), text: ran.text === undefined ? undefined : m(ran.text), context }
213    }
214    return { result: mapStrings(ran.result, m) as typeof ran.result, context }
215  })
216
217  // Backstop for what the model reads and the transcript stores, whatever the row.
218  on('session.append', async ($, e, next) => {
219    const { entries } = await read($, vault)
220    if (entries.length === 0 || !holdsSecret(e.message.content, entries)) {
221      return next(e)
222    }
223    return next({ ...e, message: { ...e.message, content: e.message.content.map(b => maskBlock(b, entries)) } })
224  })
225
226  on('command.run', { command: 'vault' }, async ($, e) => {
227    const arg = e.args.trim()
228
229    if (arg === 'clear') {
230      await update($, vault, v => ({ entries: [], issued: v.issued }))
231      $.ui.status(undefined)
232      return { text: 'Vault cleared. Placeholders already in the conversation no longer resolve.' }
233    }
234    if (arg !== '' && arg !== 'list') {
235      return { text: 'Usage: /vault [list|clear]' }
236    }
237
238    const { entries } = await read($, vault)
239    if (entries.length === 0) {
240      return { text: 'The vault is empty.' }
241    }
242    return { text: entries.map(x => `${placeholder(x.n)}  ${x.kind.padEnd(12)}  ${hint(x.value)}`).join('\n') }
243  })
244}
245

1export type VaultEntry = { n: number; kind: string; value: string }
2
3// `issued` only grows, so a number is never reused after `/vault clear`.
4export type Vault = { entries: VaultEntry[]; issued: number }
5
6declare module 'claude-code' {
7  interface PluginState {
8    'secret-vault': { vault: Vault }
9  }
10}
11