telemt-bot/main.go

package main

import (
	"bytes"
	"context"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"regexp"
	"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

	TelemtLinkHost      string
	TelemtLinkPort      int
	TelemtLinkTLSDomain string
}

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"`
	Links              userLinks `json:"links"`
}

type userLinks struct {
	Classic []string `json:"classic"`
	Secure  []string `json:"secure"`
	TLS     []string `json:"tls"`
}

type createUserRequest struct {
	Username string `json:"username"`
}

type createUserResponse struct {
	User   userInfo `json:"user"`
	Secret string   `json:"secret"`
}

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)
	cfg.TelemtLinkHost = strings.TrimSpace(os.Getenv("TELEMT_LINK_HOST"))
	cfg.TelemtLinkPort = parseIntEnv("TELEMT_LINK_PORT", 9443)
	cfg.TelemtLinkTLSDomain = strings.TrimSpace(os.Getenv("TELEMT_TLS_DOMAIN"))

	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 parseIntEnv(key string, fallback int) int {
	raw := strings.TrimSpace(os.Getenv(key))
	if raw == "" {
		return fallback
	}
	v, err := strconv.Atoi(raw)
	if err != nil || v <= 0 {
		return fallback
	}
	return v
}

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\n/create_user <username>")
	case "/health":
		return b.handleHealth(ctx, chatID)
	case "/summary":
		return b.handleSummary(ctx, chatID)
	case "/users":
		return b.handleUsers(ctx, chatID)
	case "/create_user", "/createuser":
		return b.handleCreateUser(ctx, chatID, text)
	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) handleCreateUser(ctx context.Context, chatID int64, text string) error {
	parts := strings.Fields(text)
	if len(parts) < 2 {
		return b.sendMessage(ctx, chatID, "Использование: /create_user <username>")
	}

	username := strings.TrimSpace(parts[1])
	if !isValidUsername(username) {
		return b.sendMessage(ctx, chatID, "Некорректный username. Разрешены [A-Za-z0-9_.-], длина 1..64.")
	}

	req := createUserRequest{Username: username}
	var created createUserResponse
	if err := b.callTelemtJSON(ctx, http.MethodPost, "/v1/users", req, &created, http.StatusCreated); err != nil {
		return b.sendMessage(ctx, chatID, "Ошибка создания пользователя: "+err.Error())
	}

	links := b.resolveUserLinks(ctx, created.User.Username, created.Secret, created.User.Links)

	msg := fmt.Sprintf("Пользователь создан: %s\nsecret: %s", created.User.Username, created.Secret)
	if linksText := formatUserLinks(links); linksText != "" {
		msg = msg + "\n\nСсылки:\n" + linksText
	} else {
		msg = msg + "\n\nСсылки не вернулись из API."
	}
	return b.sendMessage(ctx, chatID, msg)
}

func hasAnyLinks(links userLinks) bool {
	return len(links.Classic) > 0 || len(links.Secure) > 0 || len(links.TLS) > 0
}

// serverParamRe matches server=... in tg://proxy URLs
var serverParamRe = regexp.MustCompile(`server=[^&]+`)

func deduplicateStrings(items []string) []string {
	if len(items) <= 1 {
		return items
	}
	seen := make(map[string]struct{}, len(items))
	out := make([]string, 0, len(items))
	for _, s := range items {
		if _, ok := seen[s]; ok {
			continue
		}
		seen[s] = struct{}{}
		out = append(out, s)
	}
	return out
}

func (b *bot) rewriteLinksWithHost(links userLinks) userLinks {
	host := strings.TrimSpace(b.cfg.TelemtLinkHost)
	if host == "" {
		return links
	}
	replacement := "server=" + host
	rewrite := func(items []string) []string {
		if len(items) == 0 {
			return items
		}
		out := make([]string, len(items))
		for i, s := range items {
			out[i] = serverParamRe.ReplaceAllString(s, replacement)
		}
		return deduplicateStrings(out)
	}
	return userLinks{
		Classic: rewrite(links.Classic),
		Secure:  rewrite(links.Secure),
		TLS:     rewrite(links.TLS),
	}
}

