note/tech/auto_check.go

package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
)

type Config struct {
	AccessToken  string `json:"access_token"`
	Keyword      string `json:"keyword"`
	Timeout      int    `json:"timeout"`
	MaxTime      int    `json:"max_time"`
	MaxFileSize  int    `json:"max_file_size"`
	WebhookURL   string `json:"webhook_url"`
	TaskFilePath string `json:"task_file_path"`
	Port         int    `json:"port"`         // 新增：自定义端口号
	BindAddress  string `json:"bind_address"` // 新增：绑定地址（127.0.0.1 或 0.0.0.0）
}

type CronJob struct {
	ID            int    `json:"id"`
	Name          string `json:"name"`
	URL           string `json:"url"`
	ExpectedKey   string `json:"expected_key"`
	Interval      int    `json:"interval"`
	LastExecution string `json:"last_execution"`
}

var (
	config     Config
	cronJobs   = make(map[int]CronJob)
	jobIDMutex sync.Mutex
	jobID      = 0

	// 日志文件
	checkLog  *log.Logger
	accessLog *log.Logger
)

const (
	defaultConfigFile = "~/.config/checkurl/cron.conf"
	defaultTaskFile   = "~/.config/checkurl/task.list"
	defaultWebhookURL = "https://oapi.dingtalk.com/robot/send?access_token=%s"
)

func main() {
	// 解析命令行参数
	configFilePath := flag.String("config", defaultConfigFile, "Path to the configuration file")
	flag.Parse()

	// 初始化日志
	initLogs()

	// 初始化配置
	loadConfig(*configFilePath)

	// 加载任务
	loadCronJobs()

	// 启动定时任务
	go startCronJobs()

	// 设置Gin路由
	r := gin.Default()

	// 设置Gin的日志输出到access.log
	r.Use(gin.LoggerWithWriter(accessLog.Writer()))

	r.POST("/check_url", checkURLHandler)
	r.POST("/add_cron", addCronHandler)
	r.POST("/del_cron", delCronHandler)
	r.GET("/get_cron_list", getCronListHandler)
	r.POST("/config_cron", configCronHandler)
	r.GET("/get_config", getConfigHandler)
	// 启动Gin服务
	bindAddress := fmt.Sprintf("%s:%d", config.BindAddress, config.Port)
	r.Run(bindAddress)
}

