diff --git a/cmd/safe/safe.go b/cmd/safe/safe.go index f6f9547..d607ecc 100644 --- a/cmd/safe/safe.go +++ b/cmd/safe/safe.go @@ -16,19 +16,22 @@ import ( "golang.org/x/term" ) -var SafeCmd = &cobra.Command{ - Use: "safe", - Short: "交互式安全确认,将可疑对象加入白名单", - Long: "查看当前的可疑文件和进程列表,并选择将其移入白名单。", - Run: func(cmd *cobra.Command, args []string) { - cfg, err := config.LoadConfig("./config.yaml") - if err != nil { - fmt.Printf("加载配置失败: %v\n", err) - os.Exit(1) - } +func NewSafeCmd() *cobra.Command { + cmd := &cobra.Command{ + Use: "safe", + Short: "交互式安全确认,将可疑对象加入白名单", + Long: "查看当前的可疑文件和进程列表,并选择将其移入白名单。", + Run: func(cmd *cobra.Command, args []string) { + cfg, err := config.LoadConfig("./config.yaml") + if err != nil { + fmt.Printf("加载配置失败: %v\n", err) + os.Exit(1) + } - interactiveSafe(cfg) - }, + interactiveSafe(cfg) + }, + } + return cmd } func readKeyWithESC() (string, error) { diff --git a/cmd/start/start.go b/cmd/start/start.go index df3d8b6..c68cab8 100644 --- a/cmd/start/start.go +++ b/cmd/start/start.go @@ -20,160 +20,164 @@ import ( "go.uber.org/zap" ) -var StartCmd = &cobra.Command{ - Use: "start", - Short: "启动系统监控守护服务", - Long: "sysmonitord start 命令用于启动系统监控守护服务,首次启动会进行全量扫描建立白名单。", - Run: func(cmd *cobra.Command, args []string) { - logger.Log.Info("正在启动系统监控守护服务...") +func NewStartCmd() *cobra.Command { + cmd := &cobra.Command{ + Use: "start", + Short: "启动系统监控守护服务", + Long: "sysmonitord start 命令用于启动系统监控守护服务,首次启动会进行全量扫描建立白名单。", + Run: func(cmd *cobra.Command, args []string) { + logger.Log.Info("正在启动系统监控守护服务...") - cfg, err := config.LoadConfig("./config.yaml") - if err != nil { - logger.Log.Error("加载配置文件失败", zap.Error(err)) - os.Exit(1) - } - - logger.Log.Info("配置文件加载成功", - zap.String("审计服务器地址", fmt.Sprintf("%s:%d", cfg.Audit.Server, cfg.Audit.Port)), - ) - - storageCfg := &storage.Storage{ - DataDir: cfg.Storage.DataDir, - ProcessSystemFile: cfg.Storage.ProcessSystemFile, - FileSystemFile: cfg.Storage.FileSystemFile, - } - - // ====== 进程扫描和存储 ====== - - startTime := time.Now() - procs, err := process.ScanAllProcesses(cfg) - - logger.Log.Info("进程扫描完成", - zap.Int("进程数量", len(procs)), - zap.Duration("扫描耗时", time.Since(startTime)), - ) - - if err != nil { - logger.Log.Error("扫描进程失败", zap.Error(err)) - os.Exit(1) - } else { - if err := storage.SaveProcessSystem(procs, storageCfg.DataDir, storageCfg.ProcessSystemFile); err != nil { - logger.Log.Error("保存进程白名单失败", zap.Error(err)) - } - } - - // ====== 文件扫描和存储 ====== - logger.Log.Info("正在扫描文件系统...") - - startTime = time.Now() - fileScanner := file.NewScanner(cfg) - - files, err := fileScanner.Scan() - if err != nil { - logger.Log.Error("扫描文件系统失败", zap.Error(err)) - os.Exit(1) - } else { - if err := storage.SaveFileSystem(files, storageCfg.DataDir, storageCfg.FileSystemFile); err != nil { - logger.Log.Error("保存文件系统白名单失败", zap.Error(err)) + cfg, ok := cmd.Context().Value("config").(*config.Config) + if !ok { + logger.Log.Error("无法获取配置") os.Exit(1) } - } - duration := time.Since(startTime) - logger.Log.Info("文件系统扫描完成", - zap.Int("文件数量", len(files)), - zap.Duration("扫描耗时", duration), - ) - - // ====== 启动文件监听 ====== - logger.Log.Info("正在启动文件监听...") - - fileMon, err := watcher.NewWatcher(cfg) - if err != nil { - logger.Log.Error("启动文件监听失败", zap.Error(err)) - os.Exit(1) - } - - fileMon.Start() - - // ====== 初始化文件检测器 ====== - fileDetector, err := detector.NewFileDetector(cfg) - if err != nil { - logger.Log.Error("初始化文件检测器失败", zap.Error(err)) - os.Exit(1) - } - - // ====== 启动进程检测定时任务 ====== - procDetector, err := detector.NewProcessDetector(cfg) - if err != nil { - logger.Log.Error("初始化进程检测器失败", zap.Error(err)) - os.Exit(1) - } - - procEventChan := procDetector.Event() - procScheduler := timer.NewScheduler(time.Duration(cfg.Scanner.Process.Interval)*time.Second, procDetector) - procScheduler.Start() - - // ====== 启动告警管理器 ====== - alerter := notifier.NewAlerter(cfg.Notification) - alerter.Start() - - logger.Log.Info("系统监控守护服务已启动,正在监控系统变化...") - - quit := make(chan os.Signal, 1) - signal.Notify(quit, os.Interrupt, syscall.SIGTERM, syscall.SIGINT) - - for { - select { - - case event := <-fileMon.Events(): - logger.Log.Info("文件系统事件", - zap.String("path", event.Path), - zap.String("op", event.Op.String()), - ) - - if event.FileInfo != nil { - logger.Log.Debug("文件详情", zap.Int64("size", event.FileInfo.Size())) - fileDetector.HandleEvent(event.Path, event.Op.String()) - } - - // test - alerter.PushAlert(notifier.AlertEvent{ - Type: "File", - Path: event.Path, - Reason: event.Op.String(), - Details: "To test", - }) - - case procEvents := <-procEventChan: - logger.Log.Info("可疑进程事件", - zap.Int32("pid", procEvents.PID), - zap.String("name", procEvents.Name), - zap.String("path", procEvents.Path), - ) - - procDetector.HandleDubiousProcesses(procEvents) - - // test - alerter.PushAlert(notifier.AlertEvent{ - Type: "Process", - Path: procEvents.Path, - Reason: "可疑进程", - Details: "To test", - }) - - case err := <-fileMon.Errors(): - logger.Log.Error("文件监听错误", zap.Error(err)) - - case <-quit: - logger.Log.Info("正在停止系统监控守护服务...") - fileMon.Stop() - procScheduler.Stop() - logger.Log.Info("系统监控守护服务已停止") - return + logger.Log.Info("配置文件加载成功", + zap.String("审计服务器地址", fmt.Sprintf("%s:%d", cfg.Audit.Server, cfg.Audit.Port)), + ) + storageCfg := &storage.Storage{ + DataDir: cfg.Storage.DataDir, + ProcessSystemFile: cfg.Storage.ProcessSystemFile, + FileSystemFile: cfg.Storage.FileSystemFile, } - } - }, + // ====== 进程扫描和存储 ====== + + startTime := time.Now() + procs, err := process.ScanAllProcesses(cfg) + + logger.Log.Info("进程扫描完成", + zap.Int("进程数量", len(procs)), + zap.Duration("扫描耗时", time.Since(startTime)), + ) + + if err != nil { + logger.Log.Error("扫描进程失败", zap.Error(err)) + os.Exit(1) + } else { + if err := storage.SaveProcessSystem(procs, storageCfg.DataDir, storageCfg.ProcessSystemFile); err != nil { + logger.Log.Error("保存进程白名单失败", zap.Error(err)) + } + } + + // ====== 文件扫描和存储 ====== + logger.Log.Info("正在扫描文件系统...") + + startTime = time.Now() + fileScanner := file.NewScanner(cfg) + + files, err := fileScanner.Scan() + if err != nil { + logger.Log.Error("扫描文件系统失败", zap.Error(err)) + os.Exit(1) + } else { + if err := storage.SaveFileSystem(files, storageCfg.DataDir, storageCfg.FileSystemFile); err != nil { + logger.Log.Error("保存文件系统白名单失败", zap.Error(err)) + os.Exit(1) + } + } + + duration := time.Since(startTime) + logger.Log.Info("文件系统扫描完成", + zap.Int("文件数量", len(files)), + zap.Duration("扫描耗时", duration), + ) + + // ====== 启动文件监听 ====== + logger.Log.Info("正在启动文件监听...") + + fileMon, err := watcher.NewWatcher(cfg) + if err != nil { + logger.Log.Error("启动文件监听失败", zap.Error(err)) + os.Exit(1) + } + + fileMon.Start() + + // ====== 初始化文件检测器 ====== + fileDetector, err := detector.NewFileDetector(cfg) + if err != nil { + logger.Log.Error("初始化文件检测器失败", zap.Error(err)) + os.Exit(1) + } + + // ====== 启动进程检测定时任务 ====== + procDetector, err := detector.NewProcessDetector(cfg) + if err != nil { + logger.Log.Error("初始化进程检测器失败", zap.Error(err)) + os.Exit(1) + } + + procEventChan := procDetector.Event() + procScheduler := timer.NewScheduler(time.Duration(cfg.Scanner.Process.Interval)*time.Second, procDetector) + procScheduler.Start() + + // ====== 启动告警管理器 ====== + alerter := notifier.NewAlerter(cfg.Notification) + alerter.Start() + + logger.Log.Info("系统监控守护服务已启动,正在监控系统变化...") + + quit := make(chan os.Signal, 1) + signal.Notify(quit, os.Interrupt, syscall.SIGTERM, syscall.SIGINT) + + for { + select { + + case event := <-fileMon.Events(): + logger.Log.Info("文件系统事件", + zap.String("path", event.Path), + zap.String("op", event.Op.String()), + ) + + if event.FileInfo != nil { + logger.Log.Debug("文件详情", zap.Int64("size", event.FileInfo.Size())) + fileDetector.HandleEvent(event.Path, event.Op.String()) + } + + // test + alerter.PushAlert(notifier.AlertEvent{ + Type: "File", + Path: event.Path, + Reason: event.Op.String(), + Details: "To test", + }) + + case procEvents := <-procEventChan: + logger.Log.Info("可疑进程事件", + zap.Int32("pid", procEvents.PID), + zap.String("name", procEvents.Name), + zap.String("path", procEvents.Path), + ) + + procDetector.HandleDubiousProcesses(procEvents) + + // test + alerter.PushAlert(notifier.AlertEvent{ + Type: "Process", + Path: procEvents.Path, + Reason: "可疑进程", + Details: "To test", + }) + + case err := <-fileMon.Errors(): + logger.Log.Error("文件监听错误", zap.Error(err)) + + case <-quit: + logger.Log.Info("正在停止系统监控守护服务...") + fileMon.Stop() + procScheduler.Stop() + logger.Log.Info("系统监控守护服务已停止") + return + + } + } + + }, + } + + return cmd } diff --git a/cmd/status/status.go b/cmd/status/status.go index 7c0eb78..d5507b2 100644 --- a/cmd/status/status.go +++ b/cmd/status/status.go @@ -10,20 +10,23 @@ import ( "github.com/spf13/cobra" ) -var StatusCmd = &cobra.Command{ - Use: "status", - Short: "显示系统状态", - Long: "显示Sysmonitod的当前状态", - Run: func(cmd *cobra.Command, args []string) { - cfg, err := config.LoadConfig("./config.yaml") +func NewStatusCmd() *cobra.Command { + cmd := &cobra.Command{ + Use: "status", + Short: "显示系统状态", + Long: "显示Sysmonitod的当前状态", + Run: func(cmd *cobra.Command, args []string) { + cfg, err := config.LoadConfig("./config.yaml") - if err != nil { - fmt.Printf("加载配置失败: %v\n", err) - os.Exit(1) - } + if err != nil { + fmt.Printf("加载配置失败: %v\n", err) + os.Exit(1) + } - printStatus(cfg) - }, + printStatus(cfg) + }, + } + return cmd } func printStatus(cfg *config.Config) { diff --git a/cmd/version/version.go b/cmd/version/version.go index 8a608b0..63ae331 100644 --- a/cmd/version/version.go +++ b/cmd/version/version.go @@ -7,11 +7,14 @@ import ( "github.com/spf13/cobra" ) -var VersionCmd = &cobra.Command{ - Use: "version", - Short: "显示 sysmonitord 的版本信息", - Long: "sysmonitord version 命令用于显示当前 sysmonitord 的版本、Git 提交信息和构建时间。", - Run: func(cmd *cobra.Command, args []string) { - fmt.Println(version.Info()) - }, +func NewVersionCmd() *cobra.Command { + cmd := &cobra.Command{ + Use: "version", + Short: "显示 sysmonitord 的版本信息", + Long: "sysmonitord version 命令用于显示当前 sysmonitord 的版本、Git 提交信息和构建时间。", + Run: func(cmd *cobra.Command, args []string) { + fmt.Println(version.Info()) + }, + } + return cmd } diff --git a/main.go b/main.go index 79a6d6b..fab0c81 100644 --- a/main.go +++ b/main.go @@ -1,6 +1,8 @@ package main import ( + "context" + "fmt" "os" "sysmonitord/cmd/safe" "sysmonitord/cmd/start" @@ -13,39 +15,44 @@ import ( "go.uber.org/zap" ) -func getConfigPath() string { - if _, err := os.Stat("./config.yaml"); err == nil { - return "./config.yaml" - } - - if _, err := os.Stat("/etc/sysmonitord/config.yaml"); err == nil { - return "/etc/sysmonitord/config.yaml" - } - - return "./config.yaml" -} +var ( + cfgFile string + cfg *config.Config +) func main() { logger.InitLogger() defer logger.Sync() - cfg, err := config.LoadConfig(getConfigPath()) - if err != nil { - logger.Log.Error("加载配置文件失败", zap.Error(err)) - os.Exit(1) - } else { - logger.SetLogLevel(cfg.Log.Level) - } - var rootCmd = &cobra.Command{ Use: "sysmonitord", Short: "Sysmonitord 是一个 Linux 系统安全监控工具", + PersistentPreRunE: func(cmd *cobra.Command, args []string) error { + if cfgFile == "" { + if _, err := os.Stat("./config.yaml"); err == nil { + cfgFile = "./config.yaml" + } else if _, err := os.Stat("/etc/sysmonitord/config.yaml"); err == nil { + cfgFile = "/etc/sysmonitord/config.yaml" + } + } + + cfg, err := config.LoadConfig(cfgFile) + if err != nil { + return fmt.Errorf("加载配置文件失败: %w", err) + } + + ctx := context.WithValue(cmd.Context(), "config", cfg) + cmd.SetContext(ctx) + return nil + }, } - rootCmd.AddCommand(start.StartCmd) - rootCmd.AddCommand(version.VersionCmd) - rootCmd.AddCommand(status.StatusCmd) - rootCmd.AddCommand(safe.SafeCmd) + rootCmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", "", "配置文件路径 (默认: ./config.yaml 或 /etc/sysmonitord/config.yaml)") + + rootCmd.AddCommand(start.NewStartCmd()) + rootCmd.AddCommand(version.NewVersionCmd()) + rootCmd.AddCommand(status.NewStatusCmd()) + rootCmd.AddCommand(safe.NewSafeCmd()) if err := rootCmd.Execute(); err != nil { logger.Log.Error("命令执行失败", zap.Error(err))