weechat-rocketchat/rocketchat.py

#!/usr/bin/env python3

try:
    import weechat
except:
    print("This script must be run under WeeChat.")
    print("Get WeeChat now at: http://www.weechat.org/")
    import sys
    sys.exit(1)

import asyncio
import websockets
import ssl

# globals
rocketchat_servers = []
rocketchat_config_file = None
rocketchat_config_section = {}

class RocketchatServer:
    def __init__(self, name, hostname):
        self.name = name
        self.buffer = ""
        self.client = None
        self._counter = 0

        options = dict(
            hostname=hostname,
            autoconnect="on",
            autoreconnect="on",
        )
        self.options = dict()
        for key, value in options.items():
            self.options[key] = weechat.config_new_option(
                rocketchat_config_file,
                rocketchat_config_section["server"],
                f"{self.name}.{key}",
                "string", "", "", 0, 99999999,
                "", value, 0, "", self.name, "", "", "", "",
            )

    async def connect(self):
        if not self.buffer:
            bufname = f"rc:{self.name}"
            self.buffer = weechat.buffer_search("python", bufname)
            if not self.buffer:
                self.buffer = weechat.buffer_new(bufname, "", "", "", "")
        self.disconnect()

        # connect to the server
        hostname = weechat.config_string(self.options["hostname"])
        uri = f"wss://{hostname}/websocket"
        weechat.prnt("", f"connecting to: {uri}")
        self.client = await websockets.connect(uri, ssl=True)
        weechat.prnt(self.buffer, f"rocketchat: connected to {self.name}")

        # run the listener in a loop
        asyncio.create_task(self.main_loop())

    async def main_loop(self):
        weechat.prnt(self.buffer, "starting to receive!")
        while True:
            g = await self.client.recv()
            weechat.prnt(self.buffer, f"received '{g}'")

    def disconnect(self):
        # TODO:
        pass

    def is_connected(self):
        if self.client is not None and self.client.open:
            return True
        return False

def init():
    weechat.hook_command(
        "rocketchat",
        "Rocketchat",
        "list || server || connect",
        "bruh",
        "list %(weechat_servers)"
        " || server add|del"
        " || connect %(weechat_servers)",
        "rocketchat_main_command", ""
    )
    weechat.hook_completion("rocketchat_servers", "list of rocketchat servers",
            "rocketchat_complete_servers", "")

    global rocketchat_config_file, rocketchat_config_section
    rocketchat_config_file = weechat.config_new("rocketchat", "rocketchat_reload_config", "")
    if not rocketchat_config_file:
        return
    rocketchat_config_section["server"] = weechat.config_new_section(
        rocketchat_config_file, "server", 0, 0,
        "rocketchat_config_servers_read", "",
        "rocketchat_config_servers_write", "",
        "", "", "", "", "", "",
    )

def rocketchat_main_command(data, buffer, args):
    global rocketchat_servers
    if args == "" or args == "list":
        pass # TODO: list
        return weechat.WEECHAT_RC_OK
    argv = args.split(" ")
    if argv[0] == "server":
        if argv[1] == "add":
            name = argv[2]
            hostname = argv[3]
            server = RocketchatServer(name, hostname)
            rocketchat_servers.append(server)
            weechat.prnt("", f"server '{name}' added!")
    elif argv[0] == "connect":
        for name in argv[1:]:
            server = rocketchat_search_server_by_name(name)
            if server:
                asyncio.get_event_loop().run_until_complete(server.connect())
            else:
                weechat.prnt("", f"could not find server name '{name}'")
    return weechat.WEECHAT_RC_OK

def rocketchat_reload_config(data, config_file):
    return weechat.config_reload(config_file)

def rocketchat_config_servers_read(data, config_file, section, option_name, value):
    global rocketchat_servers
    rc = weechat.WEECHAT_CONFIG_OPTION_SET_ERROR
    items = option_name.split(".", 1)
    if len(items) == 2:
        server = rocketchat_search_server_by_name(items[0])
        if not server:
            server = RocketchatServer(items[0])
            rocketchat_servers.append(server)
        if server:
            rc = weechat.config_option_set(server.options[items[1]], value, 1)
    return rc

def rocketchat_config_servers_write(data, config_file, section_name):
    global rocketchat_servers
    weechat.config_write_line(config_file, section_name, "")
    for server in rocketchat_servers:
        for name, option in sorted(server.options.items()):
            weechat.config_write_option(config_file, option)
    return weechat.WEECHAT_RC_OK

def rocketchat_save_config():
    global rocketchat_config_file
    return weechat.config_write(rocketchat_config_file)

def rocketchat_search_server_by_name(name):
    global rocketchat_servers
    for server in rocketchat_servers:
        if server.name == name:
            return server
    return None

def rocketchat_complete_servers(data, completion_item, buffer, completion):
    global rocketchat_servers
    for server in rocketchat_servers:
        weechat.hook_completion_list_add(completion, server.name,
                                         0, weechat.WEECHAT_LIST_POS_SORT)
    return weechat.WEECHAT_RC_OK

def rocketchat_unload():
    global rocketchat_servers
    rocketchat_save_config()
    return weechat.WEECHAT_RC_OK

if weechat.register(
        "rocketchat",
        "Michael Zhang <mail@mzhang.io>",
        "0.1.0",
        "MIT",
        "Rocketchat",
        "rocketchat_unload",
        "utf-8",
    ):
    init()