karaokards/internal/irc/irc.go

package irc

import (
	"bufio"
	"encoding/json"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/textproto"
	"regexp"
	"strings"
	"time"

	rgb "github.com/foresthoffman/rgblog"
	scribble "github.com/nanobox-io/golang-scribble"
	uuid "github.com/google/uuid"
)

const PSTFormat = "Jan 2 15:04:05 PST"

// Regex for parsing PRIVMSG strings.
//
// First matched group is the user's name, second is the channel? and the third matched group is the content of the
// user's message.
var MsgRegex *regexp.Regexp = regexp.MustCompile(`^:(\w+)!\w+@\w+\.tmi\.twitch\.tv (PRIVMSG) #(\w+)(?: :(.*))?$`)

// Regex for parsing user commands, from already parsed PRIVMSG strings.
//
// First matched group is the command name and the second matched group is the argument for the
// command.
var CmdRegex *regexp.Regexp = regexp.MustCompile(`^!(\w+)\s?(\w+)?`)

type OAuthCred struct {

	// The bot account's OAuth password.
	Password string `json:"password,omitempty"`

	// The developer application client ID. Used for API calls to Twitch.
	ClientID string `json:"client_id,omitempty"`
}

type KardBot struct {
	Channel     string
	conn        net.Conn
	Credentials *OAuthCred
	MsgRate     time.Duration
	Name        string
	Port        string
	PrivatePath string
	Server      string
	startTime   time.Time
	Prompts     []string
	Database	scribble.Driver
	ChannelData map[string]ChannelData
}

type ChannelData struct {
	Name          string `json:"name"`
	AdminKey	  string `json:"value,omitempty"`
	Command       string `json:"customcommand,omitempty"`
	ExtraStrings  string `json:"extrastrings,omitempty"`
	JoinTime      time.Time `json:"jointime"`
}

// Connects the bot to the Twitch IRC server. The bot will continue to try to connect until it
// succeeds or is forcefully shutdown.
func (bb *KardBot) Connect() {
	var err error
	rgb.YPrintf("[%s] Connecting to %s...\n", TimeStamp(), bb.Server)

	// makes connection to Twitch IRC server
	bb.conn, err = net.Dial("tcp", bb.Server+":"+bb.Port)
	if nil != err {
		rgb.YPrintf("[%s] Cannot connect to %s, retrying.\n", TimeStamp(), bb.Server)
		bb.Connect()
		return
	}
	rgb.YPrintf("[%s] Connected to %s!\n", TimeStamp(), bb.Server)
	bb.startTime = time.Now()
}

// Officially disconnects the bot from the Twitch IRC server.
func (bb *KardBot) Disconnect() {
	bb.conn.Close()
	upTime := time.Now().Sub(bb.startTime).Seconds()
	rgb.YPrintf("[%s] Closed connection from %s! | Live for: %fs\n", TimeStamp(), bb.Server, upTime)
}

