From 8f98cf650ede5b760a1e79fbeb2350b372d94ad9 Mon Sep 17 00:00:00 2001 From: lychiyu Date: Fri, 30 Aug 2024 23:03:08 +0800 Subject: [PATCH] =?UTF-8?q?[add]=20=E6=B7=BB=E5=8A=A0lark=E9=80=9A?= =?UTF-8?q?=E7=9F=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pkg/bbgo/environment.go | 16 ++ pkg/cmd/root.go | 2 + pkg/notifier/larknotifier/lark.go | 196 +++++++++++++++++++++++ pkg/notifier/larknotifier/logrus_look.go | 45 ++++++ 4 files changed, 259 insertions(+) create mode 100644 pkg/notifier/larknotifier/lark.go create mode 100644 pkg/notifier/larknotifier/logrus_look.go diff --git a/pkg/bbgo/environment.go b/pkg/bbgo/environment.go index 917dd7e..17eb1d6 100644 --- a/pkg/bbgo/environment.go +++ b/pkg/bbgo/environment.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "fmt" + "git.qtrade.icu/lychiyu/bbgo/pkg/notifier/larknotifier" "image/png" "io/ioutil" stdlog "log" @@ -688,6 +689,12 @@ func (environ *Environment) ConfigureNotificationSystem(ctx context.Context, use } } + larkBotToken := viper.GetString("lark-bot-token") + if len(larkBotToken) > 0 { + if err := environ.setupLark(userConfig, larkBotToken, persistence); err != nil { + } + } + if userConfig.Notifications != nil { if err := environ.ConfigureNotification(userConfig.Notifications); err != nil { return err @@ -959,6 +966,15 @@ func (environ *Environment) setupTelegram( return nil } +func (environ *Environment) setupLark(userConfig *Config, larkBotToken string, persistence service.PersistenceService, +) error { + + var notifier = larknotifier.New(larkBotToken) + Notification.AddNotifier(notifier) + log.AddHook(larknotifier.LarkNewLogHook(notifier)) + return nil +} + func writeOTPKeyAsQRCodePNG(key *otp.Key, imagePath string) error { // Convert TOTP key into a PNG var buf bytes.Buffer diff --git a/pkg/cmd/root.go b/pkg/cmd/root.go index e73b911..6f355c5 100644 --- a/pkg/cmd/root.go +++ b/pkg/cmd/root.go @@ -190,6 +190,8 @@ func init() { RootCmd.PersistentFlags().String("telegram-bot-token", "", "telegram bot token from bot father") RootCmd.PersistentFlags().String("telegram-bot-auth-token", "", "telegram auth token") + RootCmd.PersistentFlags().String("lark-bot-token", "", "lark bot token") + RootCmd.PersistentFlags().String("binance-api-key", "", "binance api key") RootCmd.PersistentFlags().String("binance-api-secret", "", "binance api secret") diff --git a/pkg/notifier/larknotifier/lark.go b/pkg/notifier/larknotifier/lark.go new file mode 100644 index 0000000..7fc2d25 --- /dev/null +++ b/pkg/notifier/larknotifier/lark.go @@ -0,0 +1,196 @@ +package larknotifier + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "git.qtrade.icu/lychiyu/bbgo/pkg/types" + "github.com/sirupsen/logrus" + "golang.org/x/time/rate" + "net/http" + "reflect" + "time" +) + +var apiLimiter = rate.NewLimiter(rate.Every(time.Second), 5) +var log = logrus.WithField("service", "lark") + +type notifyTask struct { + message string +} + +type Notifier struct { + token string + taskC chan notifyTask +} + +type Option func(notifier *Notifier) + +type TextMessage struct { + MsgType string `json:"msg_type"` + Content struct { + Text string `json:"text"` + } `json:"content"` +} + +func getLocDateTimeString() string { + loc, err := time.LoadLocation("Asia/Shanghai") + if err != nil { + fmt.Println("Error loading location:", err) + return "" + } + + // 获取当前东八区时间 + currentTime := time.Now().In(loc) + + // 格式化时间 + formattedTime := currentTime.Format("2006-01-02 15:04:05") + return formattedTime +} + +func sendLarkMessage(token, text string) error { + // 创建消息结构体 + message := TextMessage{ + MsgType: "text", + Content: struct { + Text string `json:"text"` + }{ + Text: fmt.Sprintf("通知时间：%s\n内容：\n%s", getLocDateTimeString(), text), + }, + } + + // 将消息结构体序列化为 JSON + jsonData, err := json.Marshal(message) + if err != nil { + return fmt.Errorf("failed to marshal JSON: %v", err) + } + webhookURL := fmt.Sprintf("https://open.larksuite.com/open-apis/bot/v2/hook/%s", token) + // 创建 HTTP POST 请求 + req, err := http.NewRequest("POST", webhookURL, bytes.NewBuffer(jsonData)) + if err != nil { + return fmt.Errorf("failed to create request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + + // 发送 HTTP 请求 + client := &http.Client{} + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("failed to send request: %v", err) + } + defer resp.Body.Close() + + // 检查响应状态码 + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("request failed with status code: %d", resp.StatusCode) + } + + return nil +} + +func New(token string, options ...Option) *Notifier { + notifier := &Notifier{ + token: token, + taskC: make(chan notifyTask, 100), + } + + for _, o := range options { + o(notifier) + } + + go notifier.worker() + + return notifier +} + +func (n *Notifier) worker() { + ctx := context.Background() + for { + select { + case <-ctx.Done(): + return + case task := <-n.taskC: + apiLimiter.Wait(ctx) + n.consume(task) + } + } +} + +func (n *Notifier) consume(task notifyTask) { + if task.message != "" { + if err := sendLarkMessage(n.token, task.message); err != nil { + log.WithError(err).Error("lark send error") + } + } +} + +func (n *Notifier) Notify(obj interface{}, args ...interface{}) { + n.NotifyTo("", obj, args...) +} + +func filterPlaintextMessages(args []interface{}) (texts []string, pureArgs []interface{}) { + var firstObjectOffset = -1 + for idx, arg := range args { + rt := reflect.TypeOf(arg) + if rt.Kind() == reflect.Ptr { + switch a := arg.(type) { + + case nil: + texts = append(texts, "nil") + if firstObjectOffset == -1 { + firstObjectOffset = idx + } + + case types.PlainText: + texts = append(texts, a.PlainText()) + if firstObjectOffset == -1 { + firstObjectOffset = idx + } + + case types.Stringer: + texts = append(texts, a.String()) + if firstObjectOffset == -1 { + firstObjectOffset = idx + } + } + } + } + + pureArgs = args + if firstObjectOffset > -1 { + pureArgs = args[:firstObjectOffset] + } + + return texts, pureArgs +} + +func (n *Notifier) NotifyTo(channel string, obj interface{}, args ...interface{}) { + var message string + + switch a := obj.(type) { + + case string: + message = fmt.Sprintf(a, args...) + + default: + //log.Errorf("unsupported notification format: %T %+v", a, a) + return + } + + select { + case n.taskC <- notifyTask{ + message: message, + }: + default: + log.Error("[lark] cannot send task to notify") + } +} + +func (n *Notifier) SendPhoto(buffer *bytes.Buffer) { + n.SendPhotoTo("", buffer) +} + +func (n *Notifier) SendPhotoTo(channel string, buffer *bytes.Buffer) { + // TODO +} diff --git a/pkg/notifier/larknotifier/logrus_look.go b/pkg/notifier/larknotifier/logrus_look.go new file mode 100644 index 0000000..ef73b2d --- /dev/null +++ b/pkg/notifier/larknotifier/logrus_look.go @@ -0,0 +1,45 @@ +package larknotifier + +import ( + "fmt" + "time" + + "github.com/sirupsen/logrus" + "golang.org/x/time/rate" +) + +var limiter = rate.NewLimiter(rate.Every(time.Minute), 3) + +type LogHook struct { + notifier *Notifier +} + +func LarkNewLogHook(notifier *Notifier) *LogHook { + return &LogHook{ + notifier: notifier, + } +} + +func (t *LogHook) Levels() []logrus.Level { + return []logrus.Level{ + logrus.ErrorLevel, + logrus.FatalLevel, + logrus.PanicLevel, + } +} + +func (t *LogHook) Fire(e *logrus.Entry) error { + if !limiter.Allow() { + return nil + } + + var message = fmt.Sprintf("[%s] %s", e.Level.String(), e.Message) + if errData, ok := e.Data[logrus.ErrorKey]; ok && errData != nil { + if err, isErr := errData.(error); isErr { + message += " Error: " + err.Error() + } + } + + t.notifier.Notify(message) + return nil +}