package main import ( "fmt" "html" "math" "strings" "text/tabwriter" "time" "github.com/go-errors/errors" "github.com/samber/lo" tele "gopkg.in/telebot.v3" "git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/cmds" "git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/stats" ) const ( stickerPanic = "CAACAgUAAxkBAAMjY3zoraxZGB8Xejyw86bHLSWLjVcAArMIAAL7-nhXNK7dStmRUGsrBA" stickerLoading = "CAACAgUAAxkBAAMmY3zp5UCMVRvy1isFCPHrx-UBWX8AApYHAALP8GhXEm9ZIBjn1v8rBA" ) func isFromAdmin(sender *tele.User) bool { if sender == nil { return false } _, ok := config.AdminUIDs[sender.ID] return ok } func initBot() (*tele.Bot, error) { pref := tele.Settings{ Token: config.TGBotToken, Poller: &tele.LongPoller{Timeout: 15 * time.Second}, } b, err := tele.NewBot(pref) if err != nil { return nil, errors.Wrap(err, 0) } b.Use(logMiddleware) // command routing b.Handle("/start", handleStartCmd) b.Handle("/me", handleUserInfoCmd) b.Handle("/chat", handleChatInfoCmd) b.Handle("/year_progress", handleYearProgressCmd) b.Handle(tele.OnText, handleGeneralMessage) b.Handle(tele.OnSticker, handleGeneralMessage) // admin required adminGrp := b.Group() adminGrp.Use(adminMiddleware) adminGrp.Handle("/traffic", handleTrafficCmd) adminGrp.Handle("/dig", handleDigCmd) // adminGrp.Handle("/test", handleTestCmd) return b, nil } func adminMiddleware(next tele.HandlerFunc) tele.HandlerFunc { return func(c tele.Context) error { if !isFromAdmin(c.Sender()) { return nil } return next(c) } } func logMiddleware(next tele.HandlerFunc) tele.HandlerFunc { return func(c tele.Context) error { upd := c.Update() defer func() { logger.Infow("Log middleware", "update", upd) }() return next(c) } } func handleStartCmd(c tele.Context) error { if !isFromAdmin(c.Sender()) { return c.Send("Hello, stranger :)") } return c.Send("Hi :)") } func handleTrafficCmd(c tele.Context) error { dailyTraffic, err := stats.VnstatDailyTraffic(config.WatchedInterface) if err != nil { _ = c.Reply(stickerFromID(stickerPanic), tele.Silent) return err } monthlyTraffic, err := stats.VnstatMonthlyTraffic(config.WatchedInterface) if err != nil { _ = c.Reply(stickerFromID(stickerPanic), tele.Silent) return err } var responseParts = []string{"*Traffic Usage Summaries*\n"} // Yesterday's traffic if present if len(dailyTraffic) > 1 { row := dailyTraffic[len(dailyTraffic)-2] day := fmt.Sprintf("%d-%02d-%02d", row.Date.Year, row.Date.Month, row.Date.Day) responseParts = append(responseParts, fmt.Sprintf("Yesterday (%s):` %s`", day, fmtTraffic(row))) } // Today's traffic if present if len(dailyTraffic) > 0 { row := dailyTraffic[len(dailyTraffic)-1] responseParts = append(responseParts, fmt.Sprintf("Today so far:` %s`\n", fmtTraffic(row))) } // Last month's traffic, if present if len(monthlyTraffic) > 1 { row := monthlyTraffic[len(monthlyTraffic)-2] month := fmt.Sprintf("%d-%02d", row.Date.Year, row.Date.Month) responseParts = append(responseParts, fmt.Sprintf("Last month (%s):` %s`", month, fmtTraffic(row))) } // 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))) responseParts = append(responseParts, drawBarForTrafficRecord(row)) } var respText string if len(responseParts) == 1 { respText = "No traffic data available." } else { respText = strings.Join(responseParts, "\n") } return c.Reply(respText, &tele.SendOptions{ParseMode: tele.ModeMarkdown}, tele.Silent) } func fmtTraffic(r stats.VnstatTrafficRecord) string { effective := lo.Max([]uint64{r.Rx, r.Tx}) return fmt.Sprintf("%.2f GiB", float64(effective)/1024/1024/1024) } func handleUserInfoCmd(c tele.Context) error { u := c.Sender() if u == nil { return c.Reply("Unknown.", tele.Silent) } replyText := []string{ `*User Info*`, "```", fmt.Sprintf(`ID: %d`, u.ID), fmt.Sprintf(`Username: %s`, u.Username), fmt.Sprintf(`FirstName: %s`, u.FirstName), fmt.Sprintf(`LastName: %s`, u.LastName), fmt.Sprintf(`LanguageCode: %s`, u.LanguageCode), fmt.Sprintf(`IsBot: %t`, u.IsBot), fmt.Sprintf(`IsPremium: %t`, u.IsPremium), "```", } return c.Reply(strings.Join(replyText, "\n"), &tele.SendOptions{ParseMode: tele.ModeMarkdown}, tele.Silent) } func handleChatInfoCmd(c tele.Context) error { chat := c.Chat() if chat == nil { return c.Reply("Unknown.", tele.Silent) } loc := "" if chat.ChatLocation != nil { loc = chat.ChatLocation.Address } replyText := []string{ `*Chat Info*`, "```", fmt.Sprintf(`ID: %d`, chat.ID), fmt.Sprintf(`Type: %s`, chat.Type), fmt.Sprintf(`Title: %s`, chat.Title), fmt.Sprintf(`FirstName: %s`, chat.FirstName), fmt.Sprintf(`LastName: %s`, chat.LastName), fmt.Sprintf(`Username: %s`, chat.Username), fmt.Sprintf(`SlowMode: %d`, chat.SlowMode), fmt.Sprintf(`StickerSet: %s`, chat.StickerSet), fmt.Sprintf(`CanSetStickerSet: %t`, chat.CanSetStickerSet), fmt.Sprintf(`LinkedChatID: %d`, chat.LinkedChatID), fmt.Sprintf(`ChatLocation: %s`, loc), fmt.Sprintf(`Private: %t`, chat.Private), fmt.Sprintf(`Protected: %t`, chat.Protected), fmt.Sprintf(`NoVoiceAndVideo: %t`, chat.NoVoiceAndVideo), "```", } return c.Reply(strings.Join(replyText, "\n"), &tele.SendOptions{ParseMode: tele.ModeMarkdown}, tele.Silent) } func drawBar(progress float64, length int) string { barChars := []rune("ยท-=") if length <= 0 { return "" } step := 1 / float64(length) buf := make([]rune, length+2) buf[0], buf[length+1] = '[', ']' 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+1] = 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) } func handleYearProgressCmd(c tele.Context) error { yearStart := time.Date(time.Now().Year(), 1, 1, 0, 0, 0, 0, time.Local) yearEnd := yearStart.AddDate(1, 0, 0) yearDur := yearEnd.Sub(yearStart) elapsed := time.Since(yearStart) ratio := float64(elapsed) / float64(yearDur) replyText := fmt.Sprintf( "\n%d is %2.0f%% complete.\n
%s
", time.Now().Year(), ratio*100, drawBar(ratio, 20), ) return c.Reply(replyText, &tele.SendOptions{ParseMode: tele.ModeHTML}, tele.Silent) } func handleGeneralMessage(_ tele.Context) error { // Do nothing for now return nil } func stickerFromID(id string) *tele.Sticker { return &tele.Sticker{ File: tele.File{ FileID: id, }, } } func handleDigCmd(c tele.Context) error { msg := c.Message() if msg == nil { return nil } req, err := cmds.NewDigRequest(msg.Payload) if err != nil { return c.Reply( "Invalid arguments.\nUsage: `/dig [type]`", &tele.SendOptions{ParseMode: tele.ModeMarkdown}, tele.Silent, ) } resp, err := cmds.Dig(req) if err != nil { _ = c.Reply(stickerFromID(stickerPanic), tele.Silent) return err } replyBuf := &strings.Builder{} tw := tabwriter.NewWriter(replyBuf, 0, 0, 2, ' ', 0) // Write header if len(resp.Records) > 0 { _, _ = tw.Write([]byte("Name\tTTL\tType\tData\n")) } // Write data for _, r := range resp.Records { _, _ = fmt.Fprintf(tw, "%s\t%d\t%s\t%s\n", r.Name, r.TTL, r.Type, r.Data) } _ = tw.Flush() replyText := []string{ fmt.Sprintf("Status: %s\n", resp.Status), fmt.Sprintf("Query Time: %s\n\n", resp.QueryTime), "
",
		html.EscapeString(replyBuf.String()),
		"
", } return c.Reply(strings.Join(replyText, ""), &tele.SendOptions{ParseMode: tele.ModeHTML}, tele.Silent) }