irc-client/components/composer.py

from components import IRCSocketClient
from abots.events import CoroEvent
from abots.helpers import Logger, infinitedict, isnumeric, sha256
from abots.helpers import coroutine, generator

from time import sleep
from os import remove as delete_file
from os.path import isfile
from functools import wraps

class Composer:
    def __init__(self, hooks, logger=None):
        self.hooks = hooks
        self.logger = logger
        self._listeners = dict()
        self.kill_switch = CoroEvent()

    def start_socket(self, host, port, timeout=None, secure=False):
        sock = IRCSocketClient(host, port, timeout=timeout, secure=secure)
        sock.start()
        inbox = sock.queues.get("inbox")
        outbox = sock.queues.get("outbox")
        events = sock.queues.get("events")
        return sock, inbox, outbox, events

    def register_hook(self, func):
        hook_name = func.__name__
        if hook_name in list(self.hooks):
            return
        self.hooks[hook_name] = func
        return generator(func)

    def subscribe(self, hook_name):
        key = sha256()
        hook = self.hooks[hook_name]
        if hook is not None:
            event = CoroEvent()
            gen_hook = generator(hook)
            self._listeners[key] = gen_hook(self.unsubscribe(key), event)

    def subscribe_many(self, hook_names):
        for hook_name in hook_names:
            self.subscribe(hook_name)

    def unsubscribe(self, key):
        def cancel():
            try:
                del self._listeners[key]
            except KeyError:
                pass
        return cancel

    @generator
    def send_to_listeners(self, sock):
        line = (yield)
        for key in list(self._listeners):
            listener = self._listeners.get(key, None)
            if listener is None:
                return
            try:
                listener.send((sock, line))
            except StopIteration:
                cancel = self.unsubscribe(key)
                cancel()

    def get_messages(self, sock):
        cache = ""
        listeners = self.send_to_listeners(sock)
        while not self.kill_switch.is_set():
            for recv in sock.recv():
                if recv is None or len(recv) == 0:
                    continue
                if "\n" in recv:
                    lines = recv.split("\n")
                    listeners.send(cache + lines[0].strip("\r"))
                    for line in lines[1:-1]:
                        listeners.send(line.strip("\r"))
                    cache = lines[-1]
                else:
                    cache = cache + recv
        return cache

    # from section 2.3.1 of rfc1459
    def parse_message(self, message):
        meta = dict()
        meta["raw"] = message
        if message is None or len(message) == 0:
            return meta
        has_tags = message[0] == "@"
        meta["has_tags"] = has_tags
        if has_tags:
            tags, message = message.lstrip("@").split(" ", 1)
            meta["tags"] = dict()
            for tag in tags.split(";"):
                key, value = tag.split("=", 1)
                meta["tags"][key] = value
        has_prefix = message[0] == ":"
        message = message.lstrip(":")
        meta["has_prefix"] = has_prefix
        if not has_prefix:
            command, params = message.lstrip(":").split(" ", 1)
            meta["type"] = "alpha"
        else:
            prefix, command, params = message.split(" ", 2)
            source = "nick" if "!" in prefix else "server"
            meta["prefix"] = prefix
            meta["source"] = source
            meta["type"] = "numeric" if isnumeric(command) else "alpha"
            if source == "nick":
                nick, user_info = prefix.split("!", 1)
                user, host = user_info.split("@", 1)
                meta["nick"] = nick
                meta["user"] = user
                meta["host"] = host
        meta["command"] = command
        meta["params"] = params
        return meta

    def cleanup(self, parameters, dirt=":"):
        params = parameters.partition(dirt)
        return params[0] + params[2]