twitchtts/tts.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-

import json
import yaml
import logging
import socket
import sys
import time
import datetime
import socketserver

from threading import Thread, current_thread
from http.server import BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs

class IRC:
    irc = socket.socket()

    def __init__(self):
        self.irc = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.tts_denied = []
        if 'WHITELIST_USER' in conf:
            self.tts_allowed = conf['WHITELIST_USER']
        else:
            self.tts_allowed = []
        self.tts_status = True

    def connect(self, server, port, channel, botnick, botpass):
        logging.info("Connecting to: " + server)
        try:
            self.irc.connect((server, port))
        except ConnectionResetError:
            logging.fatal('Twitch refused to connect, please check your settings and try again (or just try again, as Twitch sometimes refuses to connect for no reason).')
            sys.exit(252)

        self.irc.settimeout(1)

        self.irc.send(bytes("PASS " + botpass + "\r\n", "UTF-8"))
        self.irc.send(bytes("USER " + botnick + " " + botnick +" " + botnick + " :python\r\n", "UTF-8"))
        self.irc.send(bytes("NICK " + botnick + "\r\n", "UTF-8"))
        self.irc.send(bytes("CAP REQ :twitch.tv/commands twitch.tv/tags \r\n", "UTF-8"))
        time.sleep(5)

        try:
            self.irc.send(bytes("JOIN " + channel + "\r\n", "UTF-8"))
        except ConnectionResetError:
            self.irc.send(bytes("JOIN " + channel + "\r\n", "UTF-8"))

    def sendpriv(self, channel, user, msg):
        self.irc.send(bytes("PRIVMSG "+channel+" :"+user+" "+msg+"\r\n", "UTF-8"))

    def get_response(self):
        try:
            resp = self.irc.recv(2048).decode("UTF-8")
            logging.debug('resp:')
            logging.debug(resp)
        except socket.timeout:
            return False
        except:
            sys.exit(255)

        if resp.find('PING') != -1:
            logging.info('PING received')
            self.irc.send(bytes('PONG :tmi.twitch.tv\r\n', "UTF-8"))

        if resp.find('CLEARMSG') != -1:
            logging.info('CLEARMSG received')
            msgid = False

            global msg_queue_raw
            filtered_msg_queue = []

            tags = resp.split(';')
            for tag in tags:
                if "target-msg-id=" in tag:
                    msgid = tag.rsplit('target-msg-id=',1)[1]
                    logging.debug('Trying to suppress message')
                    logging.debug(msgid)

                    for msg in list(msg_queue_raw):
                        if msg['msgid'] == msgid:
                            logging.info('Suppressing message '+str(msgid))
                        else:
                            filtered_msg_queue.append(msg)

                    msg_queue_raw = filtered_msg_queue

            return True

        if resp.find('PRIVMSG') != -1:
            logging.debug('PRIVMSG received')
            badges = False
            subscriber = False
            msgid = False
            msg = False
            msglen = False
            user = False
            tts = False

            tags = resp.split(';')
            for tag in tags:
                if tag.startswith('badges='):
                    badges = tag.rsplit('badges=',1)[1]
                if tag.startswith('subscriber='):
                    subscriber = tag.rsplit('subscriber=',1)[1]
                    logging.debug('Subscriber: '+str(subscriber))
                if tag.startswith('id='):
                    msgid = tag.rsplit('id=',1)[1]
                    logging.debug('Message ID: '+str(msgid))
                if tag.startswith('display-name='):
                    user = tag.rsplit('display-name=',1)[1]
                    logging.debug('Username: '+str(user))

            msg = resp.rsplit('PRIVMSG #',1)[1]
            msg = msg.split(':',1)[1]
            msg = msg.replace('\r\n','')
            msglen = len(msg)

            logging.debug('Msg:')
            logging.debug(msg)
            logging.debug('Msg length: '+str(msglen))

            if 'broadcaster' in badges or 'moderator' in badges:
                if msg.startswith('!dtts'):
                    logging.debug("!dtts command detected")
                    user = msg.replace('!dtts', '').strip()
                    if user not in self.tts_denied:
                        logging.info("Adding "+str(user)+" to deny list")
                        self.tts_denied.append(user)
                    if user in self.tts_allowed:
                        logging.info("Removing "+str(user)+" from allowed list")
                        self.tts_allowed.remove(user)

                if msg.startswith('!ptts'):
                    logging.debug("!ptts command detected")
                    user = msg.replace('!ptts', '').strip()

                    self.tts_allowed.append(user)

                    if user in self.tts_denied:
                        logging.info("Removing "+str(user)+" from deny list")
                        self.tts_denied.remove(user)

                logging.debug('Deny List:')
                logging.debug(self.tts_denied)

                if msg.startswith('!toff'):
                    logging.info('TTS is now turned off')
                    msg_queue.clear()
                    self.tts_status = False
                if msg.startswith('!ton'):
                    logging.info('TTS is now turned on')
                    msg_queue.clear()
                    self.tts_status = True

            if msg.startswith('!tts'):
                logging.debug('!tts command detected')

                if msglen > conf['IRC_TTS_LEN']:
                    logging.info('TTS message is to long')
                    self.sendpriv(conf['IRC_CHANNEL'], "@"+str(user), conf['MESSAGE']['TOO_LONG'])
                    return False

                if not self.tts_status:
                    logging.info('TTS is disabled')
                    self.sendpriv(conf['IRC_CHANNEL'], "@"+str(user), conf['MESSAGE']['DISABLED'])
                    return False

                if user in self.tts_denied:
                    logging.info(str(user) + " is not allowed to use TTS")
                    self.sendpriv(conf['IRC_CHANNEL'], "@"+str(user), conf['MESSAGE']['DENIED'])
                    return False

                if conf['IRC_SUBONLY']:
                    if subscriber != "0" or 'moderator' in badges or 'broadcaster' in badges:
                        logging.debug('TTS is sub-only and user has allowance')
                    else:
                        logging.info('TTS is sub-only')
                        self.sendpriv(conf['IRC_CHANNEL'], "@"+str(user), conf['MESSAGE']['SUBONLY'])
                        return False

                if conf['IRC_MODONLY']:
                    if 'moderator' in badges or 'broadcaster' in badges:
                        logging.debug('TTS is mod-only and user has allowance')
                    else:
                        logging.info('TTS is sub-only')
                        self.sendpriv(conf['IRC_CHANNEL'], "@"+str(user), conf['MESSAGE']['MODONLY'])
                        return False

                if conf['WHITELIST']:
                    if not user in self.tts_allowed:
                        logging.info('User is not on whitelist')
                        logging.info(self.tts_allowed)
                        self.sendpriv(conf['IRC_CHANNEL'], "@"+str(user), conf['MESSAGE']['WHITELISTONLY'])
                        return False

                logging.info('Valid TTS message, adding to raw queue')
                tts = True
                now = datetime.datetime.now()
                msg = msg.replace('!tts','',1)
                msg = {"TTS": tts, "msg": msg, "badges": badges, "subscriber": subscriber, "msgid": msgid, "user": user, "length": msglen, "queuetime": now, "timestamp": str(time.time_ns())}
                msg_queue_raw.append(msg)

            return resp

