@yivan-lab/pretty-please 1.0.0 → 1.2.0

package/dist/src/remote.js

@@ -0,0 +1,634 @@
+/**
+ * 远程执行器模块
+ * 通过 SSH 在远程服务器上执行命令
+ */
+import { spawn } from 'child_process';
+import fs from 'fs';
+import path from 'path';
+import os from 'os';
+import readline from 'readline';
+import chalk from 'chalk';
+import { CONFIG_DIR, getConfig, saveConfig } from './config.js';
+import { getCurrentTheme } from './ui/theme.js';
+// 获取主题颜色
+function getColors() {
+    const theme = getCurrentTheme();
+    return {
+        primary: theme.primary,
+        secondary: theme.secondary,
+        success: theme.success,
+        error: theme.error,
+        warning: theme.warning,
+        muted: theme.text.muted,
+    };
+}
+// 远程服务器数据目录
+const REMOTES_DIR = path.join(CONFIG_DIR, 'remotes');
+// SSH ControlMaster 配置
+const SSH_CONTROL_PERSIST = '10m'; // 连接保持 10 分钟
+/**
+ * 确保远程服务器数据目录存在
+ */
+function ensureRemotesDir() {
+    if (!fs.existsSync(REMOTES_DIR)) {
+        fs.mkdirSync(REMOTES_DIR, { recursive: true });
+    }
+}
+/**
+ * 获取远程服务器数据目录
+ */
+function getRemoteDataDir(name) {
+    return path.join(REMOTES_DIR, name);
+}
+/**
+ * 获取 SSH ControlMaster socket 路径
+ */
+function getSSHSocketPath(name) {
+    return path.join(REMOTES_DIR, name, 'ssh.sock');
+}
+/**
+ * 检查 ControlMaster 连接是否存在
+ */
+function isControlMasterActive(name) {
+    const socketPath = getSSHSocketPath(name);
+    return fs.existsSync(socketPath);
+}
+/**
+ * 关闭 ControlMaster 连接
+ */
+export async function closeControlMaster(name) {
+    const remote = getRemote(name);
+    if (!remote)
+        return;
+    const socketPath = getSSHSocketPath(name);
+    if (!fs.existsSync(socketPath))
+        return;
+    // 使用 ssh -O exit 关闭 master 连接
+    const args = ['-O', 'exit', '-o', `ControlPath=${socketPath}`, `${remote.user}@${remote.host}`];
+    return new Promise((resolve) => {
+        const child = spawn('ssh', args, { stdio: 'ignore' });
+        child.on('close', () => {
+            // 确保 socket 文件被删除
+            if (fs.existsSync(socketPath)) {
+                try {
+                    fs.unlinkSync(socketPath);
+                }
+                catch {
+                    // 忽略错误
+                }
+            }
+            resolve();
+        });
+        child.on('error', () => resolve());
+    });
+}
+/**
+ * 确保远程服务器数据目录存在
+ */
+function ensureRemoteDataDir(name) {
+    const dir = getRemoteDataDir(name);
+    if (!fs.existsSync(dir)) {
+        fs.mkdirSync(dir, { recursive: true });
+    }
+}
+// ================== 远程服务器管理 ==================
+/**
+ * 获取所有远程服务器配置
+ */
+export function getRemotes() {
+    const config = getConfig();
+    return config.remotes || {};
+}
+/**
+ * 获取单个远程服务器配置
+ */
+export function getRemote(name) {
+    const remotes = getRemotes();
+    return remotes[name] || null;
+}
+/**
+ * 解析 user@host:port 格式
+ */
+function parseHostString(hostStr) {
+    let user = '';
+    let host = hostStr;
+    let port = 22;
+    // 解析 user@host
+    if (hostStr.includes('@')) {
+        const atIndex = hostStr.indexOf('@');
+        user = hostStr.substring(0, atIndex);
+        host = hostStr.substring(atIndex + 1);
+    }
+    // 解析 host:port
+    if (host.includes(':')) {
+        const colonIndex = host.lastIndexOf(':');
+        const portStr = host.substring(colonIndex + 1);
+        const parsedPort = parseInt(portStr, 10);
+        if (!isNaN(parsedPort)) {
+            port = parsedPort;
+            host = host.substring(0, colonIndex);
+        }
+    }
+    return { user, host, port };
+}
+/**
+ * 添加远程服务器
+ */
+export function addRemote(name, hostStr, options = {}) {
+    // 验证名称
+    if (!name || !name.trim()) {
+        throw new Error('服务器名称不能为空');
+    }
+    if (!/^[a-zA-Z0-9_-]+$/.test(name)) {
+        throw new Error('服务器名称只能包含字母、数字、下划线和连字符');
+    }
+    // 解析 host 字符串
+    const { user, host, port } = parseHostString(hostStr);
+    if (!host) {
+        throw new Error('主机地址不能为空');
+    }
+    if (!user) {
+        throw new Error('用户名不能为空，请使用 user@host 格式');
+    }
+    // 验证密钥文件
+    if (options.key) {
+        const keyPath = options.key.replace(/^~/, os.homedir());
+        if (!fs.existsSync(keyPath)) {
+            throw new Error(`密钥文件不存在: ${options.key}`);
+        }
+    }
+    const config = getConfig();
+    if (!config.remotes) {
+        config.remotes = {};
+    }
+    // 检查是否已存在
+    if (config.remotes[name]) {
+        throw new Error(`服务器 "${name}" 已存在，请使用其他名称或先删除`);
+    }
+    config.remotes[name] = {
+        host,
+        user,
+        port,
+        key: options.key,
+        password: options.password,
+    };
+    saveConfig(config);
+    // 创建数据目录
+    ensureRemoteDataDir(name);
+}
+/**
+ * 删除远程服务器
+ */
+export function removeRemote(name) {
+    const config = getConfig();
+    if (!config.remotes || !config.remotes[name]) {
+        return false;
+    }
+    delete config.remotes[name];
+    saveConfig(config);
+    // 删除数据目录
+    const dataDir = getRemoteDataDir(name);
+    if (fs.existsSync(dataDir)) {
+        fs.rmSync(dataDir, { recursive: true });
+    }
+    return true;
+}
+/**
+ * 显示所有远程服务器
+ */
+export function displayRemotes() {
+    const remotes = getRemotes();
+    const config = getConfig();
+    const colors = getColors();
+    const names = Object.keys(remotes);
+    console.log('');
+    if (names.length === 0) {
+        console.log(chalk.gray('  暂无远程服务器'));
+        console.log('');
+        console.log(chalk.gray('  使用 pls remote add <name> <user@host> 添加服务器'));
+        console.log('');
+        return;
+    }
+    console.log(chalk.bold('远程服务器:'));
+    console.log(chalk.gray('━'.repeat(60)));
+    for (const name of names) {
+        const remote = remotes[name];
+        const authType = remote.password ? '密码' : remote.key ? '密钥' : '默认密钥';
+        const isDefault = config.defaultRemote === name;
+        // 服务器名称，如果是默认则显示标记
+        if (isDefault) {
+            console.log(`  ${chalk.hex(colors.primary)(name)} ${chalk.hex(colors.success)('(default)')}`);
+        }
+        else {
+            console.log(`  ${chalk.hex(colors.primary)(name)}`);
+        }
+        console.log(`    ${chalk.gray('→')} ${remote.user}@${remote.host}:${remote.port}`);
+        console.log(`    ${chalk.gray('认证:')} ${authType}${remote.key ? ` (${remote.key})` : ''}`);
+        // 显示工作目录
+        if (remote.workDir) {
+            console.log(`    ${chalk.gray('工作目录:')} ${remote.workDir}`);
+        }
+        // 检查是否有缓存的系统信息
+        const sysInfo = getRemoteSysInfo(name);
+        if (sysInfo) {
+            console.log(`    ${chalk.gray('系统:')} ${sysInfo.os} ${sysInfo.osVersion} (${sysInfo.shell})`);
+        }
+        console.log('');
+    }
+    console.log(chalk.gray('━'.repeat(60)));
+    console.log(chalk.gray('使用: pls -r <name> <prompt> 在远程服务器执行'));
+    console.log('');
+}
+/**
+ * 设置远程服务器工作目录
+ */
+export function setRemoteWorkDir(name, workDir) {
+    const config = getConfig();
+    if (!config.remotes || !config.remotes[name]) {
+        throw new Error(`远程服务器 "${name}" 不存在`);
+    }
+    // 清除工作目录
+    if (!workDir || workDir === '-') {
+        delete config.remotes[name].workDir;
+    }
+    else {
+        config.remotes[name].workDir = workDir;
+    }
+    saveConfig(config);
+}
+/**
+ * 获取远程服务器工作目录
+ */
+export function getRemoteWorkDir(name) {
+    const remote = getRemote(name);
+    return remote?.workDir;
+}
+/**
+ * 读取密码（交互式）
+ */
+async function readPassword(prompt) {
+    return new Promise((resolve) => {
+        const rl = readline.createInterface({
+            input: process.stdin,
+            output: process.stdout,
+        });
+        // 隐藏输入
+        const stdin = process.stdin;
+        if (stdin.isTTY) {
+            stdin.setRawMode(true);
+        }
+        process.stdout.write(prompt);
+        let password = '';
+        stdin.on('data', (char) => {
+            const c = char.toString();
+            switch (c) {
+                case '\n':
+                case '\r':
+                case '\u0004': // Ctrl+D
+                    if (stdin.isTTY) {
+                        stdin.setRawMode(false);
+                    }
+                    console.log('');
+                    rl.close();
+                    resolve(password);
+                    break;
+                case '\u0003': // Ctrl+C
+                    if (stdin.isTTY) {
+                        stdin.setRawMode(false);
+                    }
+                    console.log('');
+                    rl.close();
+                    process.exit(0);
+                case '\u007F': // Backspace
+                    if (password.length > 0) {
+                        password = password.slice(0, -1);
+                        process.stdout.write('\b \b');
+                    }
+                    break;
+                default:
+                    password += c;
+                    process.stdout.write('*');
+                    break;
+            }
+        });
+    });
+}
+/**
+ * 构建 SSH 命令参数
+ * @param remote 远程服务器配置
+ * @param command 要执行的命令
+ * @param options.password 密码（用于首次建立连接）
+ * @param options.socketPath ControlMaster socket 路径
+ * @param options.isMaster 是否建立 master 连接
+ */
+function buildSSHArgs(remote, command, options = {}) {
+    const args = [];
+    // 使用 sshpass 处理密码认证（仅在建立新连接时需要）
+    let cmd = 'ssh';
+    if (options.password) {
+        cmd = 'sshpass';
+        args.push('-p', options.password, 'ssh');
+    }
+    // SSH 选项
+    args.push('-o', 'StrictHostKeyChecking=accept-new');
+    args.push('-o', 'ConnectTimeout=10');
+    // ControlMaster 选项
+    if (options.socketPath) {
+        if (options.isMaster) {
+            // 建立 master 连接
+            args.push('-o', 'ControlMaster=yes');
+            args.push('-o', `ControlPersist=${SSH_CONTROL_PERSIST}`);
+        }
+        else {
+            // 复用已有连接
+            args.push('-o', 'ControlMaster=no');
+        }
+        args.push('-o', `ControlPath=${options.socketPath}`);
+    }
+    // 端口
+    if (remote.port !== 22) {
+        args.push('-p', remote.port.toString());
+    }
+    // 密钥
+    if (remote.key) {
+        const keyPath = remote.key.replace(/^~/, os.homedir());
+        args.push('-i', keyPath);
+    }
+    // 目标
+    args.push(`${remote.user}@${remote.host}`);
+    // 命令
+    args.push(command);
+    return { cmd, args };
+}
+/**
+ * 执行 SSH 命令的内部实现
+ */
+function spawnSSH(cmd, args, options) {
+    return new Promise((resolve, reject) => {
+        let stdout = '';
+        let stderr = '';
+        const child = spawn(cmd, args, {
+            stdio: ['pipe', 'pipe', 'pipe'],
+        });
+        // 超时处理
+        let timeoutId = null;
+        if (options.timeout) {
+            timeoutId = setTimeout(() => {
+                child.kill('SIGTERM');
+                reject(new Error(`命令执行超时 (${options.timeout}ms)`));
+            }, options.timeout);
+        }
+        child.stdout.on('data', (data) => {
+            const str = data.toString();
+            stdout += str;
+            options.onStdout?.(str);
+        });
+        child.stderr.on('data', (data) => {
+            const str = data.toString();
+            stderr += str;
+            options.onStderr?.(str);
+        });
+        // 写入 stdin
+        if (options.stdin) {
+            child.stdin.write(options.stdin);
+            child.stdin.end();
+        }
+        child.on('close', (code) => {
+            if (timeoutId) {
+                clearTimeout(timeoutId);
+            }
+            resolve({
+                exitCode: code || 0,
+                stdout,
+                stderr,
+                output: stdout + stderr,
+            });
+        });
+        child.on('error', (err) => {
+            if (timeoutId) {
+                clearTimeout(timeoutId);
+            }
+            // 检查是否是 sshpass 未安装
+            if (err.message.includes('ENOENT') && cmd === 'sshpass') {
+                reject(new Error('密码认证需要安装 sshpass，请运行: brew install hudochenkov/sshpass/sshpass'));
+            }
+            else {
+                reject(err);
+            }
+        });
+    });
+}
+/**
+ * 通过 SSH 执行命令
+ * 使用 ControlMaster 实现连接复用，密码认证只需输入一次
+ */
+export async function sshExec(name, command, options = {}) {
+    const remote = getRemote(name);
+    if (!remote) {
+        throw new Error(`远程服务器 "${name}" 不存在`);
+    }
+    // 确保数据目录存在
+    ensureRemoteDataDir(name);
+    const socketPath = getSSHSocketPath(name);
+    const masterActive = isControlMasterActive(name);
+    // 如果需要密码认证且没有活跃的 master 连接
+    if (remote.password && !masterActive) {
+        // 读取密码并建立 master 连接
+        const password = await readPassword(`${name} 密码: `);
+        // 建立 master 连接（执行一个简单命令来建立连接）
+        const { cmd: masterCmd, args: masterArgs } = buildSSHArgs(remote, 'true', {
+            password,
+            socketPath,
+            isMaster: true,
+        });
+        try {
+            const masterResult = await spawnSSH(masterCmd, masterArgs, { timeout: 30000 });
+            if (masterResult.exitCode !== 0) {
+                throw new Error(`SSH 连接失败: ${masterResult.stderr}`);
+            }
+        }
+        catch (err) {
+            throw err;
+        }
+    }
+    // 使用 ControlMaster 连接（或直接连接）执行命令
+    const useSocket = remote.password || isControlMasterActive(name);
+    const { cmd, args } = buildSSHArgs(remote, command, {
+        socketPath: useSocket ? socketPath : undefined,
+        isMaster: false,
+    });
+    return spawnSSH(cmd, args, options);
+}
+/**
+ * 测试远程连接
+ */
+export async function testRemoteConnection(name) {
+    const colors = getColors();
+    try {
+        const result = await sshExec(name, 'echo "pls-connection-test"', { timeout: 15000 });
+        if (result.exitCode === 0 && result.stdout.includes('pls-connection-test')) {
+            return { success: true, message: chalk.hex(colors.success)('连接成功') };
+        }
+        else {
+            return { success: false, message: chalk.hex(colors.error)(`连接失败，退出码: ${result.exitCode}`) };
+        }
+    }
+    catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        return { success: false, message: chalk.hex(colors.error)(`连接失败: ${message}`) };
+    }
+}
+// ================== 系统信息采集 ==================
+/**
+ * 获取缓存的远程系统信息
+ */
+export function getRemoteSysInfo(name) {
+    const dataDir = getRemoteDataDir(name);
+    const sysInfoPath = path.join(dataDir, 'sysinfo.json');
+    if (!fs.existsSync(sysInfoPath)) {
+        return null;
+    }
+    try {
+        const content = fs.readFileSync(sysInfoPath, 'utf-8');
+        return JSON.parse(content);
+    }
+    catch {
+        return null;
+    }
+}
+/**
+ * 保存远程系统信息
+ */
+function saveRemoteSysInfo(name, sysInfo) {
+    ensureRemoteDataDir(name);
+    const dataDir = getRemoteDataDir(name);
+    const sysInfoPath = path.join(dataDir, 'sysinfo.json');
+    fs.writeFileSync(sysInfoPath, JSON.stringify(sysInfo, null, 2));
+}
+/**
+ * 采集远程系统信息
+ */
+export async function collectRemoteSysInfo(name, force = false) {
+    // 检查缓存
+    if (!force) {
+        const cached = getRemoteSysInfo(name);
+        if (cached) {
+            // 检查缓存是否过期（7天）
+            const cachedAt = new Date(cached.cachedAt);
+            const now = new Date();
+            const daysDiff = (now.getTime() - cachedAt.getTime()) / (1000 * 60 * 60 * 24);
+            if (daysDiff < 7) {
+                return cached;
+            }
+        }
+    }
+    // 采集系统信息
+    const collectScript = `
+echo "OS:$(uname -s)"
+echo "OS_VERSION:$(uname -r)"
+echo "SHELL:$(basename "$SHELL")"
+echo "HOSTNAME:$(hostname)"
+`.trim();
+    const result = await sshExec(name, collectScript, { timeout: 30000 });
+    if (result.exitCode !== 0) {
+        throw new Error(`无法采集系统信息: ${result.stderr}`);
+    }
+    // 解析输出
+    const lines = result.stdout.split('\n');
+    const info = {};
+    for (const line of lines) {
+        const colonIndex = line.indexOf(':');
+        if (colonIndex > 0) {
+            const key = line.substring(0, colonIndex).trim();
+            const value = line.substring(colonIndex + 1).trim();
+            info[key] = value;
+        }
+    }
+    const sysInfo = {
+        os: info['OS'] || 'unknown',
+        osVersion: info['OS_VERSION'] || 'unknown',
+        shell: info['SHELL'] || 'bash',
+        hostname: info['HOSTNAME'] || 'unknown',
+        cachedAt: new Date().toISOString(),
+    };
+    // 保存缓存
+    saveRemoteSysInfo(name, sysInfo);
+    return sysInfo;
+}
+/**
+ * 格式化远程系统信息供 AI 使用
+ */
+export function formatRemoteSysInfoForAI(name, sysInfo) {
+    const remote = getRemote(name);
+    if (!remote)
+        return '';
+    let info = `【远程服务器信息】
+服务器: ${name} (${remote.user}@${remote.host})
+操作系统: ${sysInfo.os} ${sysInfo.osVersion}
+Shell: ${sysInfo.shell}
+主机名: ${sysInfo.hostname}`;
+    // 如果有工作目录，告知 AI 当前工作目录（执行时会自动 cd）
+    if (remote.workDir) {
+        info += `\n当前工作目录: ${remote.workDir}`;
+    }
+    return info;
+}
+/**
+ * 批量远程执行命令
+ * 每个服务器单独生成命令，支持异构环境
+ */
+export async function generateBatchRemoteCommands(serverNames, userPrompt, options = {}) {
+    const { generateMultiStepCommand } = await import('./multi-step.js');
+    const { fetchRemoteShellHistory } = await import('./remote-history.js');
+    // 1. 验证所有服务器是否存在
+    const invalidServers = serverNames.filter(name => !getRemote(name));
+    if (invalidServers.length > 0) {
+        throw new Error(`以下服务器不存在: ${invalidServers.join(', ')}`);
+    }
+    // 2. 并发采集所有服务器的系统信息
+    const servers = await Promise.all(serverNames.map(async (name) => ({
+        name,
+        sysInfo: await collectRemoteSysInfo(name),
+        shellHistory: await fetchRemoteShellHistory(name),
+    })));
+    // 3. 并发为每个服务器生成命令
+    const commandResults = await Promise.all(servers.map(async (server) => {
+        const remoteContext = {
+            name: server.name,
+            sysInfo: server.sysInfo,
+            shellHistory: server.shellHistory,
+        };
+        const result = await generateMultiStepCommand(userPrompt, [], // 批量执行不支持多步骤，只生成单个命令
+        { debug: options.debug, remoteContext });
+        return {
+            server: server.name,
+            command: result.stepData.command,
+            sysInfo: server.sysInfo,
+        };
+    }));
+    return commandResults;
+}
+/**
+ * 执行批量远程命令
+ */
+export async function executeBatchRemoteCommands(commands) {
+    // 并发执行所有命令
+    const results = await Promise.all(commands.map(async ({ server, command, sysInfo }) => {
+        let stdout = '';
+        let stderr = '';
+        const result = await sshExec(server, command, {
+            onStdout: (data) => { stdout += data; },
+            onStderr: (data) => { stderr += data; },
+        });
+        return {
+            server,
+            command,
+            exitCode: result.exitCode,
+            stdout,
+            stderr,
+            output: stdout + stderr,
+            sysInfo,
+        };
+    }));
+    return results;
+}