telegram-antispam/bot.go

package main

import (
	"bytes"
	"context"
	"encoding/base64"
	"fmt"
	"log"
	"math"
	"net/http"
	"os"
	"regexp"
	"strings"
	"sync"

	"github.com/gotd/td/session"
	"github.com/gotd/td/telegram"
	"github.com/gotd/td/telegram/auth"
	"github.com/gotd/td/telegram/message"
	"github.com/gotd/td/telegram/message/html"
	"github.com/gotd/td/telegram/updates"
	"github.com/gotd/td/tg"
	"golang.org/x/term"
	"gopkg.in/yaml.v3"
)

type YAMLScamPattern struct {
	Name    string  `yaml:"name"`
	Pattern string  `yaml:"pattern"`
	Weight  float64 `yaml:"weight"`
}

type Config struct {
	APIId         int               `yaml:"api_id"`
	APIHash       string            `yaml:"api_hash"`
	MonitoredChat int64             `yaml:"monitored_chat"`
	AlertChat     int64             `yaml:"alert_chat"`
	NtfyHost      string            `yaml:"ntfy_host"`
	NtfyToken     string            `yaml:"ntfy_token"`
	NtfyTopic     string            `yaml:"ntfy_topic"`
	SessionPath   string            `yaml:"session_path"`
	YAMLPatterns  []YAMLScamPattern `yaml:"scam_patterns"`
	BadWords      []string          `yaml:"bad_words"`
	Patterns      []scamPattern     `yaml:"-"` // Compiled patterns, not from YAML
}

func loadConfig(path string) (*Config, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			// Create default config file
			defaultConfig := `api_id: 0
api_hash: ""
monitored_chat: 0
alert_chat: 0
ntfy_host: "https://ntfy.sh/"
ntfy_token: ""
ntfy_topic: ""
session_path: "antiscam.session"
scam_patterns: []
bad_words: []
`
			if err := os.WriteFile(path, []byte(defaultConfig), 0644); err != nil {
				return nil, fmt.Errorf("creating default config file: %w", err)
			}
			log.Printf("Created default config file at %s. Please fill in the required values.", path)
			data = []byte(defaultConfig)
		} else {
			return nil, fmt.Errorf("reading config file: %w", err)
		}
	}

	var cfg Config
	if err := yaml.Unmarshal(data, &cfg); err != nil {
		return nil, fmt.Errorf("parsing config: %w", err)
	}

	// Validate required fields
	if cfg.APIId == 0 {
		return nil, fmt.Errorf("api_id is not set in config")
	}
	if cfg.APIHash == "" {
		return nil, fmt.Errorf("api_hash is not set in config")
	}
	if cfg.MonitoredChat == 0 {
		return nil, fmt.Errorf("monitored_chat is not set in config")
	}
	if cfg.SessionPath == "" {
		return nil, fmt.Errorf("session_path is not set in config")
	}

	// Compile patterns from YAML
	if len(cfg.YAMLPatterns) > 0 {
		cfg.Patterns = make([]scamPattern, len(cfg.YAMLPatterns))
		for i, p := range cfg.YAMLPatterns {
			re, err := regexp.Compile(p.Pattern)
			if err != nil {
				return nil, fmt.Errorf("compiling pattern %q: %w", p.Name, err)
			}
			cfg.Patterns[i] = scamPattern{
				name:   p.Name,
				re:     re,
				weight: p.Weight,
			}
		}
		log.Printf("✓ Loaded %d scam patterns from YAML config:", len(cfg.Patterns))
		for _, p := range cfg.Patterns {
			log.Printf("  - %s (%s) (weight: %.1f)", p.name, p.re, p.weight)
		}
	}

	return &cfg, nil
}

type scamPattern struct {
	name   string
	re     *regexp.Regexp
	weight float64
}

type matchResult struct {
	score           float64
	matchedKeywords map[string][]string
}

