使用go写的一个api接口

记录一下使用go写的一些脚本

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	requestCount  int
	lastTimestamp time.Time
	mu            sync.Mutex
)

const (
	maxRequests   = 10
	resetDuration = time.Minute
	logFilePath   = "server.log"
)

type RequestBody struct {
	Cluster string `json:"cluster"`
	Exec    string `json:"exec"`
	Type    string `json:"type"`
	Ip      string `json:"ip"`
}

func getippath(ip string) bool {
	cmd := exec.Command("get_service_by_host", "-i", ip)

	out, err := cmd.Output()
	if err != nil {
		fmt.Println(err)
		return false
	}

	lines := strings.Split(string(out), "\n")
	for _, line := range lines {
		instance := strings.TrimSpace(line)
		if strings.Contains(instance, "k8s-eks-native-guanxing-node.K8S.all") ||
			strings.Contains(instance, "IM-GPU.SC-GPU.all") {
			return true
		}
	}
	return false
}

func FreezedHost(cluster, execCmd, freezereason, ip string) (string, error) {
	if execCmd == "freeze" {
		execCmd = "freeze"
	} else if execCmd == "unfreeze" {
		execCmd = "unfreeze"
	} else {
		return "", fmt.Errorf("error is not unfreeze or freeze")
	}

	if !getippath(ip) {
		return "", fmt.Errorf("ip is not in k8s-eks-native-guanxing-node.K8S.all or IM-GPU.SC-GPU.all")
	}

	cmd := exec.Command("matrix", "-c", cluster, "host", execCmd, "--reason", freezereason, ip)

	log.Printf("执行命令: %s", cmd.String())

	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", err
	}

	return string(out), nil
}

func fetchEtcdMetricsHandler(w http.ResponseWriter, r *http.Request) {
	mu.Lock()
	defer mu.Unlock()

	now := time.Now()
	if now.Sub(lastTimestamp) > resetDuration {
		requestCount = 0
		lastTimestamp = now
	}

	if requestCount >= maxRequests {
		http.Error(w, "请求频率过高", http.StatusTooManyRequests)
		return
	}

	requestCount++

	clientIP := strings.Split(r.RemoteAddr, ":")[0]

	var requestBody RequestBody
	err := json.NewDecoder(r.Body).Decode(&requestBody)
	if err != nil {
		http.Error(w, "Invalid request body", http.StatusBadRequest)
		return
	}

	cluster := requestBody.Cluster
	execCmd := requestBody.Exec
	freezeReason := requestBody.Type
	ip := requestBody.Ip
        log.Printf("客户端IP: %s - Received params: %s, %s, %s, %s\n", clientIP, cluster, execCmd, freezeReason, ip)
	fmt.Printf("Received params: %s, %s, %s, %s\n", cluster, execCmd, freezeReason, ip)

	cmd, err := FreezedHost(cluster, execCmd, freezeReason, ip)
	fmt.Printf("Received params from client %s: %s, %s, %s, %s\n", clientIP, cluster, execCmd, freezeReason, ip)
	if err != nil {
		log.Printf("发生错误：%v", err)
		http.Error(w, fmt.Sprintf("发生错误：%v", err), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "text/plain")
	fmt.Fprint(w, cmd)
}

func main() {
	logFile, err := os.OpenFile(logFilePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		log.Fatalf("无法打开日志文件: %v", err)
	}
	defer logFile.Close()

	log.SetOutput(logFile)
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

	http.HandleFunc("/matrix_free_api", fetchEtcdMetricsHandler)

	portStr := os.Getenv("SERVER_PORT")
	port := 6999
	if portStr != "" {
		p, err := strconv.Atoi(portStr)
		if err == nil {
			port = p
		}
	}

	fmt.Printf("服务器正在运行，地址：http://0.0.0.0:%d\n", port)
	err = http.ListenAndServe(fmt.Sprintf(":%d", port), nil)
	if err != nil {
		log.Fatalf("启动服务器时发生错误：%v", err)
	}
}