discord_bots/main.go

package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"math/rand/v2"
	"os"
	"os/signal"
	"strconv"
	"time"

	"github.com/bwmarrin/discordgo"
	"github.com/joho/godotenv"

	_ "modernc.org/sqlite"
)

var botToken string
var announceChannelName string
var cleanDelay = 30

var confFile string
var dbFile string

func init() {
	flag.StringVar(&confFile, "conf", ".env", ".env file w/ config variables")
	flag.StringVar(&dbFile, "db", "data.db", "db to store logs of events in")
	flag.Parse()
}

var joinMessages = []string{
	"<@%s> has joined <#%s>",
	"<@%s> has joined <#%s>; Join in for some nerd talk",
	"<#%[2]s> is the place to be! <@%[1]s> just joined",
	"<#%[2]s> just got a bit cooler, <@%[1]s> is now in",
	"<@%s> is hanging out in <#%s>",
}

type Channel struct {
	GuildID, ChannelID string
}

type Server struct {
	*sql.DB
}

func NewServer(dbFile string) (*Server, error) {
	db, err := sql.Open("sqlite", dbFile+"?_time_format=sqlite")
	if err != nil {
		return nil, fmt.Errorf("unable to open db %q: %w", dbFile, err)
	}
	tx, err := db.Begin()
	if err != nil {
		return nil, fmt.Errorf("unable to work on db %q: %w", dbFile, err)
	}
	tx.Exec(`CREATE TABLE IF NOT EXISTS voice_stats (
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	guildID INTEGER,
	guildName TEXT,
	memberId INTEGER,
	memberName TEXT,
	channelID INTEGER,
	channelName TEXT,
	action TEXT,
	timestamp TIMESTAMP  
	);`)
	err = tx.Commit()
	if err != nil {
		return nil, fmt.Errorf("unable to create table: %w", err)
	}
	s := Server{DB: db}
	return &s, nil
}

func main() {
	envs, err := godotenv.Read(confFile)
	if err != nil {
		log.Fatalf("Unable to get environment variables: %v", err)
	}
	s, err := NewServer(dbFile)
	if err != nil {
		log.Fatalf("can't create server: %v", err)
	}

	// set values from env
	botToken = envs["BOT_TOKEN"]
	announceChannelName = envs["ANNOUNCE_CHANNEL"]
	if delay, ok := envs["CLEAN_DELAY"]; ok {
		if d2, err := strconv.Atoi(delay); err == nil {
			cleanDelay = d2
		}
	}

	ds, err := discordgo.New("Bot " + botToken)
	if err != nil {
		log.Fatal(err)
	}

	// Log all of the servers tht this bot is installed in.
	ds.AddHandler(func(ds *discordgo.Session, r *discordgo.Ready) {
		for _, g := range r.Guilds {
			g2, _ := ds.Guild(g.ID)
			log.Printf("Bot is connected to %q.", g2.Name)
		}

	})

	// Add different capability handlers:
	ds.AddHandler(s.voiceStatus)

	err = ds.Open()
	if err != nil {
		log.Fatalf("Cannot open the session: %v", err)
	}
	defer ds.Close()

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt)
	<-stop
}

type VoiceState int

const (
	Joining VoiceState = iota
	Switching
	Leaving
)

func (s VoiceState) String() string {
	return []string{"Joined", "Switched", "Left"}[s]
}

// Handle for when someone joins a voice channel - send notification about them joining.
func (s *Server) voiceStatus(ds *discordgo.Session, m *discordgo.VoiceStateUpdate) {
	var state VoiceState
	switch {
	case len(m.ChannelID) == 0:
		state = Leaving
	case m.BeforeUpdate != nil:
		if m.BeforeUpdate.ChannelID == m.ChannelID {
			// Action like muting caused new voice state, but in the same channel.
			return
		}
		state = Switching
	}

	guild, _ := ds.Guild(m.VoiceState.GuildID)

	switch state {
	case Leaving, Switching:
		channel, _ := ds.Channel(m.BeforeUpdate.ChannelID)
		log.Default().Printf("%q has left %q in %q", m.Member.DisplayName(), channel.Name, guild.Name)

		_, err := s.Exec("INSERT INTO voice_stats (guildID, guildName, memberID, memberName, channelID, channelName, action, timestamp) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", guild.ID, guild.Name, m.VoiceState.UserID, m.VoiceState.Member.DisplayName(), channel.ID, channel.Name, "LEFT", time.Now())
		if err != nil {
			log.Printf("unable to save log to db: %s", err)
		}
	}

	switch state {
	case Joining, Switching:
		channel, _ := ds.Channel(m.ChannelID)
		log.Default().Printf("%q has joined %q in %q", m.Member.DisplayName(), channel.Name, guild.Name)

		_, err := s.Exec("INSERT INTO voice_stats (guildID, guildName, memberID, memberName, channelID, channelName, action, timestamp) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", guild.ID, guild.Name, m.VoiceState.UserID, m.VoiceState.Member.DisplayName(), channel.ID, channel.Name, "JOIN", time.Now())
		if err != nil {
			log.Printf("unable to save log to db: %s", err)
		}
	}

	switch state {
	case Leaving:
		return
	}

	channels, err := ds.GuildChannels(m.GuildID)
	if err != nil {
		log.Printf("Unable to get channels for Guild: %s", m.GuildID)
		return
	}
	var announceChannel *discordgo.Channel
	for _, c := range channels {
		if c.Name == announceChannelName {
			announceChannel = c
		}
	}
	if announceChannel == nil {
		log.Printf("Unable to get announce channel for Guild: %s", m.GuildID)
		return
	}
	msg := fmt.Sprintf(joinMessages[rand.IntN(len(joinMessages))], m.UserID, m.ChannelID)
	switch state {
	case Switching:
		msg = fmt.Sprintf("<@%s> has left <#%s> to join <#%s>", m.UserID, m.BeforeUpdate.ChannelID, m.ChannelID)
	}

	sent, err := ds.ChannelMessageSend(announceChannel.ID, msg)
	if err != nil {
		log.Printf("unable to send message: %s", err)
	}

	time.AfterFunc(time.Second*time.Duration(cleanDelay), func() {
		ds.ChannelMessageDelete(sent.ChannelID, sent.ID)
	})
}