Edm0nd/edmond/bot.py

import asyncio
import importlib
import json
import os
import signal
import sys
import time
import traceback
from pathlib import Path
from typing import Any, Iterable, Optional

import irc.client
import irc.client_aio
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from irc.client import Connection, Event, NickMask

from edmond.log import Logger
from edmond.plugin import Plugin


class Bot(irc.client_aio.AioSimpleIRCClient, Logger):
    """Main class for the IRC bot: handles connection and manages plugins."""

    CHANNELS_RUNTIME_KEY = "_channels"

    def __init__(self, config: dict, logger):
        super().__init__()
        self.config: dict = config
        self.logger = logger
        self.plugins: list[Plugin] = []
        self.values: dict[str, Any] = {}
        self.storage: dict[str, Any] = self.get_storage()
        self.tasks: list[asyncio.Task] = []
        self.done: bool = False
        self.scheduler = AsyncIOScheduler()
        self.scheduler.start()

    @property
    def nick(self) -> str:
        """Nickname validated by the server, or the configured nick."""
        if self.connection.is_connected():
            return self.connection.get_nickname()
        return self.config["nick"]

    @property
    def names(self) -> Iterable[str]:
        """Collection of names the bot should identify with."""
        return (self.nick, *self.config["alternative_nicks"])

    @property
    def channels(self) -> list[str]:
        """List of joined channels."""
        if self.CHANNELS_RUNTIME_KEY not in self.values:
            self.values[self.CHANNELS_RUNTIME_KEY] = []
        return self.values[self.CHANNELS_RUNTIME_KEY]

    def get_storage(self) -> dict:
        """Load data from storage."""
        try:
            with open(self.config["storage_file"], "rt") as storage_file:
                storage = json.load(storage_file)
                self.log_d("Loaded storage file.")
                return storage
        except (OSError, json.decoder.JSONDecodeError) as exc:
            self.log_e(f"Could not load storage file: {exc}")
            self.log_w(
                "If it's not the first time Edm0nd is run, you may lose"
                " data when closing the program."
            )
            return {}

    def save_storage(self) -> None:
        """Save storage data to disk."""
        try:
            with open(self.config["storage_file"], "wt") as storage_file:
                json.dump(self.storage, storage_file, indent=2, sort_keys=True)
                self.log_d("Saved storage file.")
        except OSError as exc:
            self.log_e(f"Could not save storage file: {exc}")

    def handle_task(self, coro):
        """Schedule a task in the event loop. Keep a reference to cancel it."""
        task = self.connection.reactor.loop.create_task(coro)
        self.tasks.append(task)

    def on_welcome(self, connection: Connection, event: Event):
        """Handle a successful connection to a server."""
        self.log_i(f"Connected to server {event.source}.")
        self.run_plugin_callbacks(event)
        for channel in self.config["channels"]:
            connection.join(channel)

    def on_join(self, connection: Connection, event: Event):
        """Handle someone, possibly the bot, joining a channel."""
        if event.source.nick == self.nick:
            self.log_i(f"Joined {event.target}.")
            self.channels.append(event.target)
        self.run_plugin_callbacks(event)

    def on_part(self, connection: Connection, event: Event):
        """Handle someone, possibly the bot, leaving a channel."""
        if event.source.nick == self.nick:
            self.log_i(f"Left {event.target} (args: {event.arguments[0]}).")
            self.channels.remove(event.target)
        self.run_plugin_callbacks(event)

    def on_pubmsg(self, connection: Connection, event: Event):
        """Handle a message received in a channel."""
        channel = event.target
        nick = NickMask(event.source).nick
        message = event.arguments[0]
        self.log_d(f"Message in {channel} from {nick}: {message}")
        self.run_plugin_callbacks(event)

    def on_privmsg(self, connection: Connection, event: Event):
        """Handle a message received privately, usually like a channel msg."""
        nick = NickMask(event.source).nick
        target = event.target
        message = event.arguments[0]
        self.log_d(f"Private message from {nick} to {target}: {message}")
        self.run_plugin_callbacks(event)

    def on_ping(self, connection: Connection, event: Event):
        """Handle a ping; can be used as a random event timer."""
        self.log_d(f"Received ping from {event.target}.")
        self.run_plugin_callbacks(event)

    def run(self):
        """Connect the bot to server, join channels and start responding."""
        self.log_i("Starting Edmond.")
        self.load_plugins()
        self.log_i("Connecting to server…")

        signal.signal(signal.SIGTERM, self.handle_sigterm)
        try:
            self.connect(self.config["host"], self.config["port"], self.nick)
            self.start()
        except irc.client.ServerConnectionError as exc:
            self.log_c(f"Connection failed: {exc}")
        except KeyboardInterrupt:
            self.log_i("Caught keyboard interrupt.")
        except Exception as exc:
            self.log_c(f"Caught unhandled {type(exc).__name__}: {exc}")
            _, _, exc_traceback = sys.exc_info()
            for line in traceback.format_tb(exc_traceback):
                self.log_d(line.rstrip())
        finally:
            self.cleanup()

    def load_plugins(self):
        """Load all installed plugins."""
        self.log_i("Loading plugins…")
        plugin_files = os.listdir(Path(__file__).parent / "plugins")
        plugin_names = map(
            lambda f: os.path.splitext(f)[0],
            filter(
                lambda f: f.endswith(".py") and f != "__init__.py",
                plugin_files,
            ),
        )
        for plugin_name in plugin_names:
            module = importlib.import_module(f"edmond.plugins.{plugin_name}")
            # Get plugin class name from its module name.
            class_name = (
                "".join(map(lambda w: w.capitalize(), plugin_name.split("_")))
                + "Plugin"
            )
            plugin_class = getattr(module, class_name)
            self.plugins.append(plugin_class(self))
            self.values[plugin_name] = {}
            self.log_d(f"Loaded {class_name}.")

    def get_plugin(self, name: str) -> Optional[Plugin]:
        """Get a loaded plugin by its name (e.g. 'mood'), or None."""
        matching_plugins = filter(
            lambda plugin: plugin.name == name,
            self.plugins,
        )
        return next(matching_plugins, None)

    def say(self, target: str, message: str) -> None:
        """Send message to target after a slight delay."""
        message = message.replace("\n", " ").replace("\r", " ")
        time.sleep(self.config["speak_delay"])
        self.log_d(f"Sending to {target}: {message}")
        try:
            if message.startswith("/me "):
                self.connection.action(target, message[4:])
            else:
                self.connection.privmsg(target, message)
        except irc.client.MessageTooLong:
            self.log_e("Could not send, message is too long.")

    def run_plugin_callbacks(self, event: Event) -> None:
        """Run appropriate callbacks for each plugin."""
        etype = event.type
        ready_plugins = filter(lambda p: p.is_ready, self.plugins)
        plugins = sorted(ready_plugins, key=lambda p: p.priority, reverse=True)
        for plugin in plugins:
            callbacks = plugin.callbacks
            if etype not in callbacks:
                continue
            if callbacks[etype](event):
                break

    def handle_sigterm(self, *args):
        """Handle SIGTERM (keyboard interrupt, systemd stop, etc)."""
        self.cleanup()
        exit("Exiting after received SIGTERM.")

    def cleanup(self) -> None:
        """Save the storage file and close the connection. Run only once."""
        if self.done:
            return
        self.log_i("Stopping Edmond.")
        self.save_storage()  # FIRST THINGS FIRST

        for task in self.tasks:
            if not task.cancelled():
                self.log_d(f"Cancelling task {task.get_name()}…")
                task.cancel()
        if self.connection.is_connected():
            self.connection.close()
        self.reactor.loop.close()
        self.done = True