var scamPatterns = []scamPattern{
	{
		name:   "seeking",
		re:     regexp.MustCompile(`(?i)(?:(?:several\s+people|seeking|looking\s+for)\s+(?:partner|(\d+\s*[-–]\s*\d+|\d+)|two)?)|(new\s+project)`),
		weight: 0.5,
	},
	{
		name:   "job_opportunity",
		re:     regexp.MustCompile(`(?i)(?:remote|earning|in\s+a)\s+(?:collaboration|growing\s+field|opportunity|work)`),
		weight: 0.4,
	},
	{
		name:   "money",
		re:     regexp.MustCompile(`(?i)(?:\d+[\s,]*(?:\$|€|₹|£)|(?:\$|€|₹|£)\s*\d+,?\d+)(\s+per\s+(week|day))?`),
		weight: 0.6,
	},
	{
		name:   "contact",
		re:     regexp.MustCompile(`((?i)(?:\bdm\b|(send|write).*\bprivate\b|get\s+in\s+touch)|@\w+)`),
		weight: 0.5,
	},
	{
		name:   "links",
		re:     regexp.MustCompile(`https://t\.me/\+\S+|https://t\.me/[a-z0-9_]+`),
		weight: 0.5,
	},
	{
		name:   "call_to_action",
		re:     regexp.MustCompile(`(?i)[⬇👇👆⬆💬🪂🎯⤵⤴📍]+|(?:click|tap|join|link|check)\s+(?:here|now|below|out)`),
		weight: 0.5,
	},
	{
		name:   "extended_latin",
		re:     regexp.MustCompile(`(?i)[äöüß]+`),
		weight: 0.2,
	},
}

// BotState holds shared runtime state that may be updated at runtime.
type BotState struct {
	cfg        *Config
	configPath string
	selfID     int64
	badWordRe  *regexp.Regexp
	mu         sync.RWMutex
}

func newBotState(cfg *Config, configPath string, selfID int64) *BotState {
	s := &BotState{cfg: cfg, configPath: configPath, selfID: selfID}
	s.compileBadWords()
	return s
}

// compileBadWords rebuilds the bad word regex from cfg.BadWords.
// Must be called with mu held (write).
func (s *BotState) compileBadWords() {
	if len(s.cfg.BadWords) == 0 {
		s.badWordRe = nil
		return
	}
	escaped := make([]string, len(s.cfg.BadWords))
	for i, w := range s.cfg.BadWords {
		escaped[i] = regexp.QuoteMeta(w)
	}
	s.badWordRe = regexp.MustCompile(`(?i)(` + strings.Join(escaped, "|") + `)`)
	log.Printf("✓ Bad word pattern: %s", s.badWordRe)
}

func (s *BotState) saveConfig() error {
	data, err := yaml.Marshal(s.cfg)
	if err != nil {
		return fmt.Errorf("marshaling config: %w", err)
	}
	return os.WriteFile(s.configPath, data, 0644)
}

func sendAdminReply(ctx context.Context, api *tg.Client, alertPeer *tg.InputPeerChannel, text string) {
	sender := message.NewSender(api)
	_, err := sender.To(alertPeer).Text(ctx, text)
	if err != nil {
		log.Printf("Failed to send admin reply: %v", err)
	}
}

