omnibot22/omnibot/bot.py

import asyncio
import logging
from typing import Sequence, Set

from asyncirc.protocol import IrcProtocol
from asyncirc.server import Server
from irclib.parser import Message

from .config import ServerConfig
from . import plugin


log = logging.getLogger(__name__)
ERRORS = {
    "401": "ERR_NOSUCHNICK",
    "402": "ERR_NOSUCHSERVER",
    "403": "ERR_NOSUCHCHANNEL",
    "404": "ERR_CANNOTSENDTOCHAN",
    "405": "ERR_TOOMANYCHANNELS",
    "406": "ERR_WASNOSUCHNICK",
    "407": "ERR_TOOMANYTARGETS",
    "408": "ERR_NOSUCHSERVICE",
    "409": "ERR_NOORIGIN",
    "411": "ERR_NORECIPIENT",
    "412": "ERR_NOTEXTTOSEND",
    "413": "ERR_NOTOPLEVEL",
    "414": "ERR_WILDTOPLEVEL",
    "415": "ERR_BADMASK",
    "421": "ERR_UNKNOWNCOMMAND",
    "422": "ERR_NOMOTD",
    "423": "ERR_NOADMININFO",
    "424": "ERR_FILEERROR",
    "431": "ERR_NONICKNAMEGIVEN",
    "432": "ERR_ERRONEUSNICKNAME",
    "433": "ERR_NICKNAMEINUSE",
    "436": "ERR_NICKCOLLISION",
    "437": "ERR_UNAVAILRESOURCE",
    "441": "ERR_USERNOTINCHANNEL",
    "442": "ERR_NOTONCHANNEL",
    "443": "ERR_USERONCHANNEL",
    "444": "ERR_NOLOGIN",
    "445": "ERR_SUMMONDISABLED",
    "446": "ERR_USERSDISABLED",
    "451": "ERR_NOTREGISTERED",
    "461": "ERR_NEEDMOREPARAMS",
    "462": "ERR_ALREADYREGISTRED",
    "463": "ERR_NOPERMFORHOST",
    "464": "ERR_PASSWDMISMATCH",
    "465": "ERR_YOUREBANNEDCREEP",
    "466": "ERR_YOUWILLBEBANNED",
    "467": "ERR_KEYSET",
    "471": "ERR_CHANNELISFULL",
    "472": "ERR_UNKNOWNMODE",
    "473": "ERR_INVITEONLYCHAN",
    "474": "ERR_BANNEDFROMCHAN",
    "475": "ERR_BADCHANNELKEY",
    "476": "ERR_BADCHANMASK",
    "477": "ERR_NOCHANMODES",
    "478": "ERR_BANLISTFULL",
    "481": "ERR_NOPRIVILEGES",
    "482": "ERR_CHANOPRIVSNEEDED",
    "483": "ERR_CANTKILLSERVER",
    "484": "ERR_RESTRICTED",
    "485": "ERR_UNIQOPPRIVSNEEDED",
    "491": "ERR_NOOPERHOST",
    "501": "ERR_UMODEUNKNOWNFLAG",
    "502": "ERR_USERSDONTMATCH",
}


