feat: add bar for traffic usage

This commit is contained in:
Yiyang Kang 2022-11-21 15:17:02 +08:00
parent b429137d20
commit 3fc427fae4
2 changed files with 46 additions and 14 deletions

37
bot.go
View File

@ -2,10 +2,12 @@ package main
import ( import (
"fmt" "fmt"
"math"
"strings" "strings"
"time" "time"
"github.com/go-errors/errors" "github.com/go-errors/errors"
"github.com/samber/lo"
tele "gopkg.in/telebot.v3" tele "gopkg.in/telebot.v3"
"git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/stats" "git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/stats"
@ -96,11 +98,12 @@ func handleTrafficCmd(c tele.Context) error {
// This month's traffic, if present // This month's traffic, if present
if len(monthlyTraffic) > 0 { if len(monthlyTraffic) > 0 {
row := monthlyTraffic[len(monthlyTraffic)-1] row := monthlyTraffic[len(monthlyTraffic)-1]
responseParts = append( responseParts = append(
responseParts, responseParts,
fmt.Sprintf("This month so far:` %s`", fmtTraffic(row)), 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 var respText string
@ -114,12 +117,8 @@ func handleTrafficCmd(c tele.Context) error {
} }
func fmtTraffic(r stats.VnstatTrafficRecord) string { func fmtTraffic(r stats.VnstatTrafficRecord) string {
biggest := r.Rx effective := lo.Max([]uint64{r.Rx, r.Tx})
if r.Tx > biggest { return fmt.Sprintf("%.2f GiB", float64(effective)/1024/1024/1024)
biggest = r.Tx
}
return fmt.Sprintf("%.2f GiB", float64(biggest)/1024/1024/1024)
} }
func handleUserInfoCmd(c tele.Context) error { 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}) 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)
}

23
cfg.go
View File

@ -13,7 +13,8 @@ type Config struct {
AdminUIDs map[int64]struct{} AdminUIDs map[int64]struct{}
TGBotToken string TGBotToken string
WatchedInterface string WatchedInterface string
MonthlyTrafficLimitGiB int
} }
var config *Config var config *Config
@ -28,8 +29,9 @@ func LoadCfg() error {
cfg.TGBotToken = token cfg.TGBotToken = token
adminUIDsEnv := os.Getenv("TG_ADMIN_UIDS") adminUIDsEnv := os.Getenv("TG_ADMIN_UIDS")
adminUIDs := lo.Filter(strings.Split(adminUIDsEnv, ","), func(s string, _ int) bool { adminUIDs := lo.FilterMap(strings.Split(adminUIDsEnv, ","), func(s string, _ int) (string, bool) {
return strings.Trim(s, "\t ") != "" trimmed := strings.TrimSpace(s)
return trimmed, trimmed != ""
}) })
cfg.AdminUIDs = make(map[int64]struct{}, len(adminUIDs)) cfg.AdminUIDs = make(map[int64]struct{}, len(adminUIDs))
for _, uidStr := range adminUIDs { for _, uidStr := range adminUIDs {
@ -40,11 +42,18 @@ func LoadCfg() error {
cfg.AdminUIDs[uid] = struct{}{} cfg.AdminUIDs[uid] = struct{}{}
} }
iface := os.Getenv("TG_WATCHED_INTERFACE") cfg.WatchedInterface = "eth0"
if iface == "" { if iface := os.Getenv("TG_WATCHED_INTERFACE"); iface != "" {
iface = "eth0" 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 config = &cfg
return nil return nil