func handleOwnerCommand(ctx context.Context, api *tg.Client, state *BotState, alertPeer *tg.InputPeerChannel, text string) {
	text = strings.TrimSpace(text)
	if !strings.HasPrefix(text, "/badword") {
		return
	}

	parts := strings.Fields(text)
	if len(parts) < 2 {
		sendAdminReply(ctx, api, alertPeer, "Usage: /badword <list|add|rem> [word]")
		return
	}

	subCmd := strings.ToLower(parts[1])

	state.mu.Lock()
	defer state.mu.Unlock()

	switch subCmd {
	case "list":
		if len(state.cfg.BadWords) == 0 {
			sendAdminReply(ctx, api, alertPeer, "No bad words configured.")
			return
		}
		var sb strings.Builder
		sb.WriteString("Bad words:\n")
		for _, w := range state.cfg.BadWords {
			sb.WriteString("  • ")
			sb.WriteString(w)
			sb.WriteString("\n")
		}
		reply := sb.String()
		if state.badWordRe != nil {
			reply += "\nPattern: " + state.badWordRe.String()
		}
		sendAdminReply(ctx, api, alertPeer, strings.TrimRight(reply, "\n"))

	case "add":
		if len(parts) < 3 {
			sendAdminReply(ctx, api, alertPeer, "Usage: /badword add <word>")
			return
		}
		word := parts[2]
		for _, w := range state.cfg.BadWords {
			if strings.EqualFold(w, word) {
				sendAdminReply(ctx, api, alertPeer, fmt.Sprintf("Word %q already exists.", word))
				return
			}
		}
		state.cfg.BadWords = append(state.cfg.BadWords, word)
		state.compileBadWords()
		if err := state.saveConfig(); err != nil {
			log.Printf("Failed to save config: %v", err)
			sendAdminReply(ctx, api, alertPeer, fmt.Sprintf("Added %q but failed to save config: %v", word, err))
			return
		}
		reply := fmt.Sprintf("Added bad word: %q", word)
		if state.badWordRe != nil {
			reply += "\nPattern: " + state.badWordRe.String()
		}
		sendAdminReply(ctx, api, alertPeer, reply)

	case "rem":
		if len(parts) < 3 {
			sendAdminReply(ctx, api, alertPeer, "Usage: /badword rem <word>")
			return
		}
		word := parts[2]
		newWords := make([]string, 0, len(state.cfg.BadWords))
		found := false
		for _, w := range state.cfg.BadWords {
			if strings.EqualFold(w, word) {
				found = true
				continue
			}
			newWords = append(newWords, w)
		}
		if !found {
			sendAdminReply(ctx, api, alertPeer, fmt.Sprintf("Word %q not found.", word))
			return
		}
		state.cfg.BadWords = newWords
		state.compileBadWords()
		if err := state.saveConfig(); err != nil {
			log.Printf("Failed to save config: %v", err)
		}
		patternStr := "none"
		if state.badWordRe != nil {
			patternStr = state.badWordRe.String()
		}
		sendAdminReply(ctx, api, alertPeer, fmt.Sprintf("Removed bad word: %q\nPattern: %s", word, patternStr))

	default:
		sendAdminReply(ctx, api, alertPeer, "Unknown subcommand. Usage: /badword <list|add|rem> [word]")
	}
}

func keywordMatchScore(input string, patterns []scamPattern) matchResult {
	matched := make(map[string][]string)
	totalScore := 0.0

	for _, p := range patterns {
		var matches []string
		for _, m := range p.re.FindAllString(input, -1) {
			m = strings.TrimSpace(m)
			if m != "" {
				matches = append(matches, m)
			}
		}
		if len(matches) > 0 {
			matched[p.name] = matches
			totalScore += p.weight
		}
	}

	score := math.Min(totalScore, 1.0)
	// Round to 2 decimal places
	score = math.Round(score*100) / 100

	return matchResult{score: score, matchedKeywords: matched}
}

func escapeMarkdown(text string) string {
	escapeChars := `\_*[]()~` + "`" + `>#+-=|{}.!`
	var result strings.Builder
	for _, r := range text {
		if strings.ContainsRune(escapeChars, r) {
			result.WriteRune('\\')
		}
		result.WriteRune(r)
	}
	return result.String()
}

func fullName(first, last string) string {
	return strings.TrimSpace(first + " " + last)
}

func formatUsername(username string) string {
	if username == "" {
		return ""
	}
	return " (@" + username + ")"
}

func notify(body string, cfg *Config, title string) {
	host := cfg.NtfyHost
	if !strings.HasSuffix(host, "/") {
		host += "/"
	}
	req, err := http.NewRequest("POST", host+cfg.NtfyTopic, bytes.NewBufferString(body))
	if err != nil {
		log.Printf("notify: creating request: %v", err)
		return
	}
	if title != "" {
		req.Header.Set("Title", title)
	}
	req.Header.Set("Priority", "5")
	req.Header.Set("Markdown", "yes")
	if cfg.NtfyToken != "" {
		encoded := base64.StdEncoding.EncodeToString([]byte(":" + cfg.NtfyToken))
		req.Header.Set("Authorization", "Basic "+encoded)
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		log.Printf("notify: sending request: %v", err)
		return
	}
	defer resp.Body.Close()
}