class Bot:
    def __init__(self, server_config: ServerConfig):
        self.__server_config = server_config
        self.__channels: Set[str] = set()
        self.__quitting = asyncio.Event()
        self.__plugins = [
            plugin.load_plugin(self, config)
            for config in server_config.plugins
            if config.get("enabled", True)
        ]

    @property
    def server_config(self) -> ServerConfig:
        return self.__server_config

    @property
    def plugins(self) -> Sequence[plugin.Plugin]:
        return self.__plugins

    @property
    def joined_channels(self) -> Set[str]:
        """
        Returns a list of all channels that this bot has joined.
        """
        return self.__channels

    def quit(self):
        self.__quitting.set()

    def channel_plugins(self, channel: str) -> Sequence[plugin.Plugin]:
        return [plugin for plugin in self.plugins if channel in plugin.channels]

    async def run(self):
        loop = asyncio.get_running_loop()
        server = Server(
            self.server_config.server,
            self.server_config.port,
            self.server_config.use_ssl,
        )
        log.info("Initializing plugins")
        await asyncio.gather(*[plugin.on_load() for plugin in self.plugins])
        self.connection = IrcProtocol([server], self.server_config.nick, loop=loop)
        # Register events
        # self.connection.register("*", self.on_message)
        self.connection.register("001", self.on_connect)
        self.connection.register("JOIN", self.on_join)
        self.connection.register("PART", self.on_part)
        self.connection.register("KICK", self.on_kick)
        self.connection.register("*", self.on_message)
        # Add errors
        for code, _message in ERRORS.items():
            self.connection.register(code, self.on_error)

        # Connect
        log.info("Connecting to %s", self.server_config.server)
        await self.connection.connect()
        # Keepalive loop
        await self.keepalive()

    async def on_connect(self, conn: IrcProtocol, message: Message):
        # Join rooms
        for ch in self.server_config.all_channels:
            msg = Message(None, None, "JOIN", ch)
            conn.send(str(msg))
        # on_connect event on all plugins
        await asyncio.gather(*[plugin.on_connect(conn) for plugin in self.plugins])

    async def on_join(self, conn: IrcProtocol, message: Message):
        log.debug("%s", message)
        channel = message.parameters[0]
        who = message.prefix
        if who.nick == self.server_config.nick:
            self.__channels |= {channel}
            if channel not in self.server_config.all_channels:
                # Try to leave this channel that we were forced to join like some kind of dog
                msg = Message(None, None, "PART", channel)
                conn.send(str(msg))

        # Pass the message along to available plugins
        plugins = self.channel_plugins(channel)
        await asyncio.gather(
            *[plugin.on_join(conn, channel, who) for plugin in plugins]
        )

    async def __on_part(self, conn: IrcProtocol, message: Message):
        "This is the common logic between on_part and on_kick. Don't call this."
        channel = message.parameters[0]
        who = message.prefix
        if who.nick == self.server_config.nick:
            self.__channels -= {channel}
            if channel not in self.server_config.all_channels:
                # Try to rejoin this channel that we were force-parted from
                msg = Message(None, None, "JOIN", channel)
                conn.send(str(msg))

    async def on_part(self, conn: IrcProtocol, message: Message):
        log.debug("%s", message)
        await self.__on_part(conn, message)
        # Pass the message along to available plugins
        channel = message.parameters[0]
        who = message.prefix
        plugins = self.channel_plugins(channel)
        await asyncio.gather(
            *[plugin.on_part(conn, channel, who) for plugin in plugins]
        )

    async def on_kick(self, conn: IrcProtocol, message: Message):
        log.debug("%s", message)
        await self.__on_part(conn, message)
        # Pass the message along to available plugins
        channel = message.parameters[0]
        who = message.prefix
        plugins = self.channel_plugins(channel)
        await asyncio.gather(
            *[plugin.on_kick(conn, channel, who) for plugin in plugins]
        )

    async def on_message(self, conn: IrcProtocol, message: Message):
        # Pass the message to the plugins
        log.trace("%s", message)
        channel = message.parameters[0]
        who = message.prefix
        if who.nick == self.server_config.nick:
            # Don't raise on_message events for ourselves
            return
        line = message.parameters[1]

        # TL;DR OF THE BELOW: if the first parameter looks like a channel in
        # addition to message type, then filter by channel. Otherwise, don't
        # filter by channel.
        #
        # Here's the issue: plugins are *usually* multiplexed by channel. But
        # that's only for messages that target channels, such as PRIVMSG and
        # JOIN. For non-channel messages, such as server status messages (such
        # as 001 on connect, or 372 for MOTD, etc) we want to ignore the channel
        # aspect of plugin multiplexing.
        # In order to accomplish this, we just check if the first parameter
        # looks like a channel - i.e., starts with an octothorpe #.
        if channel and channel[0] == "#":
            plugin_pool = self.channel_plugins(channel)
        else:
            plugin_pool = self.plugins

        # Filter plugins by get_message_types()
        plugins = [
            plugin
            for plugin in plugin_pool
            if message.command in plugin.get_message_types()
        ]
        if plugins:
            await asyncio.gather(
                *[plugin.on_message(conn, channel, who, line) for plugin in plugins]
            )

    async def on_error(self, _conn: IrcProtocol, message: Message):
        log.error("%s", message)

    async def keepalive(self):
        await self.__quitting.wait()
        log.info("Shutting down gracefully")
        await asyncio.gather(
            *[plugin.on_unload(self.connection) for plugin in self.plugins]
        )