economybot.py

import re
import json
import socket
from dataclasses import dataclass, field, asdict
from collections import UserString
from typing import Set

TRUSTED = frozenset(("beginbotbot", "isidentical"))
COMMAND_RE = re.compile(r"(?P<command>![a-z]*)+")
ECONOMY_RE = re.compile(r"Mana: (?P<mana>\d+) \| Street Cred: (?P<credit>\d+) \| Cool Points: (?P<points>\d+)")

@dataclass
class Message(UserString):
    data: str
    room: str
    author: str
    
    def __post_init__(self):
        self.message = self.data

    def infer_command(self):
        args = self.data.split()
        if len(args) > 1:
            return args[1]
        else:
            return None

@dataclass
class Stats:
    mana: int
    credit: int
    points: int
    effects: Set[str] = field(default_factory=set)
 
class Client:
    handlers = []

    def __init__(self, nick, password, *, prefix = "@", host = "irc.twitch.tv", port = 6667):
        self.connection = socket.socket()
        self.connection.connect((host, port))
        
        self.push_cmd("pass", password)
        self.push_cmd("nick", nick)
        
        self.nick = nick
        self.prefix = prefix

        self.inital = True
        self.economy = None
        self._economy_needs_me = True
        
    @classmethod
    def from_conf(cls, config_file):
        with open(config_file) as file:
            return cls(**json.load(file))

    @classmethod
    def register(cls, *commands):
        def wrapper(func):
            func.commands = commands
            cls.handlers.append(func)
            return func
        return wrapper
            
    def poll(self):
        buffer = str()
        while self._economy_needs_me:
            buffer += self.connection.recv(1024).decode("UTF-8")
            *lines, buffer = re.split(r"[~\r\n]+", buffer)
            for line in lines:
                print(line)
                self.process_single_line(line.strip())
                
    
    def process_single_line(self, line):
        args = line.split()
        if args[0] == "PING":
            self.push_cmd("PONG", line[1])
        elif len(args) > 1 and args[1] == "PRIVMSG":
            author = args.pop(0).split("!")[0][1:]
            room, *message = args[1:]
            message = " ".join(message)[1:]
            self.dispatch_message(Message(message, room, author))
        else:
            ...

    def dispatch_message(self, message):
        if message.author in TRUSTED and message.startswith(f"@{self.nick} - "):
            self._try_parse_economy(message)

        if self.inital:
            # define inital commands
            self.send_message(message.room, "!me")
            self.inital = False

        if message.startswith(self.prefix + self.nick):
            command = message.infer_command()
            for handler in self.handlers:
                if command in handler.commands:
                    handler(self, message)

    def push_cmd(self, cmd, value):
        request = f"{cmd.upper()} {value}\r\n"
        self.connection.send(request.encode("utf8"))

    def send_message(self, room, message):
        self.push_cmd("privmsg", f"{room} :{message}")

    def _try_parse_economy(self, message):
        if match := ECONOMY_RE.search(str(message)):
            self.economy = Stats(**match.groupdict())

        if self.economy is not None:
            for sound in COMMAND_RE.finditer(str(message)):
                self.economy.effects.add(sound["command"])

@Client.register("help")
def on_help(self, message):
    self.send_message(
        message.room, 
        f"My commands always startswith with "
        f"'{self.prefix}{self.nick}' prefix. "
        f"Available commands: 'help', 'stats'"
    )

@Client.register("stats")
def on_stats(self, message):
    if self.economy is None:
        self.send_message("!me")

    self.send_message(
        message.room,
        ", ".join(f"{key}: {value}" for key, value in asdict(self.economy).items())
    )

if __name__ == "__main__":
    client = Client.from_conf("../configs/stallmansbot.json")
    for channel in ('beginbot', 'isidentical'):
        client.push_cmd("join", f"#{channel}")

    client.poll()