webhook-tools/scripts/notify/main.go

package main

import (
	"bytes"
	"context"
	"encoding/base64"
	"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)

		// title_value := getEnvStr("TITLE", title)
		// if title_value == "" {
		//      title = "获取不到title或环境变量TITLE"
		// }
		t_encode := decodeBase64(title)
		title = t_encode

		m_encode := decodeBase64(msg)
		msg = m_encode

		// 如果配置文件中提供飞书链接，则通知用户
		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
// }

func decodeBase64(str string) string {
	data, err := base64.StdEncoding.DecodeString(str)
	if err != nil {
		return str
	}
	reencode := base64.StdEncoding.EncodeToString(data)
	if reencode != str {
		return str
	}
	return string(data)
}

// func getEnvStr(key, defaultValue string) string {
//      v, ok := os.LookupEnv(key)
//      if !ok {
//              return defaultValue
//      }
//      value := strings.TrimSpace(v)
//      if value == "" {
//              return ""
//      }
//      return value

// }