func resolveChannel(ctx context.Context, api *tg.Client, channelID int64) (*tg.InputPeerChannel, *tg.Channel, error) {
	req := &tg.MessagesGetDialogsRequest{
		OffsetPeer: &tg.InputPeerEmpty{},
		Limit:      100,
	}

	for {
		result, err := api.MessagesGetDialogs(ctx, req)
		if err != nil {
			return nil, nil, err
		}

		var (
			chats    []tg.ChatClass
			dialogs  []tg.DialogClass
			messages []tg.MessageClass
			done     bool
		)

		switch v := result.(type) {
		case *tg.MessagesDialogs:
			chats, dialogs, messages = v.Chats, v.Dialogs, v.Messages
			done = true
		case *tg.MessagesDialogsSlice:
			chats, dialogs, messages = v.Chats, v.Dialogs, v.Messages
		case *tg.MessagesDialogsNotModified:
			return nil, nil, fmt.Errorf("dialogs not modified")
		default:
			return nil, nil, fmt.Errorf("unexpected dialogs type: %T", result)
		}

		for _, chat := range chats {
			if channel, ok := chat.(*tg.Channel); ok {
				if channel.ID == channelID {
					return &tg.InputPeerChannel{
						ChannelID:  channel.ID,
						AccessHash: channel.AccessHash,
					}, channel, nil
				}
			}
		}

		if done || len(dialogs) == 0 {
			break
		}

		// Build offset from the last dialog for the next page.
		last, ok := dialogs[len(dialogs)-1].(*tg.Dialog)
		if !ok {
			break
		}
		msgID := last.TopMessage
		var msgDate int
		for _, m := range messages {
			if msg, ok := m.(*tg.Message); ok && msg.ID == msgID {
				msgDate = msg.Date
				break
			}
		}
		if msgDate == 0 {
			break
		}
		offsetPeer, err := toPeerInput(last.Peer, chats)
		if err != nil {
			break
		}
		req.OffsetDate = msgDate
		req.OffsetID = msgID
		req.OffsetPeer = offsetPeer
	}

	return nil, nil, fmt.Errorf("channel %d not found in dialogs", channelID)
}

func toPeerInput(peer tg.PeerClass, chats []tg.ChatClass) (tg.InputPeerClass, error) {
	switch p := peer.(type) {
	case *tg.PeerChannel:
		for _, c := range chats {
			if ch, ok := c.(*tg.Channel); ok && ch.ID == p.ChannelID {
				return &tg.InputPeerChannel{ChannelID: ch.ID, AccessHash: ch.AccessHash}, nil
			}
		}
		return nil, fmt.Errorf("channel %d not found in chats", p.ChannelID)
	case *tg.PeerChat:
		return &tg.InputPeerChat{ChatID: p.ChatID}, nil
	case *tg.PeerUser:
		return &tg.InputPeerUser{UserID: p.UserID}, nil
	default:
		return nil, fmt.Errorf("unknown peer type %T", peer)
	}
}

func extractMessage(update any) *tg.Message {
	switch u := update.(type) {
	case *tg.UpdateNewMessage:
		m, _ := u.Message.(*tg.Message)
		return m
	case *tg.UpdateNewChannelMessage:
		m, _ := u.Message.(*tg.Message)
		return m
	default:
		log.Printf("DEBUG: unknown update type %T", update)
		return nil
	}
}

