FlameInTheDark / dtbot

main.go

package main

import (
	"bytes"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"gopkg.in/robfig/cron.v2"

	"github.com/FlameInTheDark/dtbot/bot"
	"github.com/FlameInTheDark/dtbot/cmd"
	"github.com/bwmarrin/discordgo"
)

var (
	conf *bot.Config
	// CmdHandler bot command handler.
	CmdHandler *bot.CommandHandler
	// Sessions bot session manager
	Sessions        *bot.SessionManager
	botId           string

var botId should be botID

	youtube         *bot.Youtube
	botMsg          *bot.BotMessages
	dataType        *bot.DataType
	dbWorker        *bot.DBWorker
	guilds          *bot.GuildsMap
	botCron         *cron.Cron
	twitch          *bot.Twitch
	albUpdater      *bot.AlbionUpdater
	blacklist       *bot.BlackListStruct
	messagesCounter int
)

func main() {
	botCron = cron.New()
	conf = bot.LoadConfig()
	CmdHandler = bot.NewCommandHandler()
	registerCommands()
	Sessions = bot.NewSessionManager()
	youtube = &bot.Youtube{Conf: conf}
	botMsg = bot.NewMessagesMap()
	dataType = bot.NewDataType()
	discord, err := discordgo.New("Bot " + os.Getenv("BOT_TOKEN"))
	if err != nil {
		fmt.Println("Create session error, ", err)
		return
	}
	usr, err := discord.User("@me")
	if err != nil {
		fmt.Println("Error obtaining account details,", err)
		return
	}
	botId = usr.ID
	discord.AddHandler(func(discord *discordgo.Session, ready *discordgo.Ready) {
		_ = discord.UpdateStatus(0, conf.General.Game)
		guilds := discord.State.Guilds
		fmt.Println("Ready with", len(guilds), "guilds.")
	})

	err = discord.Open()
	if err != nil {
		fmt.Printf("Connection open error: %v", err)
		return
	}
	defer discord.Close()
	fmt.Println("Bot is now running.")

	sc := make(chan os.Signal, 1)
	signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
	dbWorker = bot.NewDBSession(conf.General.DatabaseName)
	guilds = dbWorker.InitGuilds(discord, conf)
	botCron.Start()
	defer botCron.Stop()
	defer dbWorker.DBSession.Close()
	twitch = bot.TwitchInit(discord, conf, dbWorker)
	albUpdater = bot.AlbionGetUpdater(dbWorker)
	blacklist = dbWorker.GetBlacklist()
	go BotUpdater(discord)
	// Init command handler
	discord.AddHandler(guildAddHandler)
	discord.AddHandler(commandHandler)
	discord.AddHandler(joinHandler)
	onStart()
	<-sc
}

// Handle new users
func joinHandler(discord *discordgo.Session, e *discordgo.GuildMemberAdd) {
	if _, ok := guilds.Guilds[e.GuildID]; !ok {
		dbWorker.InitNewGuild(e.GuildID, conf, guilds)
	} else {
		bot.Greetings(discord, e, guilds.Guilds[e.GuildID])
	}
}

// Handle new guilds
func guildAddHandler(discord *discordgo.Session, e *discordgo.GuildCreate) {
	if _, ok := guilds.Guilds[e.ID]; !ok {
		dbWorker.InitNewGuild(e.ID, conf, guilds)
	}
	emb := bot.NewEmbed("").
		Field(conf.GetLocaleLang("bot_joined_title", guilds.Guilds[e.ID].Language), conf.GetLocaleLang("bot_joined_text", guilds.Guilds[e.ID].Language), false)
	_, _ = discord.ChannelMessageSendEmbed(e.OwnerID, emb.GetEmbed())
}

// Handle discord messages
func commandHandler(discord *discordgo.Session, message *discordgo.MessageCreate) {
	if blacklist.CheckGuild(message.GuildID) || blacklist.CheckUser(message.Author.ID) {
		return
	}
	messagesCounter++
	user := message.Author
	if user.ID == botId || user.Bot {
		return
	}
	args := strings.Split(message.Content, " ")
	name := strings.ToLower(args[0])
	command, found := CmdHandler.Get(name)
	if !found {
		return
	}

	var permission = true
	var msg string
	// Checking permissions
	perm, err := discord.State.UserChannelPermissions(botId, message.ChannelID)
	if err != nil {
		msg = fmt.Sprintf("Error whilst getting bot permissions %v\n", err)
		permission = false
	} else {
		if perm&discordgo.PermissionSendMessages != discordgo.PermissionSendMessages ||
			perm&discordgo.PermissionAttachFiles != discordgo.PermissionAttachFiles {
			msg = "Permissions denied"
			permission = false
		}
	}

	channel, err := discord.State.Channel(message.ChannelID)
	if err != nil {
		fmt.Println("Error getting channel,", err)
		return
	}
	guild, err := discord.State.Guild(channel.GuildID)
	if err != nil {
		fmt.Println("Error getting guild,", err)
		return
	}

	if permission {
		ctx := bot.NewContext(
			botId,
			discord,
			guild,
			channel,
			user,
			message,
			conf,
			CmdHandler,
			Sessions,
			youtube,
			botMsg,
			dataType,
			dbWorker,
			guilds,
			botCron,
			twitch,
			albUpdater,
			blacklist)
		ctx.Args = args[1:]
		c := *command
		c(*ctx)
	} else {
		dbWorker.Log("Message", guild.ID, msg)
		query := []byte(fmt.Sprintf("logs,server=%v module=\"%v\"", guild.ID, "message"))
		addr := fmt.Sprintf("%v/write?db=%v", conf.Metrics.Address, conf.Metrics.Database)
		r := bytes.NewReader(query)
		_, _ = http.Post(addr, "", r)
	}
}

