Files
telemt-bot/main.go
T
2026-03-09 16:05:44 +07:00

383 lines
9.5 KiB
Go

package main
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"log"
"net/http"
"net/url"
"os"
"os/signal"
"strconv"
"strings"
"syscall"
"time"
)
const (
defaultTelemtTimeoutSeconds = 5
defaultTelegramTimeoutSeconds = 60
defaultListenPollTimeout = 50
)
type config struct {
TelegramBotToken string
AdminChatID int64
TelemtAPIBaseURL string
TelemtAPIAuth string
TelemtTimeout time.Duration
TelegramTimeout time.Duration
}
type bot struct {
cfg config
httpClient *http.Client
telegramBase string
offset int64
}
type tgGetUpdatesResponse struct {
OK bool `json:"ok"`
Result []tgUpdate `json:"result"`
}
type tgUpdate struct {
UpdateID int64 `json:"update_id"`
Message *tgMessage `json:"message"`
}
type tgMessage struct {
Chat tgChat `json:"chat"`
Text string `json:"text"`
}
type tgChat struct {
ID int64 `json:"id"`
}
type telemtEnvelope struct {
OK bool `json:"ok"`
Data json.RawMessage `json:"data"`
Error *struct {
Code string `json:"code"`
Message string `json:"message"`
} `json:"error"`
}
type healthData struct {
Status string `json:"status"`
ReadOnly bool `json:"read_only"`
}
type summaryData struct {
UptimeSeconds float64 `json:"uptime_seconds"`
ConnectionsTotal uint64 `json:"connections_total"`
ConnectionsBadTotal uint64 `json:"connections_bad_total"`
HandshakeTimeoutsTotal uint64 `json:"handshake_timeouts_total"`
ConfiguredUsers uint64 `json:"configured_users"`
}
type userInfo struct {
Username string `json:"username"`
CurrentConnections uint64 `json:"current_connections"`
ActiveUniqueIPs uint64 `json:"active_unique_ips"`
TotalOctets uint64 `json:"total_octets"`
}
func main() {
cfg, err := loadConfig()
if err != nil {
log.Fatalf("config error: %v", err)
}
httpClient := &http.Client{
Transport: &http.Transport{
MaxIdleConns: 64,
MaxIdleConnsPerHost: 32,
IdleConnTimeout: 90 * time.Second,
TLSHandshakeTimeout: 5 * time.Second,
ForceAttemptHTTP2: true,
DisableCompression: false,
},
Timeout: cfg.TelegramTimeout,
}
b := &bot{
cfg: cfg,
httpClient: httpClient,
telegramBase: "https://api.telegram.org/bot" + cfg.TelegramBotToken,
}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()
log.Printf("telemt-bot started for admin chat %d", cfg.AdminChatID)
if err := b.run(ctx); err != nil && !errors.Is(err, context.Canceled) {
log.Fatalf("bot stopped with error: %v", err)
}
log.Println("telemt-bot stopped")
}
func loadConfig() (config, error) {
var cfg config
cfg.TelegramBotToken = strings.TrimSpace(os.Getenv("TELEGRAM_BOT_TOKEN"))
if cfg.TelegramBotToken == "" {
return cfg, errors.New("TELEGRAM_BOT_TOKEN is required")
}
adminChat := strings.TrimSpace(os.Getenv("TELEGRAM_ADMIN_CHAT_ID"))
if adminChat == "" {
return cfg, errors.New("TELEGRAM_ADMIN_CHAT_ID is required")
}
id, err := strconv.ParseInt(adminChat, 10, 64)
if err != nil {
return cfg, fmt.Errorf("invalid TELEGRAM_ADMIN_CHAT_ID: %w", err)
}
cfg.AdminChatID = id
cfg.TelemtAPIBaseURL = strings.TrimRight(strings.TrimSpace(os.Getenv("TELEMT_API_BASE_URL")), "/")
if cfg.TelemtAPIBaseURL == "" {
return cfg, errors.New("TELEMT_API_BASE_URL is required")
}
if _, err := url.ParseRequestURI(cfg.TelemtAPIBaseURL); err != nil {
return cfg, fmt.Errorf("invalid TELEMT_API_BASE_URL: %w", err)
}
cfg.TelemtAPIAuth = strings.TrimSpace(os.Getenv("TELEMT_API_AUTH_HEADER"))
cfg.TelemtTimeout = parseSecondsEnv("TELEMT_HTTP_TIMEOUT_SECONDS", defaultTelemtTimeoutSeconds)
cfg.TelegramTimeout = parseSecondsEnv("TELEGRAM_HTTP_TIMEOUT_SECONDS", defaultTelegramTimeoutSeconds)
return cfg, nil
}
func parseSecondsEnv(key string, fallback int) time.Duration {
raw := strings.TrimSpace(os.Getenv(key))
if raw == "" {
return time.Duration(fallback) * time.Second
}
v, err := strconv.Atoi(raw)
if err != nil || v <= 0 {
return time.Duration(fallback) * time.Second
}
return time.Duration(v) * time.Second
}
func (b *bot) run(ctx context.Context) error {
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
updates, err := b.getUpdates(ctx)
if err != nil {
log.Printf("getUpdates error: %v", err)
time.Sleep(2 * time.Second)
continue
}
for i := range updates {
u := updates[i]
if u.UpdateID >= b.offset {
b.offset = u.UpdateID + 1
}
if u.Message == nil || strings.TrimSpace(u.Message.Text) == "" {
continue
}
if err := b.handleMessage(ctx, u.Message); err != nil {
log.Printf("handleMessage error: %v", err)
}
}
}
}
func (b *bot) getUpdates(ctx context.Context) ([]tgUpdate, error) {
reqBody := map[string]any{
"timeout": defaultListenPollTimeout,
"offset": b.offset,
"allowed_updates": []string{"message"},
}
var body bytes.Buffer
if err := json.NewEncoder(&body).Encode(reqBody); err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, b.telegramBase+"/getUpdates", &body)
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
resp, err := b.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("telegram getUpdates status: %s", resp.Status)
}
var parsed tgGetUpdatesResponse
if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
return nil, err
}
if !parsed.OK {
return nil, errors.New("telegram getUpdates returned ok=false")
}
return parsed.Result, nil
}
func (b *bot) handleMessage(ctx context.Context, msg *tgMessage) error {
chatID := msg.Chat.ID
text := strings.TrimSpace(msg.Text)
if text == "" {
return nil
}
if chatID != b.cfg.AdminChatID {
return nil
}
command := strings.Fields(text)[0]
if at := strings.IndexByte(command, '@'); at > 0 {
command = command[:at]
}
switch command {
case "/start", "/help":
return b.sendMessage(ctx, chatID, "Команды:\n/health\n/summary\n/users")
case "/health":
return b.handleHealth(ctx, chatID)
case "/summary":
return b.handleSummary(ctx, chatID)
case "/users":
return b.handleUsers(ctx, chatID)
default:
return b.sendMessage(ctx, chatID, "Неизвестная команда. Используй /help")
}
}
func (b *bot) handleHealth(ctx context.Context, chatID int64) error {
var h healthData
if err := b.callTelemt(ctx, "/v1/health", &h); err != nil {
return b.sendMessage(ctx, chatID, "Ошибка API: "+err.Error())
}
msg := fmt.Sprintf("Health: %s\nRead-only: %t", h.Status, h.ReadOnly)
return b.sendMessage(ctx, chatID, msg)
}
func (b *bot) handleSummary(ctx context.Context, chatID int64) error {
var s summaryData
if err := b.callTelemt(ctx, "/v1/stats/summary", &s); err != nil {
return b.sendMessage(ctx, chatID, "Ошибка API: "+err.Error())
}
msg := fmt.Sprintf(
"Uptime: %.0fs\nConn total: %d\nConn bad: %d\nHandshake timeouts: %d\nUsers: %d",
s.UptimeSeconds, s.ConnectionsTotal, s.ConnectionsBadTotal, s.HandshakeTimeoutsTotal, s.ConfiguredUsers,
)
return b.sendMessage(ctx, chatID, msg)
}
func (b *bot) handleUsers(ctx context.Context, chatID int64) error {
var users []userInfo
if err := b.callTelemt(ctx, "/v1/stats/users", &users); err != nil {
return b.sendMessage(ctx, chatID, "Ошибка API: "+err.Error())
}
if len(users) == 0 {
return b.sendMessage(ctx, chatID, "Пользователи не найдены.")
}
const maxRows = 20
var sb strings.Builder
if len(users) > maxRows {
sb.WriteString(fmt.Sprintf("Показаны первые %d из %d\n", maxRows, len(users)))
}
limit := len(users)
if limit > maxRows {
limit = maxRows
}
for i := 0; i < limit; i++ {
u := users[i]
sb.WriteString(fmt.Sprintf(
"%d) %s | conn=%d | ips=%d | octets=%d\n",
i+1, u.Username, u.CurrentConnections, u.ActiveUniqueIPs, u.TotalOctets,
))
}
return b.sendMessage(ctx, chatID, strings.TrimRight(sb.String(), "\n"))
}
func (b *bot) callTelemt(ctx context.Context, path string, out any) error {
cctx, cancel := context.WithTimeout(ctx, b.cfg.TelemtTimeout)
defer cancel()
req, err := http.NewRequestWithContext(cctx, http.MethodGet, b.cfg.TelemtAPIBaseURL+path, nil)
if err != nil {
return err
}
if b.cfg.TelemtAPIAuth != "" {
req.Header.Set("Authorization", b.cfg.TelemtAPIAuth)
}
resp, err := b.httpClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("telemt status %s", resp.Status)
}
var env telemtEnvelope
if err := json.NewDecoder(resp.Body).Decode(&env); err != nil {
return err
}
if !env.OK {
if env.Error != nil {
return fmt.Errorf("%s: %s", env.Error.Code, env.Error.Message)
}
return errors.New("telemt returned ok=false")
}
if len(env.Data) == 0 {
return errors.New("telemt empty data")
}
return json.Unmarshal(env.Data, out)
}
func (b *bot) sendMessage(ctx context.Context, chatID int64, text string) error {
reqBody := map[string]any{
"chat_id": chatID,
"text": text,
}
var body bytes.Buffer
if err := json.NewEncoder(&body).Encode(reqBody); err != nil {
return err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, b.telegramBase+"/sendMessage", &body)
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
resp, err := b.httpClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("telegram sendMessage status: %s", resp.Status)
}
return nil
}