✨ feat: add notifier (#144)

* ♻️ refactor: email refactor

* ✨ feat: add notifier

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))
+}

common/notify/channel/lark.go

@@ -0,0 +1,149 @@
+package channel
+
+import (
+	"context"
+	"crypto/hmac"
+	"crypto/sha256"
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"one-api/common/requester"
+	"one-api/types"
+	"strconv"
+	"time"
+)
+
+const larkURL = "https://open.feishu.cn/open-apis/bot/v2/hook/"
+
+type Lark struct {
+	token   string
+	secret  string
+	keyWord string
+}
+
+type larkMessage struct {
+	MessageType string          `json:"msg_type"`
+	Timestamp   string          `json:"timestamp,omitempty"`
+	Sign        string          `json:"sign,omitempty"`
+	Card        larkCardContent `json:"card"`
+}
+
+type larkCardContent struct {
+	Config struct {
+		WideScreenMode bool `json:"wide_screen_mode"`
+		EnableForward  bool `json:"enable_forward"`
+	}
+	Elements []larkMessageRequestCardElement `json:"elements"`
+}
+
+type larkMessageRequestCardElementText struct {
+	Content string `json:"content"`
+	Tag     string `json:"tag"`
+}
+
+type larkMessageRequestCardElement struct {
+	Tag  string                            `json:"tag"`
+	Text larkMessageRequestCardElementText `json:"text"`
+}
+
+type larkResponse struct {
+	Code    int    `json:"code"`
+	Message string `json:"msg"`
+}
+
+func NewLark(token, secret string) *Lark {
+	return &Lark{
+		token:  token,
+		secret: secret,
+	}
+}
+
+func NewLarkWithKeyWord(token, keyWord string) *Lark {
+	return &Lark{
+		token:   token,
+		keyWord: keyWord,
+	}
+}
+
+func (l *Lark) Name() string {
+	return "Lark"
+}
+
+func (l *Lark) Send(ctx context.Context, title, message string) error {
+	msg := larkMessage{
+		MessageType: "interactive",
+	}
+
+	if l.keyWord != "" {
+		title = fmt.Sprintf("%s(%s)", title, l.keyWord)
+	}
+
+	msg.Card.Config.WideScreenMode = true
+	msg.Card.Config.EnableForward = true
+	msg.Card.Elements = append(msg.Card.Elements, larkMessageRequestCardElement{
+		Tag: "div",
+		Text: larkMessageRequestCardElementText{
+			Content: fmt.Sprintf("**%s**\n%s", title, message),
+			Tag:     "lark_md",
+		},
+	})
+
+	if l.secret != "" {
+		t := time.Now().Unix()
+		msg.Timestamp = strconv.FormatInt(t, 10)
+		msg.Sign = l.sign(t)
+	}
+
+	uri := larkURL + l.token
+	client := requester.NewHTTPRequester("", larkErrFunc)
+	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()
+
+	larkErr := larkErrFunc(resp)
+	if larkErr != nil {
+		return fmt.Errorf("%s", larkErr.Message)
+	}
+
+	return nil
+
+}
+
+func (l *Lark) sign(timestamp int64) string {
+	//timestamp + key 做sha256, 再进行base64 encode
+	stringToSign := fmt.Sprintf("%v", timestamp) + "\n" + l.secret
+	var data []byte
+	h := hmac.New(sha256.New, []byte(stringToSign))
+	h.Write(data)
+
+	return base64.StdEncoding.EncodeToString(h.Sum(nil))
+}
+
+func larkErrFunc(resp *http.Response) *types.OpenAIError {
+	respMsg := &larkResponse{}
+	err := json.NewDecoder(resp.Body).Decode(respMsg)
+	if err != nil {
+		return nil
+	}
+
+	if respMsg.Code == 0 {
+		return nil
+	}
+
+	return &types.OpenAIError{
+		Message: fmt.Sprintf("send msg err. err msg: %s", respMsg.Message),
+		Type:    "lark_error",
+		Code:    respMsg.Code,
+	}
+}

common/notify/channel/pushdeer.go