// Adds bot commands
func registerCommands() {
	CmdHandler.Register("!r", cmd.RadioCommand)
	CmdHandler.Register("!w", cmd.WeatherCommand)
	CmdHandler.Register("!ww", cmd.WeatherWeekCommand)
	CmdHandler.Register("!t", cmd.TranslateCommand)
	CmdHandler.Register("!n", cmd.NewsCommand)
	CmdHandler.Register("!c", cmd.CurrencyCommand)
	CmdHandler.Register("!y", cmd.YoutubeCommand)
	CmdHandler.Register("!v", cmd.VoiceCommand)
	CmdHandler.Register("!b", cmd.BotCommand)
	CmdHandler.Register("!play", cmd.YoutubeShortCommand)
	CmdHandler.Register("!d", cmd.DebugCommand)
	CmdHandler.Register("!p", cmd.PollCommand)
	CmdHandler.Register("!m", cmd.YandexmapCommand)
	CmdHandler.Register("!dice", cmd.DiceCommand)
	CmdHandler.Register("!help", cmd.HelpCommand)
	CmdHandler.Register("!cron", cmd.CronCommand)
	CmdHandler.Register("!geoip", cmd.GeoIPCommand)
	CmdHandler.Register("!twitch", cmd.TwitchCommand)
	CmdHandler.Register("!greetings", cmd.GreetingsCommand)
	CmdHandler.Register("!alb", cmd.AlbionCommand)
	CmdHandler.Register("!slap", cmd.SlapCommand)
	CmdHandler.Register("!fu", cmd.FUCommand)
}

// MetricsSender sends metrics to InfluxDB and another services

comment on exported function BotUpdater should be of the form "BotUpdater ..."

func BotUpdater(d *discordgo.Session) {
	for {
		var vregions = make(map[string]int)
		go twitch.Update()
		go albUpdater.Update(d, dbWorker, conf)
		go clearSessions(d, Sessions)
		// Calculating users count
		usersCount := 0
		for _, g := range d.State.Guilds {
			vregions[g.Region] += 1

should replace vregions[g.Region] += 1 with vregions[g.Region]++

			if !blacklist.CheckGuild(g.ID) {
				usersCount += g.MemberCount
			}
		}
		// Metrics counters
		queryCounters := []byte(fmt.Sprintf("counters guilds=%d,messages=%d,users=%d,voices=%d", len(d.State.Guilds), messagesCounter, usersCount, Sessions.Count()))
		addrCounters := fmt.Sprintf("%v/write?db=%v&u=%v&p=%v",
			conf.Metrics.Address, conf.Metrics.Database, conf.Metrics.User, conf.Metrics.Password)
		rCounters := bytes.NewReader(queryCounters)
		_, _ = http.Post(addrCounters, "", rCounters)

		// Voice region metrics
		for r, c := range vregions {
			queryCounters := []byte(fmt.Sprintf("region_%s count=%d", r, c))
			addrCounters := fmt.Sprintf("%v/write?db=%v&u=%v&p=%v",
				conf.Metrics.Address, conf.Metrics.Database, conf.Metrics.User, conf.Metrics.Password)
			rCounters := bytes.NewReader(queryCounters)
			_, _ = http.Post(addrCounters, "", rCounters)
		}

		// Bot lists
		if conf.DBL.Token != "" {
			sendDBL(conf.DBL.BotID, conf.DBL.Token, len(d.State.Guilds))
		}
		messagesCounter = 0
		time.Sleep(time.Minute)
	}
}

func sendDBL(botID, token string, guilds int) {
	client := &http.Client{
		Timeout: time.Duration(1) * time.Second,
	}
	query := url.Values{}
	query.Add("server_count", fmt.Sprintf("%v", guilds))
	req, _ := http.NewRequest("POST",
		fmt.Sprintf("https://discordbots.org/api/bots/%v/stats",
			botID), strings.NewReader(query.Encode()))
	req.Header.Add("Authorization", token)
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	_, _ = client.Do(req)
}

func onStart() {
	queryCounters := []byte(fmt.Sprintf("starts value=1"))
	addrCounters := fmt.Sprintf("%v/write?db=%v&u=%v&p=%v",
		conf.Metrics.Address, conf.Metrics.Database, conf.Metrics.User, conf.Metrics.Password)
	rCounters := bytes.NewReader(queryCounters)
	_, _ = http.Post(addrCounters, "", rCounters)
}

func clearSessions(d *discordgo.Session, s *bot.SessionManager) {
	gids := s.GetGuilds()
	for _, gid := range gids {
		cs := s.GetByGuild(gid)
		g, gerr := d.Guild(gid)
		if gerr != nil {
			continue
		}
		var ok bool
		for _, vs := range g.VoiceStates {
			if vs.ChannelID == cs.ChannelID {
				if vs.UserID != botId {
					ok = true
				}
			}
		}
		if !ok && cs != nil {
			s.Leave(d, *cs)
		}
	}

	for _, sess := range s.GetAll() {
		if !sess.IsOk() {
			s.Leave(d, *sess)
		}
	}

	for i,vc := range d.VoiceConnections {
		if _,b := s.GetByChannel(vc.ChannelID); b == false {
			err := vc.Disconnect()
			if err != nil {
				log.Println(err)
			}
			vc.Close()
			delete(d.VoiceConnections, i)
		}
	}
}