[add] 添加lark通知

This commit is contained in:
lychiyu 2024-08-30 23:03:08 +08:00
parent 89562094b7
commit 8f98cf650e
4 changed files with 259 additions and 0 deletions

View File

@ -4,6 +4,7 @@ import (
"bytes" "bytes"
"context" "context"
"fmt" "fmt"
"git.qtrade.icu/lychiyu/bbgo/pkg/notifier/larknotifier"
"image/png" "image/png"
"io/ioutil" "io/ioutil"
stdlog "log" 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 userConfig.Notifications != nil {
if err := environ.ConfigureNotification(userConfig.Notifications); err != nil { if err := environ.ConfigureNotification(userConfig.Notifications); err != nil {
return err return err
@ -959,6 +966,15 @@ func (environ *Environment) setupTelegram(
return nil 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 { func writeOTPKeyAsQRCodePNG(key *otp.Key, imagePath string) error {
// Convert TOTP key into a PNG // Convert TOTP key into a PNG
var buf bytes.Buffer var buf bytes.Buffer

View File

@ -190,6 +190,8 @@ func init() {
RootCmd.PersistentFlags().String("telegram-bot-token", "", "telegram bot token from bot father") 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("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-key", "", "binance api key")
RootCmd.PersistentFlags().String("binance-api-secret", "", "binance api secret") RootCmd.PersistentFlags().String("binance-api-secret", "", "binance api secret")

View File

@ -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
}

View File

@ -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
}