add interface and fix

scripts/notify/main.go

@@ -0,0 +1,241 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/credentials"
+	"github.com/aws/aws-sdk-go-v2/service/s3"
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+)
+
+var (
+	cfgFile string
+	title   string
+	msg     string
+	file    string
+	notice  bool = false
+)
+
+var rootCmd = &cobra.Command{
+	Use:   "notify",
+	Short: "notify by jenkins",
+	Run: func(cmd *cobra.Command, args []string) {
+		// 如果配置文件存在，则读取配置文件
+		if cfgFile != "" {
+			viper.SetConfigFile(cfgFile)
+			if err := viper.ReadInConfig(); err != nil {
+				fmt.Println("Error reading config file:", err)
+				os.Exit(1)
+			}
+		}
+
+		// 获取配置
+		ak := viper.GetString("s3.ak")
+		sk := viper.GetString("s3.sk")
+		notify := viper.GetString("notify")
+		s3_endpoint := viper.GetString("s3.endpoint")
+		s3_bucket := viper.GetString("s3.bucket")
+		s3_path := viper.GetString("s3.path")
+		s3_public_url := viper.GetString("s3.public_url")
+
+		file_type := viper.GetString("type")
+
+		FILE := os.Getenv("FILE")
+		flist := strings.Split(FILE, "/")
+		f := fmt.Sprintf("%s.%s", flist[len(flist)-1], file_type)
+
+		if title == "" {
+			if os.Getenv("TITLE") != "" {
+				title, _ = readFile(os.Getenv("TITLE"))
+			} else {
+				title = "获取不到title或环境变量TITLE"
+			}
+		}
+		if msg == "" {
+			if os.Getenv("MSG") != "" {
+				msg, _ = readFile(os.Getenv("MSG"))
+			} else {
+				msg = "获取不到msg或环境变量MSG"
+			}
+		}
+
+		// 如果配置文件中提供飞书链接，则通知用户
+		if notify != "" {
+			// 是否通知
+			notice = true
+		}
+		//发送用户格式，默认是text文本
+		post := false
+		// 富文本中的链接
+		postUrl := ""
+
+		if FILE != "" {
+			post = true
+			input_key := fmt.Sprintf("%s/%s/%s", s3_path, time.Now().Format("201601021504"), f)
+			objectUrl, err := s3PutObject(ak, sk, s3_endpoint, s3_bucket, input_key, FILE, s3_public_url)
+			if err != nil {
+				fmt.Println("Error sending message:", err)
+				os.Exit(1)
+			}
+			postUrl = objectUrl
+		}
+
+		if notice {
+			var data interface{}
+			if post {
+				data = genPostAMsg(title, msg, "查看", postUrl)
+
+			} else {
+				data = FeishuTextMessage{
+					MsgType: "text",
+					Content: TextContent{
+						Text: fmt.Sprintf("%s\n%s", title, msg),
+					},
+				}
+
+			}
+			if err := postRequest(notify, data); err != nil {
+				fmt.Println("Error sending message:", err)
+				os.Exit(1)
+			}
+
+		}
+
+	},
+}
+
+func main() {
+	cobra.OnInitialize(initConfig)
+	rootCmd.Flags().StringVarP(&cfgFile, "config", "c", "", "Config file (default is $HOME/config.yaml and ./config.yaml)")
+	rootCmd.Flags().StringVarP(&title, "title", "t", "", "title")
+	rootCmd.Flags().StringVarP(&msg, "msg", "m", "", "message")
+	rootCmd.Flags().StringVarP(&file, "file", "f", "", "file")
+
+	if err := rootCmd.Execute(); err != nil {
+		fmt.Println(err)
+		os.Exit(1)
+	}
+}
+
+func initConfig() {
+	if cfgFile != "" {
+		viper.SetConfigFile(cfgFile)
+	} else {
+		home, err := os.UserHomeDir()
+		if err != nil {
+			fmt.Println(err)
+			os.Exit(1)
+		}
+		exec, err := os.Executable()
+		if err != nil {
+			fmt.Println(err)
+			os.Exit(1)
+		}
+		exPath := filepath.Dir(exec)
+		viper.AddConfigPath(".")
+		viper.AddConfigPath(home)
+		viper.AddConfigPath(exPath)
+		viper.SetConfigName("config")
+	}
+
+	viper.AutomaticEnv() // 从环境变量中读取配置
+	if err := viper.ReadInConfig(); err == nil {
+		fmt.Println("Using config file:", viper.ConfigFileUsed())
+	}
+}
+
+func postRequest(url string, data interface{}) error {
+	jsonData, err := json.Marshal(data)
+	if err != nil {
+		return err
+	}
+	fmt.Printf("post data: %s\n", string(jsonData))
+	resp, err := http.Post(url, "application/json", bytes.NewBuffer(jsonData))
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	fmt.Printf("Response status: %d\n", resp.StatusCode)
+
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("HTTP status code: %d", resp.StatusCode)
+	}
+	return nil
+}
+
+func genPostAMsg(title, message, text, url string) *FeishuPostMessage {
+	return &FeishuPostMessage{
+		MsgType: "post",
+		Content: PostContent{
+			Post: Post{
+				ZhCn: Message{
+					Title: title,
+					Content: [][]ContentItem{
+						{
+							{
+								Tag:  "text",
+								Text: message,
+							},
+							{
+								Tag:  "a",
+								Text: text,
+								Href: url,
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+}
+
+func s3PutObject(ak, sk, endpoint, bucket, inputKey, filePath, url string) (string, error) {
+	cfg, err := config.LoadDefaultConfig(context.TODO(),
+
+		config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(ak, sk, "")),
+		config.WithRegion("auto"),
+	)
+	if err != nil {
+		return "", err
+	}
+	client := s3.NewFromConfig(cfg, func(o *s3.Options) {
+		o.BaseEndpoint = aws.String(endpoint)
+	})
+
+	fd, err := os.Open(filePath)
+	if err != nil {
+		return "", err
+	}
+	defer fd.Close()
+
+	_, err = client.PutObject(context.TODO(), &s3.PutObjectInput{
+		Bucket:      aws.String(bucket),
+		Key:         aws.String(inputKey),
+		Body:        fd,
+		ContentType: aws.String("text/html; charset=utf-8"),
+	})
+	if err == nil {
+		return fmt.Sprintf("%s/%s", url, inputKey), err
+	}
+	return "", err
+}
+
+func readFile(path string) (string, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return "", err
+	}
+	return string(data), nil
+}