// Listens for and logs messages from chat. Responds to commands from the channel owner. The bot
// continues until it gets disconnected, told to shutdown, or forcefully shutdown.
func (bb *KardBot) HandleChat() error {
	rgb.YPrintf("[%s] Watching #%s...\n", TimeStamp(), bb.Channel)

	// reads from connection
	tp := textproto.NewReader(bufio.NewReader(bb.conn))

	// listens for chat messages
	for {
		line, err := tp.ReadLine()
		if nil != err {

			// officially disconnects the bot from the server
			bb.Disconnect()

			return errors.New("bb.Bot.HandleChat: Failed to read line from channel. Disconnected.")
		}

		// logs the response from the IRC server
		rgb.YPrintf("[%s] %s\n", TimeStamp(), line)

		if "PING :tmi.twitch.tv" == line {

			// respond to PING message with a PONG message, to maintain the connection
			bb.conn.Write([]byte("PONG :tmi.twitch.tv\r\n"))
			continue
		} else {

			// handle a PRIVMSG message
			matches := MsgRegex.FindStringSubmatch(line)
			if nil != matches {
				userName := matches[1]
				msgType := matches[2]
				channel := matches[3]

				switch msgType {
				case "PRIVMSG":
					msg := matches[4]
					rgb.GPrintf("[%s] %s: %s\n", TimeStamp(), userName, msg)

					// parse commands from user message
					cmdMatches := CmdRegex.FindStringSubmatch(msg)
					if nil != cmdMatches {
						cmd := cmdMatches[1]

						rgb.YPrintf("[%s] Checking cmd %s against %s\n", TimeStamp(), cmd, bb.ChannelData[channel].Command)
						switch cmd {
						case bb.ChannelData[channel].Command:
							rgb.CPrintf("[%s] Card asked for by %s on %s' channel!\n", TimeStamp(), userName, channel)

							bb.Say("Your prompt is : "+bb.Prompts[rand.Intn(len(bb.Prompts))], channel)
						}

						// channel-owner specific commands
						if userName == channel {
							switch cmd {
							case "tbdown":
								rgb.CPrintf(
									"[%s] Shutdown command received. Shutting down now...\n",
									TimeStamp(),
								)

								bb.Disconnect()
								return nil
							case "wat":
								magicCode := bb.readOrCreateChannelKey(channel)
								rgb.CPrintf(
									"[%s] Magic code is %s - https://karaokards.ing.martyn.berlin/admin/%s/%s\n",
									TimeStamp(),
									magicCode, userName, magicCode,
								)
								bb.Say("Ack.")
							default:
								// do nothing
							}
						}
					}
				default:
					// do nothing
					rgb.YPrintf("[%s] unknown IRC message : %s\n", TimeStamp(), line)
				}
			}
		}
		time.Sleep(bb.MsgRate)
	}
}

// Login to the IRC server
func (bb *KardBot) Login() {
	rgb.YPrintf("[%s] Logging into #%s...\n", TimeStamp(), bb.Channel)
	bb.conn.Write([]byte("PASS " + bb.Credentials.Password + "\r\n"))
	bb.conn.Write([]byte("NICK " + bb.Name + "\r\n"))
}


func (bb *KardBot) LeaveChannel(channels ...string) {
	for _, channel := range channels {
		rgb.YPrintf("[%s] Leaving #%s...\n", TimeStamp(), channel)
		bb.conn.Write([]byte("PART #" + channel + "\r\n"))
		rgb.YPrintf("[%s] Left #%s as @%s!\n", TimeStamp(), channel, bb.Name)
	}
}

// Makes the bot join its pre-specified channel.
func (bb *KardBot) JoinChannel(channels ...string) {
	if len(channels) == 0 {
		channels = append(channels, bb.Channel)
	}

	for _, channel := range channels {
		rgb.YPrintf("[%s] Joining #%s...\n", TimeStamp(), channel)
		bb.conn.Write([]byte("JOIN #" + channel + "\r\n"))
		rgb.YPrintf("[%s] Joined #%s as @%s!\n", TimeStamp(), channel, bb.Name)
	}
}

// Reads from the private credentials file and stores the data in the bot's Credentials field.
func (bb *KardBot) ReadCredentials() error {

	// reads from the file
	credFile, err := ioutil.ReadFile(bb.PrivatePath)
	if nil != err {
		return err
	}

	bb.Credentials = &OAuthCred{}

	// parses the file contents
	dec := json.NewDecoder(strings.NewReader(string(credFile)))
	if err = dec.Decode(bb.Credentials); nil != err && io.EOF != err {
		return err
	}

	return nil
}

// Makes the bot send a message to the chat channel.
func (bb *KardBot) Say(msg string, channels ...string) error {
	if "" == msg {
		return errors.New("BasicBot.Say: msg was empty.")
	}

	// check if message is too large for IRC
	if len(msg) > 512 {
		return errors.New("BasicBot.Say: msg exceeded 512 bytes")
	}

	if len(channels) == 0 {
		channels = append(channels, bb.Channel)
	}

	rgb.YPrintf("[%s] sending %s to channels %v as @%s!\n", TimeStamp(), msg, channels, bb.Name)
	for _, channel := range channels {
		_, err := bb.conn.Write([]byte(fmt.Sprintf("PRIVMSG #%s :%s\r\n", channel, msg)))
		rgb.YPrintf("[%s] PRIVMSG #%s :%s\r\n", TimeStamp(), channel, msg)
		if nil != err {
			return err
		}
	}
	return nil
}

