package main import ( "fmt" "os" "os/exec" "path/filepath" "strconv" "strings" "syscall" "time" "backend-service/pkg/client" "backend-service/pkg/machine" "backend-service/pkg/service" "bash_go_service/shared/pkg/constants" "bash_go_service/shared/pkg/logger" "github.com/spf13/viper" ) // isDaemon 检查是否以守护进程模式运行 func isDaemon() bool { return len(os.Args) > 1 && os.Args[1] == constants.DaemonFlag } func init() { logger.Debug("Initializing...") if isDaemon() { logger.Debug("Starting daemon...") runDaemon() } } func main() { if err := run(); err != nil { logger.Error("Application failed: %v", err) os.Exit(1) } } // run 程序主逻辑 func run() error { // 初始化配置 if err := initConfig(); err != nil { return fmt.Errorf("init config: %w", err) } // 采集并发送机器信息 info, err := collectAndSendMachineInfo() if err != nil { return fmt.Errorf("collect and send machine info: %w", err) } // 显示欢迎信息 showWelcomeMessage(info) // 检查并启动守护进程 if err := handleDaemonProcess(); err != nil { return fmt.Errorf("handle daemon process: %w", err) } return nil } // initConfig 初始化配置 func initConfig() error { viper.AddConfigPath(constants.ConfigPath) return viper.ReadInConfig() } // collectAndSendMachineInfo 采集并发送机器信息 func collectAndSendMachineInfo() (*machine.Info, error) { info := machine.CollectMachineInfo() if err := client.SendMachineInfo(info); err != nil { return nil, fmt.Errorf("send machine info: %w", err) } return info, nil } // handleDaemonProcess 处理守护进程相关逻辑 func handleDaemonProcess() error { // 获取可执行文件路径 exePath, err := os.Executable() if err != nil { return fmt.Errorf("get executable path: %w", err) } logger.Debug("Exe path: %s", filepath.Dir(exePath)) // 检查现有进程 if err := checkExistingProcess(); err != nil { return err } // 启动新的守护进程 return startDaemonProcess() } // checkExistingProcess 检查是否已有守护进程在运行 func checkExistingProcess() error { pidFile := constants.PidFilePath pidBytes, err := os.ReadFile(pidFile) if err != nil { if os.IsNotExist(err) { return nil } return fmt.Errorf("read pid file: %w", err) } pid, err := strconv.Atoi(string(pidBytes)) if err != nil { return fmt.Errorf("parse pid: %w", err) } process, err := os.FindProcess(pid) if err != nil { return fmt.Errorf("find process: %w", err) } if err := process.Signal(syscall.Signal(0)); err == nil { logger.Info("Daemon already running with PID: %d", pid) return fmt.Errorf("daemon already running") } // 进程不存在,删除PID文件 if err := os.Remove(pidFile); err != nil { logger.Error("Failed to remove old PID file: %v", err) } return nil } func formatGoTime(format string) string { replacer := strings.NewReplacer( "%Y", "2006", "%m", "01", "%d", "02", "%H", "15", "%M", "04", "%S", "05", ) return replacer.Replace(format) } // startDaemonProcess 启动守护进程 func startDaemonProcess() error { // 确保日志目录存在 if err := os.MkdirAll(constants.LogFilePath, constants.LogFileMode); err != nil { return fmt.Errorf("create logs directory: %w", err) } // 将 strftime 风格格式转换为 Go 的时间格式 goTimeFormat := formatGoTime(constants.LogNameFormate) logFileName := time.Now().Format(goTimeFormat) logFilePath := filepath.Join(constants.LogFilePath, logFileName) logFile, err := os.OpenFile(logFilePath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, constants.LogFileMode) if err != nil { return fmt.Errorf("open log file: %w", err) } cmd := exec.Command(os.Args[0], constants.DaemonFlag) cmd.Stdout = logFile cmd.Stderr = logFile cmd.SysProcAttr = &syscall.SysProcAttr{ Setsid: true, } if err := cmd.Start(); err != nil { return fmt.Errorf("start daemon: %w", err) } // 写入 PID 文件 pidStr := []byte(fmt.Sprintf("%d", cmd.Process.Pid)) if err := os.WriteFile(constants.PidFilePath, pidStr, constants.LogFileMode); err != nil { return fmt.Errorf("write pid file: %w", err) } return nil } // runDaemon 守护进程主逻辑 func runDaemon() { logger.Info("Daemon started with PID: %d", os.Getpid()) service.Run() } // showWelcomeMessage 显示欢迎信息 func showWelcomeMessage(info *machine.Info) { message := fmt.Sprintf(`%s ╔══════════════════════════════════════╗ ║ Welcome to Bash Security Service ║ ║ Hostname: %-20s ║ ║ OS: %-32s ║ ╚══════════════════════════════════════╝%s`, constants.ColorGreen, info.Hostname, info.OS, constants.ColorReset) fmt.Println(message) }