aiops/backend/pkg/siteprobe/executor.go

package siteprobe

import (
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"

	"ezops/model"
)

// ProbeOutcome 单次拨测结果 // Logan.Li
type ProbeOutcome struct {
	Status       model.SiteResultStatus
	HTTPStatus   int
	ResponseTime float64
	TotalTime    float64
	RemoteIP     string
	ErrorMessage string
	// HTTPS 证书信息 // Logan.Li
	SSLCertExpiry    *time.Time
	SSLCertIssuer    string
	SSLCertSubject   string
	SSLDaysRemaining int
}

// Execute 按任务类型执行拨测 // Logan.Li
func Execute(ctx context.Context, task model.SiteTask) ProbeOutcome {
	timeout := time.Duration(task.Timeout) * time.Second
	if timeout <= 0 {
		timeout = 10 * time.Second
	}
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	taskType := strings.ToLower(string(task.TaskType))
	if taskType == "" {
		taskType = inferTaskType(task)
	}

	switch taskType {
	case "ssl":
		return probeSSL(ctx, task)
	case "tcp":
		return probeTCP(ctx, task.URL)
	case "dns":
		return probeDNS(ctx, task.URL, task.DNSResolver)
	case "ping":
		return probePing(ctx, task.URL, task.Timeout)
	case "http", "https":
		return probeHTTP(ctx, task)
	default:
		return probeHTTP(ctx, task)
	}
}

// InferTaskTypePublic 推断任务类型（供列表修复存量数据）// Logan.Li
func InferTaskTypePublic(task model.SiteTask) string {
	return inferTaskType(task)
}

func inferTaskType(task model.SiteTask) string {
	name := strings.ToLower(task.Name)
	switch {
	case strings.Contains(name, "tcp"):
		return "tcp"
	case strings.Contains(name, "dns"):
		return "dns"
	case strings.Contains(name, "ping"):
		return "ping"
	case strings.Contains(name, "ssl"):
		return "ssl"
	default:
		return "http"
	}
}

func probeHTTP(ctx context.Context, task model.SiteTask) ProbeOutcome {
	start := time.Now()
	url := task.URL
	if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") {
		url = "http://" + url
	}

	transport := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: !task.VerifySSL},
	}
	client := &http.Client{
		Timeout:   time.Until(deadlineFromCtx(ctx)),
		Transport: transport,
	}
	if !task.FollowRedirect {
		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		}
	}

	method := strings.ToUpper(string(task.Method))
	if method == "" {
		method = http.MethodGet
	}
	req, err := http.NewRequestWithContext(ctx, method, url, nil)
	if err != nil {
		return failOutcome(start, err.Error())
	}

	resp, err := client.Do(req)
	elapsed := msSince(start)
	if err != nil {
		if ctx.Err() == context.DeadlineExceeded {
			return ProbeOutcome{Status: model.SiteResultTimeout, TotalTime: elapsed, ResponseTime: elapsed, ErrorMessage: "请求超时"}
		}
		return failOutcome(start, err.Error())
	}
	defer resp.Body.Close()

	statusOK := resp.StatusCode == task.ExpectedStatus
	if task.ExpectedStatus == 0 {
		statusOK = resp.StatusCode >= 200 && resp.StatusCode < 400
	}
	if !statusOK {
		return ProbeOutcome{
			Status:       model.SiteResultFail,
			HTTPStatus:   resp.StatusCode,
			ResponseTime: elapsed,
			TotalTime:    elapsed,
			ErrorMessage: fmt.Sprintf("状态码 %d 不符合期望 %d", resp.StatusCode, task.ExpectedStatus),
		}
	}
	return ProbeOutcome{
		Status:       model.SiteResultSuccess,
		HTTPStatus:   resp.StatusCode,
		ResponseTime: elapsed,
		TotalTime:    elapsed,
	}
}

func probeTCP(ctx context.Context, target string) ProbeOutcome {
	start := time.Now()
	host, port, err := splitHostPort(target, "80")
	if err != nil {
		return failOutcome(start, err.Error())
	}
	d := net.Dialer{Timeout: time.Until(deadlineFromCtx(ctx))}
	conn, err := d.DialContext(ctx, "tcp", net.JoinHostPort(host, port))
	elapsed := msSince(start)
	if err != nil {
		if ctx.Err() == context.DeadlineExceeded {
			return ProbeOutcome{Status: model.SiteResultTimeout, TotalTime: elapsed, ResponseTime: elapsed, ErrorMessage: "TCP 连接超时"}
		}
		return failOutcome(start, err.Error())
	}
	defer conn.Close()
	remote := ""
	if addr := conn.RemoteAddr(); addr != nil {
		remote = addr.String()
	}
	return ProbeOutcome{
		Status:       model.SiteResultSuccess,
		ResponseTime: elapsed,
		TotalTime:    elapsed,
		RemoteIP:     remote,
	}
}

