|
| 1 | +import { existsSync, mkdirSync, readFileSync, writeFileSync } from "fs"; |
| 2 | +import type { Command } from "commander"; |
| 3 | +import { getApiKey, loadCredentials, resolveProfileName } from "../lib/config"; |
| 4 | +import { appendAudit, buildContentAudit } from "../lib/audit"; |
| 5 | +import { printError, printInfo, printJson } from "../lib/output"; |
| 6 | +import type { ClientOptions } from "../lib/http"; |
| 7 | +import { scanInbound } from "../lib/safety"; |
| 8 | +import { sanitizeData, sanitizeText } from "../lib/unicode"; |
| 9 | +import { |
| 10 | + applyServerRetryAfter, |
| 11 | + checkRateLimit, |
| 12 | + extractRetryAfterSeconds, |
| 13 | +} from "../lib/rate_limit"; |
| 14 | +import { jailbreakDir, jailbreakRemotePath } from "../lib/paths"; |
| 15 | + |
| 16 | +type Globals = { |
| 17 | + profile?: string; |
| 18 | + baseUrl?: string; |
| 19 | + json?: boolean; |
| 20 | + timeout?: string | number; |
| 21 | + retries?: string | number; |
| 22 | + verbose?: boolean; |
| 23 | + dryRun?: boolean; |
| 24 | + quiet?: boolean; |
| 25 | + yes?: boolean; |
| 26 | + allowSensitive?: boolean; |
| 27 | + noColor?: boolean; |
| 28 | + wait?: boolean; |
| 29 | + maxWait?: string | number; |
| 30 | + [key: string]: unknown; |
| 31 | +}; |
| 32 | + |
| 33 | +type BuildClientResult = { |
| 34 | + client: ClientOptions; |
| 35 | + profileName: string; |
| 36 | + profile: unknown; |
| 37 | +}; |
| 38 | + |
| 39 | +export type LogOutboundParams = { |
| 40 | + profile: string; |
| 41 | + action: string; |
| 42 | + method: string; |
| 43 | + endpoint: string; |
| 44 | + status: "sent" | "blocked" | "dry_run"; |
| 45 | + reason?: string; |
| 46 | + content?: string; |
| 47 | + safety?: unknown[]; |
| 48 | + sanitization?: string[]; |
| 49 | + meta?: Record<string, unknown>; |
| 50 | +}; |
| 51 | + |
| 52 | +export type CommandContext = { |
| 53 | + program: Command; |
| 54 | + dmsEnabled: boolean; |
| 55 | + postReminder: string; |
| 56 | + globals: () => Globals; |
| 57 | + buildClient: (requireAuth: boolean) => BuildClientResult; |
| 58 | + redactProfileData: (profile: unknown) => unknown; |
| 59 | + sanitizeFields: (fields: Record<string, string | undefined>) => { |
| 60 | + sanitized: Record<string, string | undefined>; |
| 61 | + warnings: string[]; |
| 62 | + changed: boolean; |
| 63 | + }; |
| 64 | + warnSanitization: (warnings: string[], opts: Globals, context: string) => void; |
| 65 | + handleDryRun: ( |
| 66 | + res: { dryRun?: boolean; data?: unknown }, |
| 67 | + opts: Globals, |
| 68 | + extra?: Record<string, unknown>, |
| 69 | + ) => boolean; |
| 70 | + handleDmUnavailable: ( |
| 71 | + res: { ok: boolean; status: number; data?: unknown; error?: string }, |
| 72 | + opts: Globals, |
| 73 | + action: string, |
| 74 | + ) => boolean; |
| 75 | + sleep: (ms: number) => Promise<void>; |
| 76 | + saveJailbreakRemote: (url: string) => void; |
| 77 | + readJailbreakRemote: () => { url?: string }; |
| 78 | + enforceRateLimit: ( |
| 79 | + profile: string, |
| 80 | + action: "request" | "comment" | "post", |
| 81 | + opts: Globals, |
| 82 | + ) => Promise<void>; |
| 83 | + logOutbound: (params: LogOutboundParams) => void; |
| 84 | + applyRetryAfter: (profile: string, action: "request" | "comment" | "post", data: unknown) => void; |
| 85 | + attachInboundSafety: (data: unknown) => Promise<{ |
| 86 | + data: unknown; |
| 87 | + safety: unknown[]; |
| 88 | + sanitization: string[]; |
| 89 | + meta: { truncated: boolean; qmd: string; scanned_chars?: number; total_chars?: number }; |
| 90 | + }>; |
| 91 | +}; |
| 92 | + |
| 93 | +type ContextOptions = { |
| 94 | + dmsEnabled: boolean; |
| 95 | + postReminder: string; |
| 96 | +}; |
| 97 | + |
| 98 | +function collectStrings(value: unknown, acc: string[] = [], depth = 0): string[] { |
| 99 | + if (depth > 4) return acc; |
| 100 | + if (typeof value === "string") { |
| 101 | + acc.push(value); |
| 102 | + return acc; |
| 103 | + } |
| 104 | + if (Array.isArray(value)) { |
| 105 | + for (const item of value) collectStrings(item, acc, depth + 1); |
| 106 | + } else if (value && typeof value === "object") { |
| 107 | + for (const val of Object.values(value)) collectStrings(val, acc, depth + 1); |
| 108 | + } |
| 109 | + return acc; |
| 110 | +} |
| 111 | + |
| 112 | +export function createCommandContext(program: Command, options: ContextOptions): CommandContext { |
| 113 | + const globals = () => program.opts(); |
| 114 | + |
| 115 | + const buildClient = (requireAuth: boolean): BuildClientResult => { |
| 116 | + const opts = globals(); |
| 117 | + const profileName = resolveProfileName(opts.profile); |
| 118 | + const store = loadCredentials(); |
| 119 | + const apiKey = getApiKey(store, profileName); |
| 120 | + |
| 121 | + if (requireAuth && !apiKey) { |
| 122 | + printError(`No API key found for profile '${profileName}'. Run 'mb register' first.`, opts); |
| 123 | + process.exit(1); |
| 124 | + } |
| 125 | + |
| 126 | + const client: ClientOptions = { |
| 127 | + baseUrl: opts.baseUrl, |
| 128 | + apiKey: apiKey, |
| 129 | + timeoutMs: Math.max(1, Number(opts.timeout)) * 1000, |
| 130 | + retries: Math.max(0, Number(opts.retries)), |
| 131 | + verbose: !!opts.verbose, |
| 132 | + dryRun: !!opts.dryRun, |
| 133 | + }; |
| 134 | + |
| 135 | + return { client, profileName, profile: store[profileName] }; |
| 136 | + }; |
| 137 | + |
| 138 | + const redactProfileData = (profile: unknown): unknown => { |
| 139 | + if (!profile || typeof profile !== "object") return profile; |
| 140 | + const record = { ...(profile as Record<string, unknown>) }; |
| 141 | + if (typeof record.api_key === "string") { |
| 142 | + record.api_key = "[redacted]"; |
| 143 | + } |
| 144 | + return record; |
| 145 | + }; |
| 146 | + |
| 147 | + const sanitizeFields = (fields: Record<string, string | undefined>) => { |
| 148 | + const warnings = new Set<string>(); |
| 149 | + let changed = false; |
| 150 | + const sanitized: Record<string, string | undefined> = {}; |
| 151 | + |
| 152 | + for (const [key, value] of Object.entries(fields)) { |
| 153 | + if (typeof value !== "string") { |
| 154 | + sanitized[key] = value; |
| 155 | + continue; |
| 156 | + } |
| 157 | + const result = sanitizeText(value); |
| 158 | + result.warnings.forEach((warning) => warnings.add(warning)); |
| 159 | + if (result.changed) changed = true; |
| 160 | + sanitized[key] = result.text; |
| 161 | + } |
| 162 | + |
| 163 | + return { sanitized, warnings: Array.from(warnings), changed }; |
| 164 | + }; |
| 165 | + |
| 166 | + const warnSanitization = (warnings: string[], opts: Globals, context: string) => { |
| 167 | + if (warnings.length === 0) return; |
| 168 | + printInfo(`Warning: ${context}: ${warnings.join("; ")}`, opts); |
| 169 | + }; |
| 170 | + |
| 171 | + const handleDryRun = ( |
| 172 | + res: { dryRun?: boolean; data?: unknown }, |
| 173 | + opts: Globals, |
| 174 | + extra: Record<string, unknown> = {}, |
| 175 | + ) => { |
| 176 | + if (!res.dryRun) return false; |
| 177 | + if (opts.json) { |
| 178 | + printJson({ ...extra, result: res.data, dry_run: true }); |
| 179 | + return true; |
| 180 | + } |
| 181 | + printInfo("Dry run: request not sent.", opts); |
| 182 | + if (res.data) { |
| 183 | + printInfo(JSON.stringify(res.data, null, 2), opts); |
| 184 | + } |
| 185 | + return true; |
| 186 | + }; |
| 187 | + |
| 188 | + const handleDmUnavailable = ( |
| 189 | + res: { ok: boolean; status: number; data?: unknown; error?: string }, |
| 190 | + opts: Globals, |
| 191 | + action: string, |
| 192 | + ): boolean => { |
| 193 | + if (res.ok || res.status !== 404) return false; |
| 194 | + const message = "DM API unavailable on this server (404)."; |
| 195 | + if (opts.json) { |
| 196 | + printJson({ |
| 197 | + unavailable: true, |
| 198 | + action, |
| 199 | + status: res.status, |
| 200 | + error: res.error || res.data || message, |
| 201 | + }); |
| 202 | + } else { |
| 203 | + printInfo(message, opts); |
| 204 | + } |
| 205 | + return true; |
| 206 | + }; |
| 207 | + |
| 208 | + const sleep = (ms: number): Promise<void> => new Promise((resolve) => setTimeout(resolve, ms)); |
| 209 | + |
| 210 | + const saveJailbreakRemote = (url: string): void => { |
| 211 | + const dir = jailbreakDir(); |
| 212 | + if (!existsSync(dir)) mkdirSync(dir, { recursive: true }); |
| 213 | + const payload = { url, updated_at: new Date().toISOString() }; |
| 214 | + writeFileSync(jailbreakRemotePath(), JSON.stringify(payload, null, 2), { mode: 0o600 }); |
| 215 | + }; |
| 216 | + |
| 217 | + const readJailbreakRemote = (): { url?: string } => { |
| 218 | + const path = jailbreakRemotePath(); |
| 219 | + if (!existsSync(path)) return {}; |
| 220 | + try { |
| 221 | + const raw = readFileSync(path, "utf-8"); |
| 222 | + return JSON.parse(raw) as { url?: string }; |
| 223 | + } catch { |
| 224 | + return {}; |
| 225 | + } |
| 226 | + }; |
| 227 | + |
| 228 | + const enforceRateLimit = async ( |
| 229 | + profile: string, |
| 230 | + action: "request" | "comment" | "post", |
| 231 | + opts: Globals, |
| 232 | + ): Promise<void> => { |
| 233 | + if (opts.dryRun) return; |
| 234 | + const maxWait = Math.max(1, Number(opts.maxWait)) * 1000; |
| 235 | + while (true) { |
| 236 | + const decision = checkRateLimit(profile, action); |
| 237 | + if (decision.allowed) return; |
| 238 | + if (!opts.wait) { |
| 239 | + printError( |
| 240 | + `Rate limit hit: ${decision.reason}. Retry after ${Math.ceil(decision.waitMs / 1000)}s.`, |
| 241 | + opts, |
| 242 | + ); |
| 243 | + throw new Error("rate_limit"); |
| 244 | + } |
| 245 | + if (decision.waitMs > maxWait) { |
| 246 | + printError(`Rate limit wait exceeds max (${Math.ceil(decision.waitMs / 1000)}s).`, opts); |
| 247 | + throw new Error("rate_limit"); |
| 248 | + } |
| 249 | + printInfo( |
| 250 | + `Rate limited (${decision.reason}). Waiting ${Math.ceil(decision.waitMs / 1000)}s...`, |
| 251 | + opts, |
| 252 | + ); |
| 253 | + await sleep(decision.waitMs); |
| 254 | + } |
| 255 | + }; |
| 256 | + |
| 257 | + const logOutbound = (params: LogOutboundParams): void => { |
| 258 | + const contentAudit = buildContentAudit(params.content); |
| 259 | + appendAudit({ |
| 260 | + timestamp: new Date().toISOString(), |
| 261 | + profile: params.profile, |
| 262 | + action: params.action, |
| 263 | + method: params.method, |
| 264 | + endpoint: params.endpoint, |
| 265 | + status: params.status, |
| 266 | + reason: params.reason, |
| 267 | + safety_matches: params.safety, |
| 268 | + sanitization: params.sanitization, |
| 269 | + ...contentAudit, |
| 270 | + meta: params.meta, |
| 271 | + }); |
| 272 | + }; |
| 273 | + |
| 274 | + const applyRetryAfter = (profile: string, action: "request" | "comment" | "post", data: unknown) => { |
| 275 | + const retry = extractRetryAfterSeconds(data); |
| 276 | + if (retry && retry > 0) { |
| 277 | + applyServerRetryAfter(profile, action, retry); |
| 278 | + } |
| 279 | + }; |
| 280 | + |
| 281 | + const attachInboundSafety = async (data: unknown) => { |
| 282 | + const sanitized = sanitizeData(data); |
| 283 | + const strings = collectStrings(sanitized.value); |
| 284 | + if (strings.length === 0) { |
| 285 | + return { |
| 286 | + data: sanitized.value, |
| 287 | + safety: [] as unknown[], |
| 288 | + sanitization: sanitized.warnings, |
| 289 | + meta: { truncated: false, qmd: "skipped" }, |
| 290 | + }; |
| 291 | + } |
| 292 | + const combined = strings.join("\n"); |
| 293 | + const maxChars = Math.max(2000, Number(process.env.MB_INBOUND_MAX_CHARS ?? "20000")); |
| 294 | + const maxQmdChars = Math.max(2000, Number(process.env.MB_INBOUND_QMD_MAX_CHARS ?? "8000")); |
| 295 | + const truncated = combined.length > maxChars; |
| 296 | + const sample = truncated ? combined.slice(0, maxChars) : combined; |
| 297 | + const useQmd = sample.length <= maxQmdChars; |
| 298 | + const matches = await scanInbound(sample, { useQmd }); |
| 299 | + const warnings = [...sanitized.warnings]; |
| 300 | + if (truncated) warnings.push("Inbound safety scan truncated (large payload)"); |
| 301 | + if (!useQmd) warnings.push("Inbound safety qmd scan skipped (large payload)"); |
| 302 | + return { |
| 303 | + data: sanitized.value, |
| 304 | + safety: matches, |
| 305 | + sanitization: warnings, |
| 306 | + meta: { |
| 307 | + truncated, |
| 308 | + qmd: useQmd ? "used" : "skipped", |
| 309 | + scanned_chars: sample.length, |
| 310 | + total_chars: combined.length, |
| 311 | + }, |
| 312 | + }; |
| 313 | + }; |
| 314 | + |
| 315 | + return { |
| 316 | + program, |
| 317 | + dmsEnabled: options.dmsEnabled, |
| 318 | + postReminder: options.postReminder, |
| 319 | + globals, |
| 320 | + buildClient, |
| 321 | + redactProfileData, |
| 322 | + sanitizeFields, |
| 323 | + warnSanitization, |
| 324 | + handleDryRun, |
| 325 | + handleDmUnavailable, |
| 326 | + sleep, |
| 327 | + saveJailbreakRemote, |
| 328 | + readJailbreakRemote, |
| 329 | + enforceRateLimit, |
| 330 | + logOutbound, |
| 331 | + applyRetryAfter, |
| 332 | + attachInboundSafety, |
| 333 | + }; |
| 334 | +} |
0 commit comments