class HTTPserv(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        return

    def do_GET(self):
        if self.path == '/':
            self.send_response(200)
            self.send_header('Content-type', 'text/html')
            self.end_headers()
            with open("tts.html", "rb") as fh:
                html = fh.read()
                self.wfile.write(html)

        if self.path == '/tts.js':
            self.send_response(200)
            self.send_header('Content-type', 'text/javascript')
            self.end_headers()
            with open("tts.js", "rb") as fh:
                html = fh.read()
                self.wfile.write(html)

        if self.path == '/jquery.js':
            self.send_response(200)
            self.send_header('Content-type', 'text/javascript')
            self.end_headers()
            with open("jquery.js", "rb") as fh:
                html = fh.read()
                self.wfile.write(html)

        if self.path == '/bootstrap.min.css':
            self.send_response(200)
            self.send_header('Content-type', 'text/css')
            self.end_headers()
            with open("bootstrap.min.css", "rb") as fh:
                html = fh.read()
                self.wfile.write(html)

        if self.path.startswith('/tts_queue'):
            tts_json = ""
            tts = {}
            self.send_response(200)
            self.send_header('Content-type', 'text/json')
            self.end_headers()

            usermap = conf['USERMAP']
            sorted_tts = {k: msg_queue[k] for k in sorted(msg_queue, reverse=False)}

            for key in list(sorted_tts.keys()):
                if key not in tts_done:
                    if msg_queue[key][0].lower() in usermap:
                        tts = {key: usermap[msg_queue[key][0].lower()]+" "+str(conf['MESSAGE']['SAYS'])+":"+msg_queue[key][1]}
                    else:
                        tts = {key: msg_queue[key][0]+" "+str(conf['MESSAGE']['SAYS'])+":"+msg_queue[key][1]}

            tts_json = json.dumps(tts)
            self.wfile.write(bytes(str(tts_json)+"\n", "utf-8"))

        if self.path.startswith('/tts_done'):
            get_params = parse_qs(self.path)
            if '/tts_done?id' in get_params:
                logging.info("Removing message from queue")
                tts_done.append(get_params['/tts_done?id'][0])
                self.send_response(200)
                self.send_header('Content-type', 'text/plain')
                self.end_headers()
                self.wfile.write(bytes("OK\n", "utf-8"))
            else:
                self.send_response(500)
                self.send_header('Content-type', 'text/plain')
                self.end_headers()
                self.wfile.write(bytes("Internal Server error\n", "utf-8"))

        return

def http_serve_forever(httpd):
  httpd.serve_forever()

def load_config():
    logging.info("Loading configfile")

    try:
        with open("config.yml", "r") as ymlfile:
            cfg = yaml.load(ymlfile, Loader=yaml.Loader)
    except FileNotFoundError:
        logging.fatal('Your config file is missing, please copy config-dist.yml to config.yml and review your settings.')
        sys.exit(253)

    for section in cfg:
        try:
            logging.debug('Fetching config: '+str(section))
            conf['IRC_CHANNEL'] = cfg['irc']['channel']
            conf['IRC_USERNAME'] = cfg['irc']['username']
            conf['IRC_OAUTH_TOKEN'] = cfg['irc']['oauth_token']
            conf['IRC_SERVER'] = cfg['irc']['server']
            conf['IRC_CLEARMSG_TIMEOUT'] = cfg['irc']['clearmsg_timeout']
            
            conf['IRC_SUBONLY'] = cfg['bot']['subonly']
            conf['IRC_MODONLY'] = cfg['bot']['modonly']
            conf['IRC_TTS_LEN'] = cfg['bot']['message_length']
            
            conf['LOG_LEVEL'] = cfg['log']['level']
            conf['HTTP_PORT'] = cfg['http']['port']
            conf['HTTP_BIND'] = cfg['http']['bind']

            conf['MESSAGE'] = {}
            conf['MESSAGE']['TOO_LONG'] = cfg['messages']['too_long']
            conf['MESSAGE']['DISABLED'] = cfg['messages']['disabled']
            conf['MESSAGE']['DENIED'] = cfg['messages']['denied']
            conf['MESSAGE']['SUBONLY'] = cfg['messages']['subonly']
            conf['MESSAGE']['MODONLY'] = cfg['messages']['modonly']
            conf['MESSAGE']['READY'] = cfg['messages']['ready']
            conf['MESSAGE']['WHITELISTONLY'] = cfg['messages']['whitelist']
            conf['MESSAGE']['SAYS'] = cfg['messages']['says']

            conf['USERMAP'] = cfg['usermapping']

            if 'whitelist' in cfg:
                conf['WHITELIST'] = True
                conf['WHITELIST_USER'] = cfg['whitelist']
            else:
                conf['WHITELIST'] = False

        except KeyError as e:
            logging.exception('Your config file is invalid, please check and try again.')
            sys.exit(254)

    if conf['WHITELIST']:
        logging.info('Whitelist mode enabled')
        logging.debug('Whitelist:')
        logging.debug(conf['WHITELIST_USER'])

    return conf

conf = {}
tts_done = []
msg_queue_raw = []
msg_queue = {}

logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(module)s %(threadName)s %(levelname)s: %(message)s')

def main():
    conf = load_config()
    lastreload = datetime.datetime.now()
    logging.getLogger().setLevel(conf['LOG_LEVEL'])

    logging.info("Starting Webserver")
    httpd = socketserver.TCPServer((conf['HTTP_BIND'], conf['HTTP_PORT']), HTTPserv)
    httpd.allow_reuse_port = True
    httpd.allow_reuse_address = True

    http_thread = Thread(target=http_serve_forever, daemon=True, args=(httpd, ))
    http_thread.start()

    logging.info("Starting IRC bot")
    irc = IRC()
    irc.connect(conf['IRC_SERVER'], 6667, conf['IRC_CHANNEL'], conf['IRC_USERNAME'], conf['IRC_OAUTH_TOKEN'])
    irc.sendpriv(conf['IRC_CHANNEL'], 'MrDestructoid', conf['MESSAGE']['READY'])

    logging.info("Please open your browser and visit: http://"+str(conf['HTTP_BIND']+":"+str(conf['HTTP_PORT'])+"/"))

    while True:
        if conf['LOG_LEVEL'] == "DEBUG":
            time.sleep(1)

        try:
            irc.get_response()

            if not irc.tts_status:
                logging.debug("TTS is disabled")
                if conf['LOG_LEVEL'] == "DEBUG":
                    time.sleep(1)
                continue

            confreload = datetime.datetime.now()
            if confreload - lastreload > datetime.timedelta(seconds=60):
                conf = load_config()
                lastreload = datetime.datetime.now()

            logging.debug('Raw message queue:')
            logging.debug(msg_queue_raw)

            for raw_msg in msg_queue_raw:
                logging.debug('Raw msg:')
                logging.debug(msg_queue_raw)

                now = datetime.datetime.now()
                if now - raw_msg['queuetime'] > datetime.timedelta(seconds=conf['IRC_CLEARMSG_TIMEOUT']):
                    logging.debug('clearmsg_timeout reached')
                    if raw_msg['timestamp'] not in msg_queue:
                        logging.info('Sending TTS message')
                        msg_queue[raw_msg['timestamp']] = [raw_msg['user'], raw_msg['msg']]
                        logging.debug(msg_queue)
                    else:
                        logging.debug('Msg is already in queue')
        except KeyboardInterrupt:
            logging.info('Exiting...')
            sys.exit()

if __name__ == "__main__":
    main()