diff --git a/pkg/agent/config/config.go b/pkg/agent/config/config.go index c46301e..95a1ec1 100644 --- a/pkg/agent/config/config.go +++ b/pkg/agent/config/config.go @@ -123,8 +123,12 @@ func DefaultConfig() *Config { // GetDefaultConfigPath 获取默认配置文件路径 func GetDefaultConfigPath() string { - var homeDir = utils.GetSafeHomeDir() - return filepath.Join(homeDir, ".pika", "agent.yaml") + return filepath.Join(GetDataDir(), "agent.yaml") +} + +// GetDataDir 获取 Agent 数据目录 +func GetDataDir() string { + return filepath.Join(utils.GetSafeHomeDir(), ".pika") } // Load 加载配置文件 diff --git a/pkg/agent/service/service.go b/pkg/agent/service/service.go index 96bfc07..2230f79 100644 --- a/pkg/agent/service/service.go +++ b/pkg/agent/service/service.go @@ -2,17 +2,18 @@ package service import ( "context" + "errors" "fmt" "log/slog" "os" "os/signal" + "path/filepath" "runtime" "strings" "syscall" "github.com/dushixiang/pika/pkg/agent" "github.com/dushixiang/pika/pkg/agent/config" - "github.com/dushixiang/pika/pkg/agent/id" "github.com/dushixiang/pika/pkg/agent/sshmonitor" "github.com/dushixiang/pika/pkg/agent/sysutil" "github.com/dushixiang/pika/pkg/agent/updater" @@ -353,19 +354,41 @@ func UninstallAgent(cfgPath string) error { // 清理 SSH 监控配置 monitor := sshmonitor.NewMonitor() - if err := monitor.Stop(); err != nil { - slog.Warn("清理SSH监控配置失败", "error", err) + return cleanupAgentArtifacts(cfgPath, monitor.Uninstall) +} + +func cleanupAgentArtifacts(cfgPath string, uninstallSSHMonitor func() error) error { + var cleanupErr error + + if err := uninstallSSHMonitor(); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("uninstall SSH monitor failed: %w", err)) } - // 删除配置文件 - if err := os.Remove(cfgPath); err != nil { - slog.Warn("删除配置文件失败", "error", err) + // SSH 清理失败不应阻止配置、ID、日志和历史缓存的删除。 + if err := removeAgentData(cfgPath); err != nil { + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("remove agent data failed: %w", err)) } - // 删除探针 ID 文件 - idPath := id.GetIDFilePath() - if err := os.Remove(idPath); err != nil { - slog.Warn("删除探针 ID 文件失败", "error", err) + return cleanupErr +} + +func removeAgentData(cfgPath string) error { + if cfgPath == "" { + cfgPath = config.GetDefaultConfigPath() + } + + dataDir := filepath.Clean(config.GetDataDir()) + cleanCfgPath := filepath.Clean(cfgPath) + + // 自定义配置可能位于数据目录之外,只删除配置文件本身。 + if cleanCfgPath != dataDir && !strings.HasPrefix(cleanCfgPath, dataDir+string(os.PathSeparator)) { + if err := os.Remove(cleanCfgPath); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("remove config file %s: %w", cleanCfgPath, err) + } + } + + if err := os.RemoveAll(dataDir); err != nil { + return fmt.Errorf("remove data directory %s: %w", dataDir, err) } return nil diff --git a/pkg/agent/service/uninstall_test.go b/pkg/agent/service/uninstall_test.go new file mode 100644 index 0000000..6484a27 --- /dev/null +++ b/pkg/agent/service/uninstall_test.go @@ -0,0 +1,100 @@ +package service + +import ( + "errors" + "os" + "path/filepath" + "testing" + + "github.com/dushixiang/pika/pkg/agent/config" +) + +func TestRemoveAgentDataDefaultConfig(t *testing.T) { + homeDir := t.TempDir() + t.Setenv("HOME", homeDir) + + dataDir := config.GetDataDir() + configPath := config.GetDefaultConfigPath() + writeTestFile(t, configPath) + writeTestFile(t, filepath.Join(dataDir, "agent.id")) + writeTestFile(t, filepath.Join(dataDir, "logs", "agent.log")) + writeTestFile(t, filepath.Join(dataDir, "metrics_buffer.db")) + + if err := removeAgentData(configPath); err != nil { + t.Fatalf("removeAgentData() error = %v", err) + } + + if _, err := os.Stat(dataDir); !os.IsNotExist(err) { + t.Fatalf("data directory still exists after uninstall: %v", err) + } +} + +func TestRemoveAgentDataEmptyConfigPath(t *testing.T) { + homeDir := t.TempDir() + t.Setenv("HOME", homeDir) + + writeTestFile(t, config.GetDefaultConfigPath()) + + if err := removeAgentData(""); err != nil { + t.Fatalf("removeAgentData() error = %v", err) + } + if _, err := os.Stat(config.GetDataDir()); !os.IsNotExist(err) { + t.Fatalf("data directory still exists after uninstall: %v", err) + } +} + +func TestRemoveAgentDataCustomConfig(t *testing.T) { + homeDir := t.TempDir() + t.Setenv("HOME", homeDir) + + customDir := t.TempDir() + configPath := filepath.Join(customDir, "agent.yaml") + otherPath := filepath.Join(customDir, "keep.txt") + writeTestFile(t, configPath) + writeTestFile(t, otherPath) + writeTestFile(t, filepath.Join(config.GetDataDir(), "agent.id")) + + if err := removeAgentData(configPath); err != nil { + t.Fatalf("removeAgentData() error = %v", err) + } + + if _, err := os.Stat(configPath); !os.IsNotExist(err) { + t.Fatalf("custom config still exists after uninstall: %v", err) + } + if _, err := os.Stat(config.GetDataDir()); !os.IsNotExist(err) { + t.Fatalf("data directory still exists after uninstall: %v", err) + } + if _, err := os.Stat(otherPath); err != nil { + t.Fatalf("unrelated file was removed: %v", err) + } +} + +func TestCleanupAgentArtifactsRemovesDataWhenSSHCleanupFails(t *testing.T) { + homeDir := t.TempDir() + t.Setenv("HOME", homeDir) + + dataDir := config.GetDataDir() + writeTestFile(t, config.GetDefaultConfigPath()) + writeTestFile(t, filepath.Join(dataDir, "logs", "agent.log")) + + sshErr := errors.New("SSH cleanup failed") + err := cleanupAgentArtifacts(config.GetDefaultConfigPath(), func() error { + return sshErr + }) + if !errors.Is(err, sshErr) { + t.Fatalf("cleanupAgentArtifacts() error = %v, want wrapped SSH error", err) + } + if _, err := os.Stat(dataDir); !os.IsNotExist(err) { + t.Fatalf("data directory still exists after SSH cleanup failure: %v", err) + } +} + +func writeTestFile(t *testing.T, path string) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil { + t.Fatalf("create parent directory: %v", err) + } + if err := os.WriteFile(path, []byte("test"), 0644); err != nil { + t.Fatalf("write test file: %v", err) + } +} diff --git a/pkg/agent/sshmonitor/hook_manager.go b/pkg/agent/sshmonitor/hook_manager.go index 48cbfcc..437dc93 100644 --- a/pkg/agent/sshmonitor/hook_manager.go +++ b/pkg/agent/sshmonitor/hook_manager.go @@ -3,7 +3,6 @@ package sshmonitor import ( "bufio" "fmt" - "io" "log/slog" "os" "os/exec" @@ -53,15 +52,15 @@ func (h *HookManager) Install() error { return fmt.Errorf("配置 SSH UsePAM 失败: %w", err) } - // 安装 Hook 可执行文件(软链) + // 安装 Hook 可执行文件(软链接) if err := h.ensureHookBinary(); err != nil { return fmt.Errorf("安装 Hook 可执行文件失败: %w", err) } // 修改 PAM 配置 if err := h.modifyPAMConfig(true); err != nil { - // 回滚软链 - h.removeHookBinary() + // 回滚 Hook 可执行文件 + _ = h.removeHookBinary() return fmt.Errorf("修改 PAM 配置失败: %w", err) } @@ -71,36 +70,56 @@ func (h *HookManager) Install() error { // Uninstall 卸载 PAM Hook func (h *HookManager) Uninstall() error { - // 移除 PAM 配置 - if err := h.modifyPAMConfig(false); err != nil { - slog.Warn("移除 PAM 配置失败", "error", err) + // 仅在配置确实存在时改写 PAM 文件,避免未启用过 SSH 监控时产生备份等副作用。 + installed, legacyInstalled, err := h.installedConfigState() + if err != nil { + return fmt.Errorf("检查 PAM 配置失败: %w", err) + } + if installed || legacyInstalled { + if err := h.modifyPAMConfig(false); err != nil { + return fmt.Errorf("移除 PAM 配置失败: %w", err) + } } - // 删除软链 - h.removeHookBinary() + if err := h.removeHookBinary(); err != nil { + return fmt.Errorf("移除 PAM Hook 可执行文件失败: %w", err) + } slog.Info("PAM Hook 卸载成功") return nil } -// isInstalled 检查是否已安装 -func (h *HookManager) isInstalled() bool { - // 检查 PAM 配置文件是否包含我们的配置行 +func (h *HookManager) installedConfigState() (installed, legacyInstalled bool, err error) { f, err := os.Open(PAMConfigFile) if err != nil { - return false + if os.IsNotExist(err) { + return false, false, nil + } + return false, false, err } defer f.Close() scanner := bufio.NewScanner(f) for scanner.Scan() { line := strings.TrimSpace(scanner.Text()) - if line == pamConfigLine() { - return true + switch line { + case pamConfigLine(): + installed = true + case legacyPamConfigLine(): + legacyInstalled = true } } - return false + if err := scanner.Err(); err != nil { + return false, false, err + } + return installed, legacyInstalled, nil +} + +// isInstalled 检查是否已安装 +func (h *HookManager) isInstalled() bool { + installed, _, _ := h.installedConfigState() + return installed } // ensureHookBinary 确保 Hook 可执行文件存在 @@ -123,46 +142,41 @@ func (h *HookManager) ensureHookBinary() error { if target, err := os.Readlink(HookBinaryPath); err == nil && target == execPath { return nil } + if err := os.Remove(HookBinaryPath); err != nil { + return err + } + } else { + return fmt.Errorf("Hook 路径已存在普通文件,拒绝覆盖: %s", HookBinaryPath) } + } else if !os.IsNotExist(err) { + return err } - _ = os.Remove(HookBinaryPath) - - if err := os.Symlink(execPath, HookBinaryPath); err == nil { - return nil + if err := os.Symlink(execPath, HookBinaryPath); err != nil { + return fmt.Errorf("创建 Hook 软链接失败: %w", err) } - - return copyFile(execPath, HookBinaryPath, 0755) + return nil } -func (h *HookManager) removeHookBinary() { - info, err := os.Lstat(HookBinaryPath) +func (h *HookManager) removeHookBinary() error { + return removeHookSymlink(HookBinaryPath) +} + +func removeHookSymlink(path string) error { + info, err := os.Lstat(path) if err != nil { - return + if os.IsNotExist(err) { + return nil + } + return err } + if info.Mode()&os.ModeSymlink != 0 { - _ = os.Remove(HookBinaryPath) - } -} - -func copyFile(src, dst string, perm os.FileMode) error { - in, err := os.Open(src) - if err != nil { - return err - } - defer in.Close() - - out, err := os.Create(dst) - if err != nil { - return err - } - defer out.Close() - - if _, err := io.Copy(out, in); err != nil { - return err + return os.Remove(path) } - return os.Chmod(dst, perm) + // 兼容旧版本:普通文件可能是主程序、旧版复制品或用户文件,无法确认所有权时保留。 + return nil } // modifyPAMConfig 修改 PAM 配置 diff --git a/pkg/agent/sshmonitor/hook_manager_test.go b/pkg/agent/sshmonitor/hook_manager_test.go new file mode 100644 index 0000000..b7cdfd5 --- /dev/null +++ b/pkg/agent/sshmonitor/hook_manager_test.go @@ -0,0 +1,43 @@ +package sshmonitor + +import ( + "os" + "path/filepath" + "testing" +) + +func TestRemoveHookSymlinkRemovesSymlink(t *testing.T) { + dir := t.TempDir() + target := filepath.Join(dir, "agent") + link := filepath.Join(dir, "pika-agent") + if err := os.WriteFile(target, []byte("test"), 0755); err != nil { + t.Fatalf("write target: %v", err) + } + if err := os.Symlink(target, link); err != nil { + t.Fatalf("create symlink: %v", err) + } + + if err := removeHookSymlink(link); err != nil { + t.Fatalf("removeHookSymlink() error = %v", err) + } + if _, err := os.Lstat(link); !os.IsNotExist(err) { + t.Fatalf("hook symlink still exists: %v", err) + } + if _, err := os.Stat(target); err != nil { + t.Fatalf("symlink target was removed: %v", err) + } +} + +func TestRemoveHookSymlinkKeepsRegularFile(t *testing.T) { + path := filepath.Join(t.TempDir(), "pika-agent") + if err := os.WriteFile(path, []byte("test"), 0755); err != nil { + t.Fatalf("write regular file: %v", err) + } + + if err := removeHookSymlink(path); err != nil { + t.Fatalf("removeHookSymlink() error = %v", err) + } + if _, err := os.Stat(path); err != nil { + t.Fatalf("regular executable was removed: %v", err) + } +} diff --git a/pkg/agent/sshmonitor/monitor.go b/pkg/agent/sshmonitor/monitor.go index 31b4230..c0b2931 100644 --- a/pkg/agent/sshmonitor/monitor.go +++ b/pkg/agent/sshmonitor/monitor.go @@ -91,6 +91,31 @@ func (m *Monitor) Stop() error { return m.stopInternal() } +// Uninstall 清理 SSH 登录监控安装的持久配置和运行时文件。 +func (m *Monitor) Uninstall() error { + m.mu.Lock() + defer m.mu.Unlock() + + if m.enabled { + if err := m.stopInternal(); err != nil { + return err + } + } else { + if err := os.Remove(m.socketPath); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("移除 socket 失败: %w", err) + } + if err := m.hookManager.Uninstall(); err != nil { + return err + } + } + + if err := os.Remove(DefaultSocketDir); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("移除 socket 目录失败: %w", err) + } + + return nil +} + // stopInternal 内部停止方法(不加锁) func (m *Monitor) stopInternal() error { if !m.enabled { @@ -109,18 +134,19 @@ func (m *Monitor) stopInternal() error { m.listener = nil } + var cleanupErr error if err := os.Remove(m.socketPath); err != nil && !os.IsNotExist(err) { - slog.Warn("移除 socket 失败", "error", err, "path", m.socketPath) + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("移除 socket 失败: %w", err)) } // 卸载 PAM Hook if err := m.hookManager.Uninstall(); err != nil { - slog.Warn("卸载 PAM Hook 失败", "error", err) + cleanupErr = errors.Join(cleanupErr, fmt.Errorf("卸载 PAM Hook 失败: %w", err)) } m.enabled = false slog.Info("SSH登录监控已停止") - return nil + return cleanupErr } // GetEvents 获取事件通道