irc_bot/main.go

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	"log"

	"strconv"

	irc "github.com/thoj/go-ircevent"
)

var channel = "#ctf"
var sendHi = ""
var writeHistory = false

const (
	pattern   = "(?i)\\b(cat|gato|miau|meow|garfield|lolcat)[s|z]{0,1}\\b"
	nyastat   = "nyastat"
	msgPrefix = "I love cats! Here's a fact: %s ^_^"
	gifPrefix = "Meow! Here's a gif: %s ^_^"
)

// environment variables
const (
	userEnv        = "IRC_USER"
	nickEnv        = "IRC_NICK"
	himsgEnv       = "IRC_SAY_HI"
	passEnv        = "IRC_PASS"
	chanEnv        = "IRC_CHAN"
	urlEnv         = "IRC_URL"
	invitionsEnv   = "IRC_ENABLE_INVITIONS"
	chanHistoryEnv = "IRC_CHAN_HISTORY"
	dbPathEnv      = "IRC_DB_PATH"
)

type catFact struct {
	Fact   string `json:"fact"`
	Length int    `json:"length"`
}

var (
	re          = regexp.MustCompile(pattern)
	catFactsURL = "http://catfact.ninja/fact"
)

type stats struct {
	userMessages map[string]uint
	invites      map[string]bool
	history      *history
	mux          *sync.Mutex
}

func newStats() *stats {
	return &stats{
		userMessages: make(map[string]uint),
		invites:      make(map[string]bool),
		mux:          &sync.Mutex{},
	}
}

func (s *stats) addUser(user string) {
	s.mux.Lock()
	defer s.mux.Unlock()

	if v, ok := s.userMessages[user]; ok {
		log.Printf("User %s writes %d message", user, v+1)
		s.userMessages[user]++
		return
	}
	log.Printf("First message from user %s", user)
	s.userMessages[user] = 1
}

func (s *stats) online(e *irc.Event) int {
	e.Connection.Mutex.Lock()
	defer e.Connection.Mutex.Unlock()

	return -1
}

func (s *stats) realStats(e *irc.Event) {
	e.Connection.Privmsg(e.Nick, "🙀🙀🙀 WoW 🙀🙀🙀 NYAStat 🙀🙀🙀")
	t := "User %s sent %d messages"
	for u, v := range s.userMessages {
		e.Connection.Privmsg(e.Nick, fmt.Sprintf(t, u, v))
	}
}

func (s *stats) printStats(e *irc.Event) {
	s.online(e)
	s.realStats(e)
}

func (s *stats) invitionRequest(e *irc.Event) {
	s.mux.Lock()
	defer s.mux.Unlock()

	if !isInvitionsEnabled() {
		return
	}

	if sended, ok := s.invites[e.Nick]; !ok || !sended {
		log.Printf("Sending invite request to %s", e.Nick)

		s.invites[e.Nick] = true
		e.Connection.Privmsg(e.Nick, fmt.Sprintf("Привет, %s! В тиму 0х0 нужны два опытных человека. + в приват @difrex", e.Nick))
	}
}

func addCallbacks(irccon *irc.Connection, s *stats) {
	irccon.AddCallback("001", func(e *irc.Event) {
		log.Println("Welcome message:", e.Message())

		// Login
		e.Connection.Privmsg("nickserv", "identify "+os.Getenv(passEnv))

		// Join to channel
		e.Connection.Join(channel)

		e.Connection.Action(channel, "=^_^= https://gitea.difrex.ru/CTF2021/irc_bot")
		e.Connection.Action(channel, "😸🦄😸 Now with HiStOrY!!! 😸🦄😸")

		// Greetings
		if sendHi != "" {
			e.Connection.Privmsg(channel, `Hi!`)
			e.Connection.Privmsg(channel, `I'm a cat facts bot!`)
			e.Connection.Privmsg(channel, `Source code and issue tracker: https://gitea.difrex.ru/CTF2021/irc_bot`)
			e.Connection.Privmsg(channel, `Now with gifs! ^_^`)
		}
	})

	irccon.AddCallback("JOIN", s.invitionRequest)

	irccon.AddCallback("PRIVMSG", func(e *irc.Event) {
		log.Printf("Message from %s received: %s", e.Nick, e.Message())
		s.addUser(e.Nick)
		doCommand(e.Message(), e, s)
	})
}

func main() {
	nick := os.Getenv(nickEnv)
	irccon := irc.IRC(nick, os.Getenv(userEnv))
	if irccon == nil {
		log.Fatal("Connection is nil")
	}

	s := newStats()
	history, err := Open(os.Getenv(dbPathEnv))
	if err != nil {
		writeHistory = false
	}
	s.history = history

	if history != nil {
		defer history.db.Close()
	}

	irccon.VerboseCallbackHandler = false
	irccon.Debug = false
	irccon.UseTLS = false

	addCallbacks(irccon, s)

	if err := irccon.Connect(os.Getenv(urlEnv)); err != nil {
		fmt.Printf("Err %s", err)
		return
	}
	irccon.Loop()
}

