discord_bots/main.go

478 lines
12 KiB
Go

//go:generate stringer -type=FeatureName
package main
import (
"database/sql"
"encoding/json"
"flag"
"fmt"
"io"
"io/fs"
"log"
"math/rand/v2"
"os"
"os/signal"
"path/filepath"
"strings"
"time"
"cuelang.org/go/cue"
"cuelang.org/go/cue/cuecontext"
"github.com/bwmarrin/discordgo"
"github.com/joho/godotenv"
_ "embed"
_ "modernc.org/sqlite"
)
var botToken string
var confFile string
var dbFile string
var dataPath 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.StringVar(&dataPath, "data", "./", "file path for storing data")
flag.Parse()
}
//go:embed schema.cue
var schemaFile string
type FeatureName int
const (
UnknownFeature FeatureName = iota
VoiceChatAnnounceFeature
BirthdayAnnounceFeature
)
type Feature interface {
Name() FeatureName
}
func (c *Config) UnmarshalJSON(in []byte) error {
c2 := BaseConfig{}
c.featureMap = map[FeatureName]Feature{}
err := json.Unmarshal(in, &c2)
if err != nil {
return err
}
c.Server = c2.Server
for _, f := range c2.Features {
base := BaseFeature{}
err := json.Unmarshal(f, &base)
if err != nil {
fmt.Println(err.Error())
continue
}
switch base.Feature {
case VoiceChatAnnounceFeature:
f2 := VoiceChatAnnounce{}
err := json.Unmarshal(f, &f2)
if err != nil {
fmt.Println(err.Error())
continue
}
c.Features = append(c.Features, f2)
c.featureMap[VoiceChatAnnounceFeature] = f2
case BirthdayAnnounceFeature:
f2 := BirthdayAnnounce{}
err := json.Unmarshal(f, &f2)
if err != nil {
fmt.Println(err.Error())
continue
}
c.Features = append(c.Features, f2)
c.featureMap[BirthdayAnnounceFeature] = f2
default:
fmt.Println("Unknown feature")
}
}
return nil
}
type BaseConfig struct {
Server string
Features []json.RawMessage
}
type BaseFeature struct {
Feature FeatureName
Enabled bool
}
type AccounceFeature struct {
BaseFeature `json:",inline"`
AnnounceChannel string
}
type VoiceChatAnnounce struct {
AccounceFeature `json:",inline"`
JoinMessages []string
CleanUpDelay int
}
func (VoiceChatAnnounce) Name() FeatureName { return VoiceChatAnnounceFeature }
type BirthdayAnnounce struct {
AccounceFeature `json:",inline"`
Birthdays []Birthday
}
func (BirthdayAnnounce) Name() FeatureName { return BirthdayAnnounceFeature }
type Config struct {
Server string
Features []Feature
featureMap map[FeatureName]Feature
}
func (c Config) String() string {
s := strings.Builder{}
s.WriteString(fmt.Sprintf("Guild ID: %s\n", c.Server))
for _, f := range c.Features {
switch f2 := f.(type) {
case VoiceChatAnnounce:
if f2.Enabled {
s.WriteString("\tVoice Chat Accounce: ✅ \n")
s.WriteString(fmt.Sprintf("\t\tTo Channel: %s\n", f2.AnnounceChannel))
s.WriteString(fmt.Sprintf("\t\t %d Custom messages\n", len(f2.JoinMessages)))
} else {
s.WriteString("\tVoice Chat Accounce: ❌ \n")
}
case BirthdayAnnounce:
if f2.Enabled {
s.WriteString("\tBirthday Accounce: ✅ \n")
s.WriteString(fmt.Sprintf("\t\tTo Channel: %s\n", f2.AnnounceChannel))
s.WriteString(fmt.Sprintf("\t\t %d Birthdays to Announce\n", len(f2.Birthdays)))
} else {
s.WriteString("\tBirthday Accounce: ❌ \n")
}
}
}
return s.String()
}
type Birthday struct {
Name string //Optional
Member string
Date string // Format MM/DD
server string
}
var joinMessages = []string{
"<@%s> has joined <#%s>",
"<#%[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 Server struct {
*sql.DB
configs map[string]Config
}
func NewServer(dbFile string) (*Server, error) {
db, err := sql.Open("sqlite", filepath.Join(dataPath, 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, configs: map[string]Config{}}
return &s, nil
}
func (s *Server) AddConfig(c Config) {
c.featureMap = map[FeatureName]Feature{}
for _, f := range c.Features {
c.featureMap[f.Name()] = f
}
s.configs[c.Server] = c
log.Printf("Added server: \n%s\n", c)
}
func main() {
s, err := NewServer(dbFile)
if err != nil {
log.Fatalf("can't create server: %v", err)
}
ctx := cuecontext.New()
schema := ctx.CompileString(schemaFile)
vfs := os.DirFS(dataPath)
err = fs.WalkDir(vfs, ".", func(path string, d fs.DirEntry, err error) error {
if err != nil {
log.Fatal(err)
}
if filepath.Ext(d.Name()) != ".cue" {
return nil
}
if d.Name() == "schema.cue" {
return nil
}
f, err := vfs.Open(path)
if err != nil {
return err
}
data, err := io.ReadAll(f)
if err != nil {
return err
}
conf := ctx.CompileBytes(data, cue.Scope(schema))
if err := conf.Validate(); err != nil {
return fmt.Errorf("unable to validate config %q: %w", path, err)
}
cfg := Config{}
fields, err := conf.Fields()
if err != nil {
return err
}
fields.Next()
if err := fields.Value().Decode(&cfg); err != nil {
return fmt.Errorf("unable to decode config %q: %w", path, err)
}
s.AddConfig(cfg)
fmt.Printf("Loaded conf from %q\n", path)
return nil
})
if err != nil {
log.Printf("Error(s) loading conf files: %v", err)
}
time.Sleep(500 * time.Millisecond)
envs, err := godotenv.Read(filepath.Join(dataPath, confFile))
if err != nil {
log.Fatalf("Unable to get environment variables: %v", err)
}
// set values from env
botToken = envs["BOT_TOKEN"]
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()
if err = s.setupBirthdayWatch(ds); err != nil {
log.Fatalf("can't setup birthday watcher: %v", err)
}
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) {
guild, err := ds.Guild(m.VoiceState.GuildID)
if err != nil {
log.Printf("Unable to get Guild info: %s", err.Error())
return
}
server, serverConfigured := s.configs[guild.ID]
if !serverConfigured {
log.Printf("Server %q not configured", guild.Name)
return
}
vc, ok := server.featureMap[VoiceChatAnnounceFeature]
if !ok {
log.Printf("Server %q not configured for voice chat announcements", guild.Name)
return
}
config := vc.(VoiceChatAnnounce)
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
}
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
}
ch, err := ds.Channel(config.AnnounceChannel)
if err != nil {
log.Printf("Unable to get announce channel for Guild: %s", m.GuildID)
return
}
messages := append(joinMessages, config.JoinMessages...)
msg := fmt.Sprintf(messages[rand.IntN(len(messages))], 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(ch.ID, msg)
if err != nil {
log.Printf("unable to send message: %s", err)
}
time.AfterFunc(time.Second*time.Duration(config.CleanUpDelay), func() {
ds.ChannelMessageDelete(sent.ChannelID, sent.ID)
})
}
func (s *Server) setupBirthdayWatch(ds *discordgo.Session) error {
birthdays := map[string][]Birthday{} // Use a list in case there are collisions (See birthday paradox)
for _, guild := range s.configs {
ba, ok := guild.featureMap[BirthdayAnnounceFeature]
if !ok {
continue
}
config, ok := ba.(BirthdayAnnounce)
if !ok || !config.Enabled {
continue
}
for _, b := range config.Birthdays {
b.server = guild.Server
birthdays[b.Date] = append(birthdays[b.Date], b)
}
}
targetHour := 7
targetMinute := 35
now := time.Now()
location, err := time.LoadLocation("America/Los_Angeles")
if err != nil {
return fmt.Errorf("unable to get timezone data: %w", err)
}
targetTime := time.Date(now.Year(), now.Month(), now.Day(), targetHour, targetMinute, 0, 0, location)
if now.After(targetTime) {
targetTime = targetTime.Add(24 * time.Hour)
}
initialDelay := targetTime.Sub(now)
log.Printf("Waiting for %s for starting birthday watcher", initialDelay)
birthdayMatcher := func(ds *discordgo.Session) {
dateString := time.Now().Format("01/02")
log.Printf("Looking for matching birthdays on %s", dateString)
if bd, ok := birthdays[dateString]; ok {
for _, b := range bd {
log.Printf("It is %q's birthday today (%s)", b.Member, b.Date)
ba, ok := s.configs[b.server].featureMap[BirthdayAnnounceFeature]
if !ok {
continue
}
config, ok := ba.(BirthdayAnnounce)
if !ok || !config.Enabled {
continue
}
if err := announceBirthday(ds, config.AnnounceChannel, b.Member); err != nil {
log.Printf("Error w/ announcing: %v", err)
}
}
}
}
time.AfterFunc(initialDelay, func() {
log.Printf("Birthday watcher initiated")
// Run now, but also set up a daily job
birthdayMatcher(ds)
c := time.Tick(24 * time.Hour)
for range c {
birthdayMatcher(ds)
}
})
return nil
}
func announceBirthday(ds *discordgo.Session, channelID, userId string) error {
_, err := ds.ChannelMessageSend(channelID, fmt.Sprintf("Happy Birthday to <@%s>!!", userId))
if err != nil {
return fmt.Errorf("unable to send message: %w", err)
}
return nil
}