func initLogs() {
	// 创建check.log文件
	checkLogFile, err := os.OpenFile("check.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		log.Fatalf("无法打开check.log文件: %v", err)
	}
	checkLog = log.New(checkLogFile, "[CHECK] ", log.Ldate|log.Ltime|log.Lshortfile)

	// 创建access.log文件
	accessLogFile, err := os.OpenFile("access.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		log.Fatalf("无法打开access.log文件: %v", err)
	}
	accessLog = log.New(accessLogFile, "[ACCESS] ", log.Ldate|log.Ltime|log.Lshortfile)
}
func loadConfig(configFilePath string) {
	// 默认配置
	config = Config{
		AccessToken:  getEnv("ACCESS_TOKEN", "b2a92184affebb3f26dfcab1eff9571b46c27a85380fd9ff56ca4c66cf93e1a1"),
		Keyword:      getEnv("KEYWORD", "BWH:"),
		Timeout:      getEnvInt("TIMEOUT", 5),
		MaxTime:      getEnvInt("MAX_TIME", 10),
		MaxFileSize:  getEnvInt("MAX_FILESIZE", 300),
		WebhookURL:   getEnv("WEBHOOK_URL", defaultWebhookURL),
		TaskFilePath: getEnv("TASK_FILE_PATH", defaultTaskFile),
		Port:         getEnvInt("PORT", 2809),
		BindAddress:  getEnv("BIND_ADDRESS", "127.0.0.1"),
	}

	// 检查用户家目录下的 .config/checkurl/cron.conf 文件
	homeDir, err := os.UserHomeDir()
	if err != nil {
		checkLog.Printf("无法获取用户家目录: %v\n", err)
	} else {
		userConfigPath := filepath.Join(homeDir, ".config", "checkurl", "cron.conf")
		if _, err := os.Stat(userConfigPath); os.IsNotExist(err) {
			// 如果文件不存在，创建目录和文件
			createUserConfigFile(userConfigPath)
		} else {
			// 如果文件存在，从该文件读取配置
			loadUserConfig(userConfigPath)
		}
	}

	// 从指定的配置文件中读取配置项
	if configFilePath != defaultConfigFile {
		file, err := os.Open(configFilePath)
		if err != nil {
			checkLog.Printf("未找到配置文件 %s，将使用默认配置\n", configFilePath)
		} else {
			defer file.Close()
			scanner := bufio.NewScanner(file)
			for scanner.Scan() {
				line := scanner.Text()
				if strings.HasPrefix(line, "#") || line == "" {
					continue
				}

				fields := strings.Fields(line)
				if len(fields) != 2 {
					checkLog.Printf("忽略无效配置行: %s\n", line)
					continue
				}

				key := fields[0]
				value := strings.Trim(fields[1], `"`)

				switch key {
				case "ACCESS_TOKEN":
					config.AccessToken = value
				case "KEYWORD":
					config.Keyword = value
				case "TIMEOUT":
					if timeout, err := strconv.Atoi(value); err == nil {
						config.Timeout = timeout
					}
				case "MAX_TIME":
					if maxTime, err := strconv.Atoi(value); err == nil {
						config.MaxTime = maxTime
					}
				case "MAX_FILESIZE":
					if maxFileSize, err := strconv.Atoi(value); err == nil {
						config.MaxFileSize = maxFileSize
					}
				case "WEBHOOK_URL":
					config.WebhookURL = value
				case "TASK_FILE_PATH":
					config.TaskFilePath = value
				case "PORT":
					if port, err := strconv.Atoi(value); err == nil {
						config.Port = port
					}
				case "BIND_ADDRESS":
					config.BindAddress = value
				default:
					checkLog.Printf("忽略未知配置项: %s\n", key)
				}
			}
		}
	}

	// 如果 TaskFilePath 没有定义，则默认放在 ~/.config/checkurl/ 文件夹中
	if config.TaskFilePath == "" {
		config.TaskFilePath = filepath.Join(homeDir, ".config", "checkurl", "task.list")
	}
}

// 创建用户配置文件
func createUserConfigFile(userConfigPath string) {
	// 创建目录
	configDir := filepath.Dir(userConfigPath)
	if err := os.MkdirAll(configDir, 0755); err != nil {
		checkLog.Printf("无法创建用户配置目录 %s: %v\n", configDir, err)
		return
	}

	// 创建文件
	file, err := os.Create(userConfigPath)
	if err != nil {
		checkLog.Printf("无法创建用户配置文件 %s: %v\n", userConfigPath, err)
		return
	}
	defer file.Close()

	// 写入默认配置
	writer := bufio.NewWriter(file)
	_, _ = writer.WriteString("# 用户配置文件\n")
	_, _ = writer.WriteString(fmt.Sprintf("ACCESS_TOKEN %s\n", config.AccessToken))
	_, _ = writer.WriteString(fmt.Sprintf("KEYWORD %s\n", config.Keyword))
	_, _ = writer.WriteString(fmt.Sprintf("TIMEOUT %d\n", config.Timeout))
	_, _ = writer.WriteString(fmt.Sprintf("MAX_TIME %d\n", config.MaxTime))
	_, _ = writer.WriteString(fmt.Sprintf("MAX_FILESIZE %d\n", config.MaxFileSize))
	_, _ = writer.WriteString(fmt.Sprintf("WEBHOOK_URL %s\n", config.WebhookURL))
	_, _ = writer.WriteString(fmt.Sprintf("TASK_FILE_PATH %s\n", config.TaskFilePath))
	_, _ = writer.WriteString(fmt.Sprintf("PORT %d\n", config.Port))
	_, _ = writer.WriteString(fmt.Sprintf("BIND_ADDRESS %s\n", config.BindAddress))
	writer.Flush()

	checkLog.Printf("已创建用户配置文件: %s\n", userConfigPath)
}