func isHelpCommand(command string, e *irc.Event) bool {
	if e.Nick != channel && strings.ToLower(command) == "help" {
		return true
	}
	return false
}

func isHistoryCommand(command string, e *irc.Event) bool {
	if e.Nick != channel && strings.HasPrefix(e.Message(), "history") {
		return true
	}
	return false
}

func parseHistoryCommandArgs(e *irc.Event) [2]string {
	var args [2]string

	message := strings.Split(e.Message(), " ")
	if len(message) == 3 {
		args[0] = message[1]
		args[1] = message[2]
	}
	if len(message) == 2 {
		args[0] = message[1]
	}

	return args
}

func showLogs(e *irc.Event, s *stats) {
	var logs []LogLine
	args := parseHistoryCommandArgs(e)

	if args[0] == "all" {
		l, err := s.history.GetAllLog()
		if err != nil {
			log.Println("Error:", err)
			return
		}
		logs = l
	}

	if args[0] != "" && args[1] != "" {
		l, err := s.history.GetRangedLogs(args[0], args[1])
		if err != nil {
			log.Println("Error:", err)
			return
		}
		logs = l
	}

	if len(logs) == 0 {
		e.Connection.Privmsg(e.Nick, ":( such empty")
	}

	e.Connection.Privmsg(e.Nick, "WOW!!!  HiStOrY!!!")
	for _, l := range logs {
		e.Connection.Privmsg(e.Nick, fmt.Sprintf("%s [%s] %s", l.Time.Format("2006-01-02T15:04:05"), l.Nick, l.Message))
	}
}

func showHelp(e *irc.Event) {
	e.Connection.Privmsg(e.Nick, "HELP -- show this help")
	e.Connection.Privmsg(e.Nick, "cat|meow|etc -- show cat fact or gif")
	e.Connection.Privmsg(e.Nick, "nyastat -- show #ctf channel stats")
	e.Connection.Privmsg(e.Nick, "history -- show #ctf channel history logs")
	e.Connection.Privmsg(e.Nick, "        -- EXAMPLE")
	e.Connection.Privmsg(e.Nick, "history all -- show all logs")
	e.Connection.Privmsg(e.Nick, "history 2021-01-27T12:19:55 2021-01-27T12:20:00 -- show date ranged logs")
}

func doCommand(command string, e *irc.Event, s *stats) {
	if ok := statCommand(command); ok {
		s.printStats(e)
		return
	}

	if isHelpCommand(command, e) {
		showHelp(e)
		return
	}

	if isHistoryCommand(command, e) {
		showLogs(e, s)
		return
	}

	if ok := checkCommad(command, e.Connection); !ok {
		err := s.history.WriteLog(&LogLine{
			Nick:    e.Nick,
			Message: e.Message(),
			Time:    time.Now(),
		})
		if err != nil {
			log.Println("Error:", err)
		}
		return
	}

	rand.Seed(time.Now().UnixNano())
	i := rand.Intn(50)
	fmt.Println(i)
	if i > 30 {
		catGif(e)
		return
	}

	catFacts(e)
}

func statCommand(command string) bool {
	if command == nyastat {
		return true
	}
	return false
}

func checkCommad(command string, con *irc.Connection) bool {
	return re.MatchString(command)
}

func catGif(e *irc.Event) (string, error) {
	res, err := http.Get("http://thecatapi.com/api/images/get?format=src&type=gif")
	if err != nil {
		return "", err
	}
	log.Println("Send gif")
	e.Connection.Privmsg(e.Nick, fmt.Sprintf(gifPrefix, res.Request.URL.String()))
	return fmt.Sprintf(gifPrefix, res.Request.URL.String()), nil
}

func getJson(url string, v interface{}) error {
	res, err := http.Get(url)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}

	return json.Unmarshal(body, v)
}

func catFacts(e *irc.Event) (string, error) {
	log.Println("Send fact")
	data := &catFact{}
	err := getJson(catFactsURL, data)
	if err != nil {
		return "", err
	}

	if len(data.Fact) == 0 {
		return "", nil
	}

	e.Connection.Privmsg(e.Nick, fmt.Sprintf(msgPrefix, data.Fact))
	return fmt.Sprintf(msgPrefix, data.Fact), nil
}

func getChanName(e *irc.Event) string {
	c := channel
	if _, err := strconv.Atoi(e.Nick); err != nil {
		c = e.Nick
	}

	return c
}

func isFetureEnabled(feature string) bool {
	if feature != "" && strings.ToLower(feature) == "yes" {
		return true
	}

	return false
}

func isInvitionsEnabled() bool {
	return isFetureEnabled(os.Getenv(invitionsEnv))
}

func isHistoryEnabled() bool {
	return isFetureEnabled(os.Getenv(chanHistoryEnv))
}

func getNick() string {
	return os.Getenv(nickEnv)
}

func init() {
	ch := os.Getenv(chanEnv)
	if ch != "" {
		channel = ch
	}

	sendHi = os.Getenv(himsgEnv)

	writeHistory = isFetureEnabled(os.Getenv(chanHistoryEnv))
}