func probeDNS(ctx context.Context, target, resolver string) ProbeOutcome {
	start := time.Now()
	domain := normalizeDomain(target)
	if domain == "" {
		return failOutcome(start, "域名不能为空")
	}
	r := &net.Resolver{}
	if resolver != "" {
		r = &net.Resolver{
			PreferGo: true,
			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
				d := net.Dialer{Timeout: time.Until(deadlineFromCtx(ctx))}
				return d.DialContext(ctx, network, resolver+":53")
			},
		}
	}
	addrs, err := r.LookupHost(ctx, domain)
	elapsed := msSince(start)
	if err != nil {
		if ctx.Err() == context.DeadlineExceeded {
			return ProbeOutcome{Status: model.SiteResultTimeout, TotalTime: elapsed, ResponseTime: elapsed, ErrorMessage: "DNS 解析超时"}
		}
		return failOutcome(start, err.Error())
	}
	return ProbeOutcome{
		Status:       model.SiteResultSuccess,
		ResponseTime: elapsed,
		TotalTime:    elapsed,
		RemoteIP:     strings.Join(addrs, ","),
	}
}

func probePing(ctx context.Context, target string, timeoutSec int) ProbeOutcome {
	start := time.Now()
	host := normalizeDomain(target)
	if host == "" {
		return failOutcome(start, "主机不能为空")
	}
	if timeoutSec <= 0 {
		timeoutSec = 5
	}
	cmd := exec.CommandContext(ctx, "ping", "-c", "1", "-W", strconv.Itoa(timeoutSec), host)
	out, err := cmd.CombinedOutput()
	elapsed := msSince(start)
	if err != nil {
		msg := strings.TrimSpace(string(out))
		if msg == "" {
			msg = err.Error()
		}
		if ctx.Err() == context.DeadlineExceeded {
			return ProbeOutcome{Status: model.SiteResultTimeout, TotalTime: elapsed, ResponseTime: elapsed, ErrorMessage: "PING 超时"}
		}
		return ProbeOutcome{Status: model.SiteResultFail, TotalTime: elapsed, ResponseTime: elapsed, ErrorMessage: msg}
	}
	if m := regexp.MustCompile(`time[=<]([\d.]+)\s*ms`).FindStringSubmatch(string(out)); len(m) == 2 {
		if v, parseErr := strconv.ParseFloat(m[1], 64); parseErr == nil {
			elapsed = v
		}
	}
	return ProbeOutcome{Status: model.SiteResultSuccess, ResponseTime: elapsed, TotalTime: elapsed, RemoteIP: host}
}

func splitHostPort(raw, defaultPort string) (string, string, error) {
	raw = strings.TrimSpace(raw)
	for _, prefix := range []string{"ssl://", "tcp://", "http://", "https://"} {
		raw = strings.TrimPrefix(raw, prefix)
	}
	if host, port, err := net.SplitHostPort(raw); err == nil {
		return host, port, nil
	}
	if raw == "" {
		return "", "", fmt.Errorf("TCP 地址不能为空")
	}
	return raw, defaultPort, nil
}

func normalizeDomain(raw string) string {
	raw = strings.TrimSpace(raw)
	raw = strings.TrimPrefix(raw, "http://")
	raw = strings.TrimPrefix(raw, "https://")
	if host, _, err := net.SplitHostPort(raw); err == nil {
		return host
	}
	return strings.TrimSuffix(raw, "/")
}

func failOutcome(start time.Time, msg string) ProbeOutcome {
	elapsed := msSince(start)
	return ProbeOutcome{
		Status:       model.SiteResultFail,
		ResponseTime: elapsed,
		TotalTime:    elapsed,
		ErrorMessage: msg,
	}
}

func msSince(start time.Time) float64 {
	return float64(time.Since(start).Milliseconds())
}

func deadlineFromCtx(ctx context.Context) time.Time {
	if d, ok := ctx.Deadline(); ok {
		return d
	}
	return time.Now().Add(10 * time.Second)
}