func handleNewMessage(ctx context.Context, api *tg.Client, alertPeer *tg.InputPeerChannel, state *BotState, entities tg.Entities, update any) error {
	msg := extractMessage(update)
	if msg == nil {
		return nil
	}

	// Handle commands from the bot owner (messages sent by the connected account).
	if msg.Out {
		handleOwnerCommand(ctx, api, state, alertPeer, msg.Message)
		return nil
	}

	// Check if from monitored supergroup (channels in gotd terminology)
	peerChannel, ok := msg.PeerID.(*tg.PeerChannel)
	if !ok || int64(peerChannel.ChannelID) != state.cfg.MonitoredChat {
		return nil
	}
	chatID := int64(peerChannel.ChannelID)

	// Check bad words first — any match sets score to 1.0 immediately.
	state.mu.RLock()
	badWordRe := state.badWordRe
	state.mu.RUnlock()

	var result matchResult
	if badWordRe != nil && badWordRe.MatchString(msg.Message) {
		result = matchResult{
			score:           1.0,
			matchedKeywords: map[string][]string{"bad_words": badWordRe.FindAllString(msg.Message, -1)},
		}
		log.Printf("Bad word match in message, score set to 1.0")
	} else {
		// Use YAML patterns if available, otherwise use defaults.
		patterns := state.cfg.Patterns
		if len(patterns) == 0 {
			patterns = scamPatterns
			log.Printf("Using hardcoded scam patterns (%d patterns)", len(patterns))
		} else {
			log.Printf("Using YAML-loaded scam patterns (%d patterns)", len(patterns))
		}

		result = keywordMatchScore(msg.Message, patterns)

		if extended_latin, ok := result.matchedKeywords["extended_latin"]; ok {
			for range extended_latin {
				result.score += 0.1
				if result.score >= 1.0 {
					result.score = 1.0
					break
				}
			}
		}
	}

	if result.score < 1.0 {
		return nil
	}

	log.Printf("Matched message with score %.2f", result.score)
	for key, values := range result.matchedKeywords {
		for _, value := range values {
			log.Printf("  %s: %s", key, value)
		}
	}

	// Delete the message from supergroup
	channel, ok := entities.Channels[chatID]
	if !ok {
		return fmt.Errorf("channel %d not found in entities", chatID)
	}
	_, err := api.ChannelsDeleteMessages(ctx, &tg.ChannelsDeleteMessagesRequest{
		Channel: &tg.InputChannel{
			ChannelID:  chatID,
			AccessHash: channel.AccessHash,
		},
		ID: []int{msg.ID},
	})
	if err != nil {
		log.Printf("Failed to delete message: %v", err)
	}

	// Get sender info
	fromID, ok := msg.FromID.(*tg.PeerUser)
	if !ok {
		return fmt.Errorf("could not determine sender")
	}
	senderID := int64(fromID.UserID)

	user, ok := entities.Users[senderID]
	if !ok {
		return fmt.Errorf("user %d not found in entities", senderID)
	}

	displayName := fullName(user.FirstName, user.LastName)
	userDisplay := displayName
	if user.Username != "" {
		userDisplay += " (@" + user.Username + ")"
	} else {
		userDisplay += " (no username)"
	}

	// Restrict sender in supergroup
	_, err = api.ChannelsEditBanned(ctx, &tg.ChannelsEditBannedRequest{
		Channel: &tg.InputChannel{
			ChannelID:  chatID,
			AccessHash: channel.AccessHash,
		},
		Participant: &tg.InputPeerUser{
			UserID:     senderID,
			AccessHash: user.AccessHash,
		},
		BannedRights: tg.ChatBannedRights{
			SendMessages: true,
			SendMedia:    true,
			SendStickers: true,
			SendGifs:     true,
			UntilDate:    0,
		},
	})
	if err != nil {
		log.Printf("Failed to restrict user %d: %v", senderID, err)
	}

	chatName := channel.Title
	chatDisplay := escapeMarkdown(chatName) + escapeMarkdown(formatUsername(channel.Username))

	matchMessageHTML := fmt.Sprintf("🚨 Matched\n<b>Score</b>: %.2f\n<b>Chat</b>: %s (ID: %d)\n<b>User</b>: %s (ID: %d)\n",
		result.score, chatDisplay, chatID, userDisplay, senderID)

	if state.cfg.NtfyToken != "" || state.cfg.NtfyTopic != "" {
		plainMessage := fmt.Sprintf("🚨 Matched\nScore: %.2f\nChat: %s (ID: %d)\nUser: %s (ID: %d)\n\n%s",
			result.score, chatName+formatUsername(channel.Username), chatID, userDisplay, senderID, msg.Message)
		notify(plainMessage, state.cfg, fmt.Sprintf("Scam Alert: %s", chatName))
	}

	userResolver := func(id int64) (tg.InputUserClass, error) {
		return &tg.InputUserFromMessage{Peer: alertPeer, MsgID: 0, UserID: id}, nil
	}
	sender := message.NewSender(api)
	_, err = sender.To(alertPeer).StyledText(ctx, html.String(userResolver, matchMessageHTML))
	if err != nil {
		log.Printf("Failed to send alert message: %v", err)
	}

	return nil
}