@@ -0,0 +1,96 @@
+package channel
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"one-api/common/requester"
+	"one-api/types"
+	"strings"
+)
+
+const pushdeerURL = "https://api2.pushdeer.com"
+
+type Pushdeer struct {
+	url     string
+	pushkey string
+}
+
+type pushdeerMessage struct {
+	Text string `json:"text"`
+	Desp string `json:"desp"`
+	Type string `json:"type"`
+}
+
+type pushdeerResponse struct {
+	Code    int    `json:"code,omitempty"`
+	Error   string `json:"error,omitempty"`
+	Message string `json:"message,omitempty"`
+}
+
+func NewPushdeer(pushkey, url string) *Pushdeer {
+	return &Pushdeer{
+		url:     url,
+		pushkey: pushkey,
+	}
+}
+
+func (p *Pushdeer) Name() string {
+	return "Pushdeer"
+}
+
+func (p *Pushdeer) Send(ctx context.Context, title, message string) error {
+	msg := pushdeerMessage{
+		Text: title,
+		Desp: message,
+		Type: "markdown",
+	}
+
+	url := p.url
+	if url == "" {
+		url = pushdeerURL
+	}
+
+	// 去除最后一个/
+	url = strings.TrimSuffix(url, "/")
+	uri := fmt.Sprintf("%s/message/push?pushkey=%s", url, p.pushkey)
+
+	client := requester.NewHTTPRequester("", pushdeerErrFunc)
+	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
+	}
+
+	respMsg := &pushdeerResponse{}
+	_, errWithOP := client.SendRequest(req, respMsg, false)
+	if errWithOP != nil {
+		return fmt.Errorf("%s", errWithOP.Message)
+	}
+
+	if respMsg.Code != 0 {
+		return fmt.Errorf("send msg err. err msg: %s", respMsg.Error)
+	}
+
+	return nil
+}
+
+func pushdeerErrFunc(resp *http.Response) *types.OpenAIError {
+	respMsg := &pushdeerResponse{}
+	err := json.NewDecoder(resp.Body).Decode(respMsg)
+	if err != nil {
+		return nil
+	}
+
+	if respMsg.Message == "" {
+		return nil
+	}
+
+	return &types.OpenAIError{
+		Message: fmt.Sprintf("send msg err. err msg: %s", respMsg.Message),
+		Type:    "pushdeer_error",
+	}
+}

common/notify/channel/telegram.go