func loadUserConfig(userConfigPath string) {
	file, err := os.Open(userConfigPath)
	if err != nil {
		checkLog.Printf("无法打开用户配置文件 %s: %v\n", userConfigPath, err)
		return
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "#") || line == "" {
			continue
		}

		fields := strings.Fields(line)
		if len(fields) != 2 {
			checkLog.Printf("忽略无效配置行: %s\n", line)
			continue
		}

		key := fields[0]
		value := strings.Trim(fields[1], `"`)

		switch key {
		case "ACCESS_TOKEN":
			config.AccessToken = value
		case "KEYWORD":
			config.Keyword = value
		case "TIMEOUT":
			if timeout, err := strconv.Atoi(value); err == nil {
				config.Timeout = timeout
			}
		case "MAX_TIME":
			if maxTime, err := strconv.Atoi(value); err == nil {
				config.MaxTime = maxTime
			}
		case "MAX_FILESIZE":
			if maxFileSize, err := strconv.Atoi(value); err == nil {
				config.MaxFileSize = maxFileSize
			}
		case "WEBHOOK_URL":
			config.WebhookURL = value
		case "TASK_FILE_PATH":
			config.TaskFilePath = value
		default:
			checkLog.Printf("忽略未知配置项: %s\n", key)
		}
	}

	checkLog.Printf("已从用户配置文件 %s 加载配置\n", userConfigPath)
}
func loadCronJobs() {
	file, err := os.Open(config.TaskFilePath)
	if err != nil {
		checkLog.Printf("未找到任务文件 %s，将创建新文件\n", config.TaskFilePath)
		return
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := scanner.Text()

		// 跳过注释行（以 # 开头的行）
		if strings.HasPrefix(line, "#") || line == "" {
			continue
		}

		// 手动解析任务行
		fields := parseTaskLine(line)
		if len(fields) != 5 {
			checkLog.Printf("忽略无效任务行: %s\n", line)
			continue
		}

		id, _ := strconv.Atoi(fields[0])
		name := fields[1]
		url := fields[2]
		expectedKey := fields[3]
		interval, _ := strconv.Atoi(fields[4])

		cronJobs[id] = CronJob{
			ID:          id,
			Name:        name,
			URL:         url,
			ExpectedKey: expectedKey,
			Interval:    interval,
		}

		if id > jobID {
			jobID = id
		}
	}
}

// 手动解析任务行
func parseTaskLine(line string) []string {
	var fields []string
	var currentField string
	inQuotes := false

	for _, char := range line {
		switch char {
		case ' ':
			if inQuotes {
				currentField += " "
			} else if currentField != "" {
				fields = append(fields, currentField)
				currentField = ""
			}
		case '"':
			inQuotes = !inQuotes
		default:
			currentField += string(char)
		}
	}

	if currentField != "" {
		fields = append(fields, currentField)
	}

	return fields
}

func saveCronJobs() error {
	file, err := os.Create(config.TaskFilePath)
	if err != nil {
		return err
	}
	defer file.Close()

	writer := bufio.NewWriter(file)
	_, _ = writer.WriteString("# task_id task_name api_url expected_keyword cron_time\n")

	for _, job := range cronJobs {
		line := fmt.Sprintf("%d \"%s\" \"%s\" \"%s\" %d\n", job.ID, job.Name, job.URL, job.ExpectedKey, job.Interval)
		_, _ = writer.WriteString(line)
	}

	return writer.Flush()
}

func checkURLHandler(c *gin.Context) {
	var request struct {
		URL         string `json:"url"`
		ExpectedKey string `json:"expected_key"`
	}
	if err := c.ShouldBindJSON(&request); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	result := checkAPI(request.URL, request.ExpectedKey)
	c.JSON(http.StatusOK, gin.H{"result": result})
}

func addCronHandler(c *gin.Context) {
	var request struct {
		Name        string `json:"name"`
		URL         string `json:"url"`
		ExpectedKey string `json:"expected_key"`
		Interval    int    `json:"interval"`
	}
	if err := c.ShouldBindJSON(&request); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	jobIDMutex.Lock()
	jobID++
	newJob := CronJob{
		ID:          jobID,
		Name:        request.Name,
		URL:         request.URL,
		ExpectedKey: request.ExpectedKey,
		Interval:    request.Interval,
	}
	cronJobs[jobID] = newJob
	jobIDMutex.Unlock()

	// 保存任务到任务文件
	if err := saveCronJobs(); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "无法保存任务到任务文件"})
		return
	}

	c.JSON(http.StatusOK, gin.H{"message": "Cron job added", "job_id": jobID})
}

