tgbot_misaka_5882f7/bot.go

package main

import (
	"fmt"
	"html"
	"math"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/go-errors/errors"
	tele "gopkg.in/telebot.v3"

	"git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/cmds"
)

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("/xr", handleExchangeRateCmd)

	b.Handle(tele.OnText, handleGeneralMessage)
	b.Handle(tele.OnSticker, handleGeneralMessage)

	// admin required
	adminGrp := b.Group()
	adminGrp.Use(adminMiddleware)
	adminGrp.Handle("/traffic", handleTrafficCmd)
	adminGrp.Handle(&trafficBtnDays, handleTrafficBtnDays)
	adminGrp.Handle(&trafficBtnMonths, handleTrafficBtnMonths)

	adminGrp.Handle("/dig", handleDigCmd)

	// adminGrp.Handle("/test", handleTestCmd)

	return b, nil
}

func adminMiddleware(next tele.HandlerFunc) tele.HandlerFunc {
	return func(c tele.Context) error {
		u := c.Sender()
		if u == nil {
			if cb := c.Callback(); cb != nil {
				u = cb.Sender
			}
		}
		if !isFromAdmin(u) {
			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 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 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 <b>%2.0f%%</b> complete.\n<pre>%s</pre>",
		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 <name> [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("<i>Status: <b>%s</b></i>\n", resp.Status),
		fmt.Sprintf("<i>Query Time: <b>%s</b></i>\n\n", resp.QueryTime),
		"<pre>",
		html.EscapeString(replyBuf.String()),
		"</pre>",
	}
	return c.Reply(strings.Join(replyText, ""), &tele.SendOptions{ParseMode: tele.ModeHTML}, tele.Silent)
}