bash_go_service/backend-service/cmd/main.go

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...")
		err := checkExistingProcess()
		if err != nil {
			logger.Error("Error starting daemon.")
			return
		}
		runDaemon()
	}
}

func main() {
	if isDaemon() {
		return
	}
	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)
}