func main() {
	ctx := context.Background()

	// Load configuration
	configPath := "config.yaml"
	if len(os.Args) > 1 {
		configPath = os.Args[1]
	}

	cfg, err := loadConfig(configPath)
	if err != nil {
		log.Fatalf("Failed to load config: %v", err)
	}

	dispatcher := tg.NewUpdateDispatcher()
	gapManager := updates.New(updates.Config{Handler: dispatcher})

	client := telegram.NewClient(cfg.APIId, cfg.APIHash, telegram.Options{
		UpdateHandler:  gapManager,
		SessionStorage: &session.FileStorage{Path: cfg.SessionPath},
	})

	err = client.Run(ctx, func(ctx context.Context) error {
		// Check if session exists; if not, authenticate
		info, _ := os.Stat(cfg.SessionPath)
		if info.Size() == 0 {
			// Session doesn't exist, prompt for credentials
			fmt.Print("Enter phone number: ")
			var phone string
			fmt.Scanln(&phone)

			fmt.Print("Enter 2FA password: ")
			pwBytes, err := term.ReadPassword(0)
			if err != nil {
				return fmt.Errorf("reading password: %w", err)
			}
			fmt.Println()
			password := string(pwBytes)

			flow := auth.NewFlow(
				auth.Constant(
					phone,
					password,
					auth.CodeAuthenticatorFunc(func(ctx context.Context, sentCode *tg.AuthSentCode) (string, error) {
						fmt.Print("Enter code: ")
						var code string
						fmt.Scanln(&code)
						return code, nil
					}),
				),
				auth.SendCodeOptions{},
			)

			if err := client.Auth().IfNecessary(ctx, flow); err != nil {
				return err
			}
		}

		api := client.API()

		// Resolve alert channel peer at startup
		alertPeer, alertChannel, err := resolveChannel(ctx, api, cfg.AlertChat)
		if err != nil {
			return fmt.Errorf("resolving alert channel: %w", err)
		}
		_, monitoredChannel, err := resolveChannel(ctx, api, cfg.MonitoredChat)
		if err != nil {
			return fmt.Errorf("resolving monitored channel: %w", err)
		}

		// Get self ID for state and updates manager
		status, err := client.Auth().Status(ctx)
		if err != nil {
			return err
		}

		state := newBotState(cfg, configPath, int64(status.User.ID))

		// Register message handler for private/group messages
		dispatcher.OnNewMessage(func(ctx context.Context, entities tg.Entities, update *tg.UpdateNewMessage) error {
			return handleNewMessage(ctx, api, alertPeer, state, entities, update)
		})

		// Register handler for channel messages (supergroups)
		dispatcher.OnNewChannelMessage(func(ctx context.Context, entities tg.Entities, update *tg.UpdateNewChannelMessage) error {
			return handleNewMessage(ctx, api, alertPeer, state, entities, update)
		})

		user := status.User
		log.Printf("✓ Logged in as: %s%s (ID: %d)", fullName(user.FirstName, user.LastName), formatUsername(user.Username), user.ID)
		log.Printf("✓ Monitoring chat: %s%s (ID: %d)", monitoredChannel.Title, formatUsername(monitoredChannel.Username), monitoredChannel.ID)
		log.Printf("✓ Sending alerts to: %s%s (ID: %d)", alertChannel.Title, formatUsername(alertChannel.Username), alertChannel.ID)
		return gapManager.Run(ctx, api, status.User.ID, updates.AuthOptions{})
	})

	if err != nil {
		log.Fatal(err)
	}
}