From 3fc427fae4fe798a52d2b2d849e95a41669cb17e Mon Sep 17 00:00:00 2001 From: Yiyang Kang Date: Mon, 21 Nov 2022 15:17:02 +0800 Subject: [PATCH] feat: add bar for traffic usage --- bot.go | 37 ++++++++++++++++++++++++++++++------- cfg.go | 23 ++++++++++++++++------- 2 files changed, 46 insertions(+), 14 deletions(-) diff --git a/bot.go b/bot.go index c5015c9..1ca3862 100644 --- a/bot.go +++ b/bot.go @@ -2,10 +2,12 @@ package main import ( "fmt" + "math" "strings" "time" "github.com/go-errors/errors" + "github.com/samber/lo" tele "gopkg.in/telebot.v3" "git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/stats" @@ -96,11 +98,12 @@ func handleTrafficCmd(c tele.Context) error { // This month's traffic, if present if len(monthlyTraffic) > 0 { row := monthlyTraffic[len(monthlyTraffic)-1] - responseParts = append( responseParts, fmt.Sprintf("This month so far:` %s`", fmtTraffic(row)), ) + bar := fmt.Sprintf("\n```\n%s\n```", drawBarForTrafficRecord(row)) + responseParts = append(responseParts, bar) } var respText string @@ -114,12 +117,8 @@ func handleTrafficCmd(c tele.Context) error { } func fmtTraffic(r stats.VnstatTrafficRecord) string { - biggest := r.Rx - if r.Tx > biggest { - biggest = r.Tx - } - - return fmt.Sprintf("%.2f GiB", float64(biggest)/1024/1024/1024) + effective := lo.Max([]uint64{r.Rx, r.Tx}) + return fmt.Sprintf("%.2f GiB", float64(effective)/1024/1024/1024) } func handleUserInfoCmd(c tele.Context) error { @@ -174,3 +173,27 @@ func handleChatInfoCmd(c tele.Context) error { } return c.Reply(strings.Join(replyText, "\n"), &tele.SendOptions{ParseMode: tele.ModeMarkdown}) } + +func drawBar(progress float64, length int) string { + barChars := []rune("░▒▓█") + + if length <= 0 { + return "" + } + step := 1 / float64(length) + buf := make([]rune, length) + for i := 0; i < length; i++ { + fill := (progress - float64(i)*step) / step + fill = math.Min(math.Max(fill, 0), 1) + idx := int(math.Round(fill * float64(len(barChars)-1))) + buf[i] = barChars[idx] + } + return string(buf) +} + +func drawBarForTrafficRecord(r stats.VnstatTrafficRecord) string { + effective := lo.Max([]uint64{r.Rx, r.Tx}) + max := config.MonthlyTrafficLimitGiB + ratio := float64(effective) / 1024 / 1024 / 1024 / float64(max) + return fmt.Sprintf("%s %2.0f%%", drawBar(ratio, 16), ratio*100) +} diff --git a/cfg.go b/cfg.go index b2f65d8..86cfc33 100644 --- a/cfg.go +++ b/cfg.go @@ -13,7 +13,8 @@ type Config struct { AdminUIDs map[int64]struct{} TGBotToken string - WatchedInterface string + WatchedInterface string + MonthlyTrafficLimitGiB int } var config *Config @@ -28,8 +29,9 @@ func LoadCfg() error { cfg.TGBotToken = token adminUIDsEnv := os.Getenv("TG_ADMIN_UIDS") - adminUIDs := lo.Filter(strings.Split(adminUIDsEnv, ","), func(s string, _ int) bool { - return strings.Trim(s, "\t ") != "" + adminUIDs := lo.FilterMap(strings.Split(adminUIDsEnv, ","), func(s string, _ int) (string, bool) { + trimmed := strings.TrimSpace(s) + return trimmed, trimmed != "" }) cfg.AdminUIDs = make(map[int64]struct{}, len(adminUIDs)) for _, uidStr := range adminUIDs { @@ -40,11 +42,18 @@ func LoadCfg() error { cfg.AdminUIDs[uid] = struct{}{} } - iface := os.Getenv("TG_WATCHED_INTERFACE") - if iface == "" { - iface = "eth0" + cfg.WatchedInterface = "eth0" + if iface := os.Getenv("TG_WATCHED_INTERFACE"); iface != "" { + cfg.WatchedInterface = iface + } + + cfg.MonthlyTrafficLimitGiB = 1000 + if trafficLimitStr := os.Getenv("TG_MONTHLY_TRAFFIC_LIMIT_GIB"); trafficLimitStr != "" { + var err error + if cfg.MonthlyTrafficLimitGiB, err = strconv.Atoi(trafficLimitStr); err != nil { + return errors.New("invalid traffic limit: " + trafficLimitStr) + } } - cfg.WatchedInterface = iface config = &cfg return nil