my-pi/scripts/search-auth.mjs

#!/usr/bin/env node
/**
 * Read-modify-write the tavily/exa/keenable api_key credentials in Pi's
 * auth.json. Preserves every other provider entry (OAuth tokens included),
 * writes atomically with a lock directory, and keeps the file mode 600.
 *
 * Usage:
 *   node search-auth.mjs --set tavily=KEY [--set exa=KEY] [--set keenable=KEY]
 *   node search-auth.mjs --unset tavily [--unset exa]
 *   node search-auth.mjs --migrate [search.env path]
 *   node search-auth.mjs --list
 *   node search-auth.mjs --has tavily        (exit 0 when set, 1 otherwise)
 *
 * --auth FILE overrides the auth.json location (default:
 * $PI_CODING_AGENT_DIR/auth.json or ~/.pi/agent/auth.json).
 */

import {
	chmodSync,
	closeSync,
	existsSync,
	mkdirSync,
	openSync,
	readFileSync,
	renameSync,
	rmSync,
	writeFileSync,
} from "node:fs";
import { homedir } from "node:os";
import { dirname, join } from "node:path";
import { randomBytes } from "node:crypto";

const PROVIDERS = ["tavily", "exa", "keenable"];
const ENV_TO_PROVIDER = {
	TAVILY_API_KEY: "tavily",
	EXA_API_KEY: "exa",
	KEENABLE_API_KEY: "keenable",
};
const SAFE_VALUE = /^[A-Za-z0-9._-]+$/;
const LOCK_RETRIES = 50;
const LOCK_DELAY_MS = 100;

function die(message) {
	console.error(`错误: ${message}`);
	process.exit(1);
}

function usage() {
	console.log(`用法:
  node search-auth.mjs --set <provider>=<key> [--set ...]
  node search-auth.mjs --unset <provider> [--unset ...]
  node search-auth.mjs --migrate [search.env 路径]
  node search-auth.mjs --list
  node search-auth.mjs --has <provider>

选项:
  --auth FILE   覆盖 auth.json 路径（默认 $PI_CODING_AGENT_DIR/auth.json 或 ~/.pi/agent/auth.json）`);
}

function defaultAuthPath(env = process.env) {
	const agentDir = env.PI_CODING_AGENT_DIR?.trim() || join(env.HOME?.trim() || homedir(), ".pi", "agent");
	return join(agentDir, "auth.json");
}

function defaultSearchEnvPath(env = process.env) {
	if (env.MY_PI_SEARCH_CONFIG?.trim()) return env.MY_PI_SEARCH_CONFIG.trim();
	const configHome = env.XDG_CONFIG_HOME?.trim() || join(env.HOME?.trim() || homedir(), ".config");
	return join(configHome, "my-pi", "search.env");
}

function parseArgs(argv) {
	const result = { sets: {}, unsets: new Set(), migrate: undefined, list: false, has: undefined, authPath: undefined };
	for (let i = 0; i < argv.length; i++) {
		const arg = argv[i];
		if (arg === "--set") {
			const spec = argv[++i];
			const separator = spec ? spec.indexOf("=") : -1;
			if (separator < 1) die("--set 需要 <provider>=<key>");
			const provider = spec.slice(0, separator);
			const value = spec.slice(separator + 1);
			if (!PROVIDERS.includes(provider)) die(`未知 provider：${provider}（支持 ${PROVIDERS.join("/")}）`);
			result.sets[provider] = value;
		} else if (arg === "--unset") {
			const provider = argv[++i];
			if (!PROVIDERS.includes(provider)) die(`未知 provider：${provider}（支持 ${PROVIDERS.join("/")}）`);
			result.unsets.add(provider);
		} else if (arg === "--migrate") {
			const next = argv[i + 1];
			result.migrate = next && !next.startsWith("--") ? (i++, next) : true;
		} else if (arg === "--list") {
			result.list = true;
		} else if (arg === "--has") {
			const provider = argv[++i];
			if (!PROVIDERS.includes(provider)) die(`未知 provider：${provider}（支持 ${PROVIDERS.join("/")}）`);
			result.has = provider;
		} else if (arg === "--auth") {
			result.authPath = argv[++i];
			if (!result.authPath) die("--auth 缺少 FILE");
		} else if (arg === "-h" || arg === "--help") {
			usage();
			process.exit(0);
		} else {
			die(`未知参数：${arg}`);
		}
	}
	return result;
}

function acquireLock(authPath) {
	mkdirSync(dirname(authPath), { recursive: true });
	const lockDir = `${authPath}.lock`;
	for (let attempt = 0; attempt < LOCK_RETRIES; attempt++) {
		try {
			mkdirSync(lockDir);
			return () => {
				try {
					rmSync(lockDir, { recursive: true });
				} catch {}
			};
		} catch (error) {
			if (error.code !== "EEXIST") throw error;
			Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, LOCK_DELAY_MS);
		}
	}
	die(`无法获取 ${lockDir} 锁；可能有其他进程正在写 auth.json（残留锁目录可手动删除）。`);
}