func delCronHandler(c *gin.Context) {
	var request struct {
		ID int `json:"id"`
	}
	if err := c.ShouldBindJSON(&request); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	jobIDMutex.Lock()
	delete(cronJobs, request.ID)
	jobIDMutex.Unlock()

	// 保存任务到任务文件
	if err := saveCronJobs(); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "无法保存任务到任务文件"})
		return
	}

	c.JSON(http.StatusOK, gin.H{"message": "Cron job deleted"})
}

func getCronListHandler(c *gin.Context) {
	var jobList []CronJob
	for _, job := range cronJobs {
		jobList = append(jobList, job)
	}
	c.JSON(http.StatusOK, gin.H{"cron_jobs": jobList})
}

func configCronHandler(c *gin.Context) {
	var newConfig Config
	if err := c.ShouldBindJSON(&newConfig); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	config = newConfig
	c.JSON(http.StatusOK, gin.H{"message": "Config updated"})
}

func getConfigHandler(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{"config": config})
}

func startCronJobs() {
	for {
		for _, job := range cronJobs {
			go func(job CronJob) {
				result := checkAPI(job.URL, job.ExpectedKey)
				if result != "success" {
					sendAlert(job.Name, result)
				}
				job.LastExecution = time.Now().Format(time.RFC3339)
				jobIDMutex.Lock()
				cronJobs[job.ID] = job
				jobIDMutex.Unlock()
			}(job)
			time.Sleep(time.Duration(job.Interval) * time.Second)
		}
		time.Sleep(1 * time.Second)
	}
}

func checkAPI(apiURL, expectedKeyword string) string {
	client := &http.Client{
		Timeout: time.Duration(config.Timeout) * time.Second,
	}

	resp, err := client.Get(apiURL)
	if err != nil {
		checkLog.Printf("任务失败: URL=%s, 错误=%s\n", apiURL, err.Error())
		return "服务请求失败（可能是网络问题或端口未开放）"
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		checkLog.Printf("任务失败: URL=%s, 错误=%s\n", apiURL, err.Error())
		return "服务请求失败（读取响应失败）"
	}

	if !strings.Contains(string(body), expectedKeyword) {
		checkLog.Printf("任务失败: URL=%s, 期望关键词=%s, 响应=%s\n", apiURL, expectedKeyword, string(body))
		return fmt.Sprintf("服务返回异常内容（不包含期望的关键词 '%s'），状态码：%d 响应: %s", expectedKeyword, resp.StatusCode, string(body))
	}

	checkLog.Printf("任务成功: URL=%s, 期望关键词=%s\n", apiURL, expectedKeyword)
	return "success"
}

func sendAlert(jobName, msg string) {
	msg = strings.ReplaceAll(msg, "{", "<")
	msg = strings.ReplaceAll(msg, "}", ">")
	msg = strings.ReplaceAll(msg, "\"", "'")
	alertMessage := fmt.Sprintf("任务：%s, 结果： %s", jobName, msg)
	webhookURL := fmt.Sprintf(config.WebhookURL, config.AccessToken)
	payload := map[string]interface{}{
		"msgtype": "text",
		"text": map[string]string{
			"content": config.Keyword + alertMessage,
		},
	}
	payloadBytes, _ := json.Marshal(payload)

	http.Post(webhookURL, "application/json", strings.NewReader(string(payloadBytes)))
}

func getEnv(key, defaultValue string) string {
	value := os.Getenv(key)
	if value == "" {
		return defaultValue
	}
	return value
}

func getEnvInt(key string, defaultValue int) int {
	value := os.Getenv(key)
	if value == "" {
		return defaultValue
	}
	intValue, err := strconv.Atoi(value)
	if err != nil {
		return defaultValue
	}
	return intValue
}