Merge pull request #62 from woodchen-ink/MartialBE-main

Martial be main

common/notify/channel/channel_test.go

@@ -0,0 +1,128 @@
+package channel_test
+
+import (
+ "context"
+ "fmt"
+ "testing"
+
+ "one-api/common/notify/channel"
+ "one-api/common/requester"
+
+ "github.com/spf13/viper"
+ "github.com/stretchr/testify/assert"
+)
+
+func InitConfig() {
+ viper.AddConfigPath("/one-api")
+ viper.SetConfigName("config")
+ viper.ReadInConfig()
+ requester.InitHttpClient()
+}
+
+func TestDingTalkSend(t *testing.T) {
+ InitConfig()
+ access_token := viper.GetString("notify.dingtalk.token")
+ secret := viper.GetString("notify.dingtalk.secret")
+ dingTalk := channel.NewDingTalk(access_token, secret)
+
+ err := dingTalk.Send(context.Background(), "Test Title", "*Test Message*")
+ fmt.Println(err)
+ assert.Nil(t, err)
+}
+
+func TestDingTalkSendWithKeyWord(t *testing.T) {
+ InitConfig()
+ access_token := viper.GetString("notify.dingtalk.token")
+ keyWord := viper.GetString("notify.dingtalk.keyWord")
+
+ dingTalk := channel.NewDingTalkWithKeyWord(access_token, keyWord)
+
+ err := dingTalk.Send(context.Background(), "Test Title", "Test Message")
+ assert.Nil(t, err)
+}
+
+func TestDingTalkSendError(t *testing.T) {
+ InitConfig()
+ access_token := viper.GetString("notify.dingtalk.token")
+ secret := "test"
+ dingTalk := channel.NewDingTalk(access_token, secret)
+
+ err := dingTalk.Send(context.Background(), "Test Title", "*Test Message*")
+ fmt.Println(err)
+ assert.Error(t, err)
+}
+
+func TestLarkSend(t *testing.T) {
+ InitConfig()
+ access_token := viper.GetString("notify.lark.token")
+ secret := viper.GetString("notify.lark.secret")
+ dingTalk := channel.NewLark(access_token, secret)
+
+ err := dingTalk.Send(context.Background(), "Test Title", "*Test Message*")
+ fmt.Println(err)
+ assert.Nil(t, err)
+}
+
+func TestLarkSendWithKeyWord(t *testing.T) {
+ InitConfig()
+ access_token := viper.GetString("notify.lark.token")
+ keyWord := viper.GetString("notify.lark.keyWord")
+
+ dingTalk := channel.NewLarkWithKeyWord(access_token, keyWord)
+
+ err := dingTalk.Send(context.Background(), "Test Title", "Test Message\n\n- 111\n- 222")
+ assert.Nil(t, err)
+}
+
+func TestLarkSendError(t *testing.T) {
+ InitConfig()
+ access_token := viper.GetString("notify.lark.token")
+ secret := "test"
+ dingTalk := channel.NewLark(access_token, secret)
+
+ err := dingTalk.Send(context.Background(), "Title", "*Message*")
+ fmt.Println(err)
+ assert.Error(t, err)
+}
+
+func TestPushdeerSend(t *testing.T) {
+ InitConfig()
+ pushkey := viper.GetString("notify.pushdeer.pushkey")
+ dingTalk := channel.NewPushdeer(pushkey, "")
+
+ err := dingTalk.Send(context.Background(), "Test Title", "*Test Message*")
+ fmt.Println(err)
+ assert.Nil(t, err)
+}
+
+func TestPushdeerSendError(t *testing.T) {
+ InitConfig()
+ pushkey := "test"
+ dingTalk := channel.NewPushdeer(pushkey, "")
+
+ err := dingTalk.Send(context.Background(), "Test Title", "*Test Message*")
+ fmt.Println(err)
+ assert.Error(t, err)
+}
+
+func TestTelegramSend(t *testing.T) {
+ InitConfig()
+ secret := viper.GetString("notify.telegram.bot_api_key")
+ chatID := viper.GetString("notify.telegram.chat_id")
+ dingTalk := channel.NewTelegram(secret, chatID)
+
+ err := dingTalk.Send(context.Background(), "Test Title", "*Test Message*")
+ fmt.Println(err)
+ assert.Nil(t, err)
+}
+
+func TestTelegramSendError(t *testing.T) {
+ InitConfig()
+ secret := "test"
+ chatID := viper.GetString("notify.telegram.chat_id")
+ dingTalk := channel.NewTelegram(secret, chatID)
+
+ err := dingTalk.Send(context.Background(), "Test Title", "*Test Message*")
+ fmt.Println(err)
+ assert.Error(t, err)
+}

common/notify/channel/dingTalk.go