function readAuthData(authPath) {
	let raw;
	try {
		raw = readFileSync(authPath, "utf8");
	} catch (error) {
		if (error.code === "ENOENT") return {};
		throw error;
	}
	let data;
	try {
		data = JSON.parse(raw);
	} catch {
		die(`auth.json 不是合法 JSON，已中止以免破坏现有凭据：${authPath}`);
	}
	if (!data || typeof data !== "object" || Array.isArray(data)) {
		die(`auth.json 顶层必须是对象，已中止以免破坏现有凭据：${authPath}`);
	}
	return data;
}

function writeAuthData(authPath, data) {
	mkdirSync(dirname(authPath), { recursive: true });
	const temp = join(dirname(authPath), `.auth.json.${process.pid}.${randomBytes(6).toString("hex")}.tmp`);
	try {
		const fd = openSync(temp, "w", 0o600);
		writeFileSync(fd, `${JSON.stringify(data, null, 2)}\n`);
		closeSync(fd);
		renameSync(temp, authPath);
		chmodSync(authPath, 0o600);
	} finally {
		rmSync(temp, { force: true });
	}
}

function hasKey(data, provider) {
	const entry = data[provider];
	return Boolean(entry && typeof entry === "object" && !Array.isArray(entry) && entry.type === "api_key" && entry.key);
}

function setKey(data, provider, key) {
	const existing = data[provider];
	const preserved =
		existing && typeof existing === "object" && !Array.isArray(existing) && existing.type === "api_key"
			? existing
			: {};
	data[provider] = { ...preserved, type: "api_key", key };
}

function parseSearchEnv(text) {
	const parsed = {};
	for (const rawLine of text.split(/\r?\n/u)) {
		const line = rawLine.trim().replace(/^export[ \t]+/u, "");
		if (!line || line.startsWith("#")) continue;
		const separator = line.indexOf("=");
		if (separator < 1) continue;
		const name = line.slice(0, separator).trim();
		const value = line.slice(separator + 1).trim();
		const provider = ENV_TO_PROVIDER[name];
		if (!provider || !SAFE_VALUE.test(value)) continue;
		parsed[provider] = value;
	}
	return parsed;
}

function validateKey(provider, key) {
	if (!key || /[\r\n]/u.test(key)) die(`${provider} 的 key 不能为空且不能包含换行。`);
}

const args = parseArgs(process.argv.slice(2));
const authPath = args.authPath ?? defaultAuthPath();

if (args.list) {
	const data = readAuthData(authPath);
	for (const provider of PROVIDERS) {
		console.log(`${provider}: ${hasKey(data, provider) ? "已配置" : "未配置"}`);
	}
	process.exit(0);
}

if (args.has) {
	process.exit(hasKey(readAuthData(authPath), args.has) ? 0 : 1);
}

const mutations = Object.keys(args.sets).length + args.unsets.size + (args.migrate ? 1 : 0);
if (mutations === 0) {
	usage();
	process.exit(2);
}

const release = acquireLock(authPath);
try {
	const data = readAuthData(authPath);
	const changed = [];

	for (const provider of args.unsets) {
		if (provider in data) {
			delete data[provider];
			changed.push(`已移除 ${provider}`);
		}
	}
	for (const [provider, key] of Object.entries(args.sets)) {
		validateKey(provider, key);
		setKey(data, provider, key);
		changed.push(`已写入 ${provider}`);
	}
	if (args.migrate) {
		const source = args.migrate === true ? defaultSearchEnvPath() : args.migrate;
		if (!existsSync(source)) die(`迁移来源不存在：${source}`);
		const parsed = parseSearchEnv(readFileSync(source, "utf8"));
		const imported = [];
		const skipped = [];
		for (const provider of PROVIDERS) {
			if (!parsed[provider]) continue;
			if (provider in data) {
				skipped.push(provider);
				continue;
			}
			setKey(data, provider, parsed[provider]);
			imported.push(provider);
		}
		if (imported.length) changed.push(`已迁移 ${imported.join(", ")}`);
		if (skipped.length) changed.push(`已跳过（auth.json 中已存在）：${skipped.join(", ")}`);
		if (!imported.length && !skipped.length) changed.push("search.env 中没有可迁移的 key");
	}

	if (changed.length === 0) {
		console.log("auth.json 无需修改。");
	} else {
		writeAuthData(authPath, data);
		for (const line of changed) console.log(line);
		console.log(`搜索凭据已写入 ${authPath}（权限 600），重启 Pi 后生效。`);
	}
} finally {
	release();
}