@@ -0,0 +1,114 @@
+package channel
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"one-api/common/requester"
+	"one-api/types"
+)
+
+const telegramURL = "https://api.telegram.org/bot"
+
+type Telegram struct {
+	secret string
+	chatID string
+}
+
+type telegramMessage struct {
+	ChatID    string `json:"chat_id"`
+	Text      string `json:"text"`
+	ParseMode string `json:"parse_mode"`
+}
+
+type telegramResponse struct {
+	Ok          bool   `json:"ok"`
+	Description string `json:"description"`
+}
+
+func NewTelegram(secret string, chatID string) *Telegram {
+	return &Telegram{
+		secret: secret,
+		chatID: chatID,
+	}
+}
+
+func (t *Telegram) Name() string {
+	return "Telegram"
+}
+
+func (t *Telegram) Send(ctx context.Context, title, message string) error {
+	const maxMessageLength = 4096
+	message = fmt.Sprintf("*%s*\n%s", title, message)
+	messages := splitTelegramMessageIntoParts(message, maxMessageLength)
+
+	client := requester.NewHTTPRequester("", telegramErrFunc)
+	client.Context = ctx
+	client.IsOpenAI = false
+
+	for _, msg := range messages {
+		err := t.sendMessage(msg, client)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (t *Telegram) sendMessage(message string, client *requester.HTTPRequester) error {
+	msg := telegramMessage{
+		ChatID:    t.chatID,
+		Text:      message,
+		ParseMode: "Markdown",
+	}
+
+	uri := telegramURL + t.secret + "/sendMessage"
+
+	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()
+
+	telegramErr := telegramErrFunc(resp)
+	if telegramErr != nil {
+		return fmt.Errorf("%s", telegramErr.Message)
+	}
+
+	return nil
+}
+
+func splitTelegramMessageIntoParts(message string, partSize int) []string {
+	var parts []string
+	for len(message) > partSize {
+		parts = append(parts, message[:partSize])
+		message = message[partSize:]
+	}
+	parts = append(parts, message)
+
+	return parts
+}
+
+func telegramErrFunc(resp *http.Response) *types.OpenAIError {
+	respMsg := &telegramResponse{}
+	err := json.NewDecoder(resp.Body).Decode(respMsg)
+	if err != nil {
+		return nil
+	}
+
+	if respMsg.Ok {
+		return nil
+	}
+
+	return &types.OpenAIError{
+		Message: fmt.Sprintf("send msg err. err msg: %s", respMsg.Description),
+		Type:    "telegram_error",
+	}
+}

common/notify/notifier.go

@@ -0,0 +1,98 @@
+package notify
+
+import (
+	"context"
+	"one-api/common"
+	"one-api/common/notify/channel"
+
+	"github.com/spf13/viper"
+)
+
+type Notifier interface {
+	Send(context.Context, string, string) error
+	Name() string
+}
+
+func InitNotifier() {
+	InitEmailNotifier()
+	InitDingTalkNotifier()
+	InitLarkNotifier()
+	InitPushdeerNotifier()
+	InitTelegramNotifier()
+}
+
+func InitEmailNotifier() {
+	if viper.GetBool("notify.email.disable") {
+		common.SysLog("email notifier disabled")
+		return
+	}
+	smtp_to := viper.GetString("notify.email.smtp_to")
+	emailNotifier := channel.NewEmail(smtp_to)
+	AddNotifiers(emailNotifier)
+	common.SysLog("email notifier enable")
+}
+
+func InitDingTalkNotifier() {
+	access_token := viper.GetString("notify.dingtalk.token")
+	secret := viper.GetString("notify.dingtalk.secret")
+	keyWord := viper.GetString("notify.dingtalk.keyWord")
+	if access_token == "" || (secret == "" && keyWord == "") {
+		return
+	}
+
+	var dingTalkNotifier Notifier
+
+	if secret != "" {
+		dingTalkNotifier = channel.NewDingTalk(access_token, secret)
+	} else {
+		dingTalkNotifier = channel.NewDingTalkWithKeyWord(access_token, keyWord)
+	}
+
+	AddNotifiers(dingTalkNotifier)
+	common.SysLog("dingtalk notifier enable")
+}
+
+func InitLarkNotifier() {
+	access_token := viper.GetString("notify.lark.token")
+	secret := viper.GetString("notify.lark.secret")
+	keyWord := viper.GetString("notify.lark.keyWord")
+	if access_token == "" || (secret == "" && keyWord == "") {
+		return
+	}
+
+	var larkNotifier Notifier
+
+	if secret != "" {
+		larkNotifier = channel.NewLark(access_token, secret)
+	} else {
+		larkNotifier = channel.NewLarkWithKeyWord(access_token, keyWord)
+	}
+
+	AddNotifiers(larkNotifier)
+	common.SysLog("lark notifier enable")
+}
+
+func InitPushdeerNotifier() {
+	pushkey := viper.GetString("notify.pushdeer.pushkey")
+	if pushkey == "" {
+		return
+	}
+
+	pushdeerNotifier := channel.NewPushdeer(pushkey, viper.GetString("notify.pushdeer.url"))
+
+	AddNotifiers(pushdeerNotifier)
+	common.SysLog("pushdeer notifier enable")
+}
+
+func InitTelegramNotifier() {
+	bot_token := viper.GetString("notify.telegram.bot_api_key")
+	chat_id := viper.GetString("notify.telegram.chat_id")
+	if bot_token == "" || chat_id == "" {
+		return
+	}
+
+	telegramNotifier := channel.NewTelegram(bot_token, chat_id)
+
+	AddNotifiers(telegramNotifier)
+	common.SysLog("telegram notifier enable")
+}

common/notify/notify.go

@@ -0,0 +1,35 @@
+package notify
+
+var notifyChannels = New()
+
+type Notify struct {
+	notifiers map[string]Notifier
+}
+
+func (n *Notify) addChannel(channel Notifier) {
+	if channel != nil {
+		channelName := channel.Name()
+		if _, ok := n.notifiers[channelName]; ok {
+			return
+		}
+		n.notifiers[channelName] = channel
+	}
+}
+
+func (n *Notify) addChannels(channel ...Notifier) {
+	for _, s := range channel {
+		n.addChannel(s)
+	}
+}
+
+func New() *Notify {
+	notify := &Notify{
+		notifiers: make(map[string]Notifier, 0),
+	}
+
+	return notify
+}
+
+func AddNotifiers(channel ...Notifier) {
+	notifyChannels.addChannels(channel...)
+}

common/notify/send.go

@@ -0,0 +1,30 @@
+package notify
+
+import (
+	"context"
+	"fmt"
+	"one-api/common"
+)
+
+func (n *Notify) Send(ctx context.Context, title, message string) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	for channelName, channel := range n.notifiers {
+		if channel == nil {
+			continue
+		}
+		err := channel.Send(ctx, title, message)
+		if err != nil {
+			common.LogError(ctx, fmt.Sprintf("%s err: %s", channelName, err.Error()))
+		}
+	}
+}
+
+func Send(title, message string) {
+	//lint:ignore SA1029 reason: 需要使用该类型作为错误处理
+	ctx := context.WithValue(context.Background(), common.RequestIdKey, "NotifyTask")
+
+	notifyChannels.Send(ctx, title, message)
+}