@@ -0,0 +1,127 @@
+package channel
+
+import (
+ "context"
+ "crypto/hmac"
+ "crypto/sha256"
+ "encoding/base64"
+ "encoding/json"
+ "fmt"
+ "net/http"
+ "net/url"
+ "one-api/common/requester"
+ "one-api/types"
+ "time"
+)
+
+const dingTalkURL = "https://oapi.dingtalk.com/robot/send?"
+
+type DingTalk struct {
+ token string
+ secret string
+ keyWord string
+}
+
+type dingTalkMessage struct {
+ MsgType string `json:"msgtype"`
+ Markdown struct {
+ Title string `json:"title"`
+ Text string `json:"text"`
+ } `json:"markdown"`
+}
+
+type dingTalkResponse struct {
+ ErrCode int `json:"errcode"`
+ ErrMsg string `json:"errmsg"`
+}
+
+func NewDingTalk(token string, secret string) *DingTalk {
+ return &DingTalk{
+ token: token,
+ secret: secret,
+ }
+}
+
+func NewDingTalkWithKeyWord(token string, keyWord string) *DingTalk {
+ return &DingTalk{
+ token: token,
+ keyWord: keyWord,
+ }
+}
+
+func (d *DingTalk) Name() string {
+ return "DingTalk"
+}
+
+func (d *DingTalk) Send(ctx context.Context, title, message string) error {
+ msg := dingTalkMessage{
+ MsgType: "markdown",
+ }
+ msg.Markdown.Title = title
+ msg.Markdown.Text = message
+
+ if d.keyWord != "" {
+ msg.Markdown.Text = fmt.Sprintf("%s\n%s", d.keyWord, msg.Markdown.Text)
+ }
+
+ query := url.Values{}
+ query.Set("access_token", d.token)
+ if d.secret != "" {
+ t := time.Now().UnixMilli()
+ query.Set("timestamp", fmt.Sprintf("%d", t))
+ query.Set("sign", d.sign(t))
+ }
+ uri := dingTalkURL + query.Encode()
+
+ client := requester.NewHTTPRequester("", dingtalkErrFunc)
+ client.Context = ctx
+ client.IsOpenAI = false
+
+ req, err := client.NewRequest(http.MethodPost, uri, client.WithHeader(requester.GetJsonHeaders()), client.WithBody(msg))
+ if err != nil {
+ return err
+ }
+
+ resp, errWithOP := client.SendRequestRaw(req)
+ if errWithOP != nil {
+ return fmt.Errorf("%s", errWithOP.Message)
+ }
+ defer resp.Body.Close()
+
+ dingtalkErr := dingtalkErrFunc(resp)
+ if dingtalkErr != nil {
+ return fmt.Errorf("%s", dingtalkErr.Message)
+ }
+
+ return nil
+}
+
+func (d *DingTalk) sign(timestamp int64) string {
+ stringToHash := fmt.Sprintf("%d\n%s", timestamp, d.secret)
+ hmac256 := hmac.New(sha256.New, []byte(d.secret))
+ hmac256.Write([]byte(stringToHash))
+ data := hmac256.Sum(nil)
+ signature := base64.StdEncoding.EncodeToString(data)
+
+ return url.QueryEscape(signature)
+}
+
+func dingtalkErrFunc(resp *http.Response) *types.OpenAIError {
+ respMsg := &dingTalkResponse{}
+
+ err := json.NewDecoder(resp.Body).Decode(respMsg)
+ if err != nil {
+ fmt.Println(err)
+ return nil
+ }
+
+ if respMsg.ErrCode == 0 {
+ return nil
+ }
+
+ return &types.OpenAIError{
+ Message: fmt.Sprintf("send msg err. err msg: %s", respMsg.ErrMsg),
+ Type: "dingtalk_error",
+ Code: fmt.Sprintf("%d", respMsg.ErrCode),
+ }
+}

common/notify/channel/email.go

@@ -0,0 +1,50 @@
+package channel
+
+import (
+ "context"
+ "errors"
+ "one-api/common"
+ "one-api/common/stmp"
+
+ "github.com/gomarkdown/markdown"
+ "github.com/gomarkdown/markdown/html"
+ "github.com/gomarkdown/markdown/parser"
+)
+
+type Email struct {
+ To string
+}
+
+func NewEmail(to string) *Email {
+ return &Email{
+ To: to,
+ }
+}
+
+func (e *Email) Name() string {
+ return "Email"
+}
+
+func (e *Email) Send(ctx context.Context, title, message string) error {
+ to := e.To
+ if to == "" {
+ to = common.RootUserEmail
+ }
+
+ if common.SMTPServer == "" || common.SMTPAccount == "" || common.SMTPToken == "" || to == "" {
+ return errors.New("smtp config is not set, skip send email notifier")
+ }
+
+ p := parser.NewWithExtensions(parser.CommonExtensions | parser.DefinitionLists | parser.OrderedListStart)
+ doc := p.Parse([]byte(message))
+
+ htmlFlags := html.CommonFlags | html.HrefTargetBlank
+ opts := html.RendererOptions{Flags: htmlFlags}
+ renderer := html.NewRenderer(opts)
+
+ body := markdown.Render(doc, renderer)
+
+ emailClient := stmp.NewStmp(common.SMTPServer, common.SMTPPort, common.SMTPAccount, common.SMTPToken, common.SMTPFrom)
+
+ return emailClient.Send(to, title, string(body))
+}