// Starts a loop where the bot will attempt to connect to the Twitch IRC server, then connect to the
// pre-specified channel, and then handle the chat. It will attempt to reconnect until it is told to
// shut down, or is forcefully shutdown.
func (bb *KardBot) Start() {
	err := bb.ReadCredentials()
	if nil != err {
		fmt.Println(err)
		fmt.Println("Aborting!")
		return
	}
	
	err = bb.readChannelData()
	if nil != err {
		fmt.Println(err)
		fmt.Println("Aborting!")
		return
	}

	for {
		bb.Connect()
		bb.Login()
		if len(bb.ChannelData) > 0 {
			for channelName := range(bb.ChannelData) {
				bb.JoinChannel(channelName)
			}
		} else {
			bb.JoinChannel()
		}
		err = bb.HandleChat()
		if nil != err {

			// attempts to reconnect upon unexpected chat error
			time.Sleep(1000 * time.Millisecond)
			fmt.Println(err)
			fmt.Println("Starting bot again...")
		} else {
			return
		}
	}
}

func (bb *KardBot) readChannelData() error {
	records, err := bb.Database.ReadAll("channelData")
	if err != nil {
		// no db? initialise one?
		record := ChannelData{Name: bb.Channel, JoinTime: time.Now(), Command: "card"}
		rgb.YPrintf("[%s] No channel table for #%s exists, creating...\n", TimeStamp(), bb.Channel)
		if err := bb.Database.Write("channelData", bb.Channel, record); err != nil {
			return err
		}
		bb.ChannelData = make(map[string]ChannelData)
		bb.ChannelData[bb.Channel] = record;
	} else {
		bb.ChannelData = make(map[string]ChannelData)
	}
	for _, data := range records {
		record := ChannelData{}
		err := json.Unmarshal([]byte(data), &record);
		if err != nil {
			return err
		}
		if record.Name != "" {
			if record.Command == "" {
				record.Command = "card"
				
				rgb.YPrintf("[%s] Rewriting data for #%s...\n", TimeStamp(), bb.Channel)
				if err := bb.Database.Write("channelData", record.Name, record); err != nil {
					return err
				}
			}
			bb.ChannelData[record.Name] = record
		}
	}
	// Managed to leave the main channel!?
	if bb.ChannelData[bb.Channel].Name == "" {
		rgb.YPrintf("[%s] No channel data for #%s exists, creating...\n", TimeStamp(), bb.Channel)
		record := ChannelData{Name: bb.Channel, JoinTime: time.Now(), Command: "card"}
		bb.ChannelData[bb.Channel] = record
		if err := bb.Database.Write("channelData", bb.Channel, record); err != nil {
			return err
		}
		records, err = bb.Database.ReadAll("channelData")
	}
	rgb.YPrintf("[%s] Read channel data for %d channels\n", TimeStamp(), len(bb.ChannelData))
	return nil
}

func (bb *KardBot) readOrCreateChannelKey(channel string) string {
	magicCode := ""
	var err error
	var record ChannelData
	if record, ok := bb.ChannelData[channel]; !ok {
		rgb.YPrintf("[%s] No channel data for #%s exists, creating\n", TimeStamp(), channel)
		err = bb.Database.Read("channelData", channel, &record);
		if err == nil {
			bb.ChannelData[channel] = record
		}
	}
	record = bb.ChannelData[channel]
	if err != nil || record.AdminKey == "" {
		rgb.YPrintf("[%s] No channel key for #%s exists, creating one\n", TimeStamp(), channel)
		newuu, _ := uuid.NewRandom()
		magicCode = base64.StdEncoding.EncodeToString([]byte(newuu.String()))
		record.AdminKey = magicCode
		if record.Name == "" {
			record.Name = channel
		}
		if err := bb.Database.Write("channelData", channel, record); err != nil {
			rgb.RPrintf("[%s] Error writing channel data for #%s\n", TimeStamp(), channel)
		}
		bb.ChannelData[record.Name] = record
		rgb.YPrintf("[%s] Cached channel key for #%s\n", TimeStamp(), record.Name)
	} else {
		magicCode = record.AdminKey
		rgb.YPrintf("[%s] Loaded data for #%s\n", TimeStamp(), channel)
	}
	return magicCode
}

func TimeStamp() string {
	return TimeStampFmt(PSTFormat)
}

func TimeStampFmt(format string) string {
	return time.Now().Format(format)
}