func (b *bot) resolveUserLinks(ctx context.Context, username, secret string, initial userLinks) userLinks {
	var result userLinks
	if hasAnyLinks(initial) {
		result = initial
	} else {
		var fetched userInfo
		path := "/v1/users/" + url.PathEscape(username)
		if err := b.callTelemt(ctx, path, &fetched); err == nil && hasAnyLinks(fetched.Links) {
			result = fetched.Links
		} else {
			var users []userInfo
			if err := b.callTelemt(ctx, "/v1/users", &users); err == nil {
				for i := range users {
					if users[i].Username == username && hasAnyLinks(users[i].Links) {
						result = users[i].Links
						break
					}
				}
			}
			if !hasAnyLinks(result) {
				if err := b.callTelemt(ctx, "/v1/stats/users", &users); err == nil {
					for i := range users {
						if users[i].Username == username && hasAnyLinks(users[i].Links) {
							result = users[i].Links
							break
						}
					}
				}
			}
			if !hasAnyLinks(result) {
				result = b.generateLinksFromSecret(secret)
			}
		}
	}
	return b.rewriteLinksWithHost(result)
}

func (b *bot) generateLinksFromSecret(secret string) userLinks {
	host := strings.TrimSpace(b.cfg.TelemtLinkHost)
	if host == "" || secret == "" {
		return userLinks{}
	}
	port := b.cfg.TelemtLinkPort
	classic := fmt.Sprintf("tg://proxy?server=%s&port=%d&secret=%s", host, port, secret)
	dd := fmt.Sprintf("tg://proxy?server=%s&port=%d&secret=dd%s", host, port, secret)

	out := userLinks{
		Classic: []string{classic},
		Secure:  []string{dd},
	}
	if b.cfg.TelemtLinkTLSDomain != "" {
		hexDomain := hex.EncodeToString([]byte(b.cfg.TelemtLinkTLSDomain))
		tls := fmt.Sprintf("tg://proxy?server=%s&port=%d&secret=ee%s%s", host, port, secret, hexDomain)
		out.TLS = []string{tls}
	}
	return out
}

func formatUserLinks(links userLinks) string {
	var sb strings.Builder
	appendLinks := func(title string, items []string) {
		if len(items) == 0 {
			return
		}
		for i := range items {
			if len(items) == 1 {
				sb.WriteString(title)
				sb.WriteString(": ")
			} else {
				sb.WriteString(title)
				sb.WriteString(" [")
				sb.WriteString(strconv.Itoa(i + 1))
				sb.WriteString("]: ")
			}
			sb.WriteString(items[i])
			sb.WriteByte('\n')
		}
	}

	appendLinks("Classic", links.Classic)
	appendLinks("DD", links.Secure)
	appendLinks("EE-TLS", links.TLS)
	return strings.TrimRight(sb.String(), "\n")
}

func isValidUsername(v string) bool {
	if len(v) < 1 || len(v) > 64 {
		return false
	}
	for i := 0; i < len(v); i++ {
		ch := v[i]
		if (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || (ch >= '0' && ch <= '9') || ch == '_' || ch == '.' || ch == '-' {
			continue
		}
		return false
	}
	return true
}

func (b *bot) callTelemt(ctx context.Context, path string, out any) error {
	return b.callTelemtJSON(ctx, http.MethodGet, path, nil, out, http.StatusOK)
}

func (b *bot) callTelemtJSON(ctx context.Context, method, path string, payload any, out any, successStatus int) error {
	cctx, cancel := context.WithTimeout(ctx, b.cfg.TelemtTimeout)
	defer cancel()

	var bodyReader *bytes.Reader
	if payload != nil {
		body, err := json.Marshal(payload)
		if err != nil {
			return err
		}
		bodyReader = bytes.NewReader(body)
	} else {
		bodyReader = bytes.NewReader(nil)
	}

	req, err := http.NewRequestWithContext(cctx, method, b.cfg.TelemtAPIBaseURL+path, bodyReader)
	if err != nil {
		return err
	}
	if payload != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	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()

	var env telemtEnvelope
	if err := json.NewDecoder(resp.Body).Decode(&env); err != nil {
		if resp.StatusCode != successStatus {
			return fmt.Errorf("telemt status %s", resp.Status)
		}
		return fmt.Errorf("telemt decode response: %w", err)
	}

	if resp.StatusCode != successStatus {
		if env.Error != nil {
			return fmt.Errorf("%s: %s", env.Error.Code, env.Error.Message)
		}
		return fmt.Errorf("telemt status %s", resp.Status)
	}
	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 out == nil {
		return nil
	}
	if len(env.Data) == 0 || string(env.Data) == "null" {
		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
}