From b983d5904b4590dbb55a499aba94320d2febce61 Mon Sep 17 00:00:00 2001 From: Will Date: Fri, 8 Jun 2018 20:31:38 -0400 Subject: [PATCH] [V3 RPC] Swap back to initial RPC library and hook into core commands (#1780) * Switch RPC libs for websockets support * Implement RPC handling for core * Black reformat * Fix docs for build on travis * Modify RPC to use a Cog base class * Refactor rpc server reference as global * Handle cogbase unload method * Add an init call to handle mutable base attributes * Move RPC server reference back to the bot object * Remove unused import * Add tests for rpc method add/removal * Add tests for rpc method add/removal and cog base unloading * Add one more test * Black reformat * Add RPC mixin...fix MRO * Correct internal rpc method names * Add rpc test html file for debugging/example purposes * Add documentation * Add get_method_info * Update docs with an example RPC call specifying parameter formatting * Make rpc methods UPPER * Black reformat * Fix doc example * Modify this to match new method naming convention * Add more tests --- Pipfile.lock | 34 ++---- docs/framework_bot.rst | 3 + docs/framework_rpc.rst | 62 +++++++--- redbot/__main__.py | 9 +- redbot/core/bot.py | 24 ++-- redbot/core/core_commands.py | 7 ++ redbot/core/events.py | 6 +- redbot/core/rpc.py | 221 ++++++++++++++++++++--------------- requirements.txt | 2 +- tests/conftest.py | 2 +- tests/core/test_rpc.py | 145 +++++++++++++++++++++++ tests/rpc_test.html | 21 ++++ 12 files changed, 379 insertions(+), 157 deletions(-) create mode 100644 tests/core/test_rpc.py create mode 100644 tests/rpc_test.html diff --git a/Pipfile.lock b/Pipfile.lock index f8959d31e..1c92e188d 100644 --- a/Pipfile.lock +++ b/Pipfile.lock @@ -32,6 +32,13 @@ ], "version": "==2.2.5" }, + "aiohttp-json-rpc": { + "hashes": [ + "sha256:9ec69ea70ce49c4af445f0ac56ac728708ccfad8b214272d2cc7e75bc0b31327", + "sha256:e2b8b49779d5d9b811f3a94e98092b1fa14af6d9adbf71c3afa6b20c641fa5d5" + ], + "version": "==0.8.7" + }, "appdirs": { "hashes": [ "sha256:9e5896d1372858f8dd3344faf4e5014d21849c756c8d5701f78f8a103b372d92", @@ -76,13 +83,6 @@ "editable": true, "path": "." }, - "funcsigs": { - "hashes": [ - "sha256:330cc27ccbf7f1e992e69fef78261dc7c6569012cf397db8d3de0234e6c937ca", - "sha256:a7bb0f2cf3a3fd1ab2732cb49eba4252c2af4240442415b4abce3b87022a8f50" - ], - "version": "==1.0.2" - }, "fuzzywuzzy": { "hashes": [ "sha256:d40c22d2744dff84885b30bbfc07fab7875f641d070374331777a4d1808b8d4e", @@ -97,19 +97,6 @@ ], "version": "==2.6" }, - "jsonrpcserver": { - "hashes": [ - "sha256:ab8013cdee3f65d59c5d3f84c75be76a3492caa0b33ecaa3f0f69906cf3d9e92" - ], - "version": "==3.5.4" - }, - "jsonschema": { - "hashes": [ - "sha256:000e68abd33c972a5248544925a0cae7d1125f9bf6c58280d37546b946769a08", - "sha256:6ff5f3180870836cae40f06fa10419f557208175f13ad7bc26caa77beb1f6e02" - ], - "version": "==2.6.0" - }, "multidict": { "hashes": [ "sha256:1a1d76374a1e7fe93acef96b354a03c1d7f83e7512e225a527d283da0d7ba5e0", @@ -166,13 +153,6 @@ ], "version": "==1.1.1" }, - "six": { - "hashes": [ - "sha256:70e8a77beed4562e7f14fe23a786b54f6296e34344c23bc42f07b15018ff98e9", - "sha256:832dc0e10feb1aa2c68dcc57dbb658f1c7e65b9b61af69048abc87a2db00a0eb" - ], - "version": "==1.11.0" - }, "websockets": { "hashes": [ "sha256:09dfec40e9b73e8808c39ecdbc1733e33915a2b26b90c54566afc0af546a9ec3", diff --git a/docs/framework_bot.rst b/docs/framework_bot.rst index 31063c0c2..0f9a6cad3 100644 --- a/docs/framework_bot.rst +++ b/docs/framework_bot.rst @@ -13,6 +13,9 @@ RedBase :members: :exclude-members: get_context + .. automethod:: register_rpc_handler + .. automethod:: unregister_rpc_handler + Red ^^^ diff --git a/docs/framework_rpc.rst b/docs/framework_rpc.rst index c0190f78c..5d69fadbc 100644 --- a/docs/framework_rpc.rst +++ b/docs/framework_rpc.rst @@ -4,36 +4,60 @@ RPC === -.. currentmodule:: redbot.core.rpc - V3 comes default with an internal RPC server that may be used to remotely control the bot in various ways. Cogs must register functions to be exposed to RPC clients. Each of those functions must only take JSON serializable parameters and must return JSON serializable objects. -To begin, register all methods using individual calls to the :func:`Methods.add` method. +To enable the internal RPC server you must start the bot with the ``--rpc`` flag. ******** Examples ******** -Coming soon to a docs page near you! +.. code-block:: Python + + def setup(bot): + c = Cog() + bot.add_cog(c) + bot.register_rpc_handler(c.rpc_method) + +******************************* +Interacting with the RPC Server +******************************* + +The RPC server opens a websocket bound to port ``6133`` on ``127.0.0.1``. +This is not configurable for security reasons as broad access to this server gives anyone complete control over your bot. +To access the server you must find a library that implements websocket based JSONRPC in the language of your choice. + +There are a few built-in RPC methods to note: + +* ``GET_METHODS`` - Returns a list of available RPC methods. +* ``GET_METHOD_INFO`` - Will return the docstring for an available RPC method. Useful for finding information about the method's parameters and return values. +* ``GET_TOPIC`` - Returns a list of available RPC message topics. +* ``GET_SUBSCRIPTIONS`` - Returns a list of RPC subscriptions. +* ``SUBSCRIBE`` - Subscribes to an available RPC message topic. +* ``UNSUBSCRIBE`` - Unsubscribes from an RPC message topic. + +All RPC methods accept a list of parameters. +The built-in methods above expect their parameters to be in list format. + +All cog-based methods expect their parameter list to take one argument, a JSON object, in the following format:: + + params = [ + { + "args": [], # A list of positional arguments + "kwargs": {}, # A dictionary of keyword arguments + } + ] + + # As an example, here's a call to "get_method_info" + rpc_call("GET_METHOD_INFO", ["get_methods",]) + + # And here's a call to "core__load" + rpc_call("CORE__LOAD", {"args": [["general", "economy", "downloader"],], "kwargs": {}}) ************* API Reference ************* -.. py:attribute:: redbot.core.rpc.methods - - An instance of the :class:`Methods` class. - All attempts to register new RPC methods **MUST** use this object. - You should never create a new instance of the :class:`Methods` class! - -RPC -^^^ -.. autoclass:: redbot.core.rpc.RPC - :members: - -Methods -^^^^^^^ -.. autoclass:: redbot.core.rpc.Methods - :members: +Please see the :class:`redbot.core.bot.RedBase` class for details on the RPC handler register and unregister methods. diff --git a/redbot/__main__.py b/redbot/__main__.py index 27030dabb..92c0344e7 100644 --- a/redbot/__main__.py +++ b/redbot/__main__.py @@ -13,8 +13,7 @@ from redbot.core.events import init_events from redbot.core.cli import interactive_config, confirm, parse_cli_flags, ask_sentry from redbot.core.core_commands import Core from redbot.core.dev_commands import Dev -from redbot.core import rpc, __version__ -import redbot.meta +from redbot.core import __version__ import asyncio import logging.handlers import logging @@ -112,7 +111,7 @@ def main(): sys.exit(1) load_basic_configuration(cli_flags.instance_name) log, sentry_log = init_loggers(cli_flags) - red = Red(cli_flags, description=description, pm_help=None) + red = Red(cli_flags=cli_flags, description=description, pm_help=None) init_global_checks(red) init_events(red, cli_flags) red.add_cog(Core(red)) @@ -166,6 +165,10 @@ def main(): pending = asyncio.Task.all_tasks(loop=red.loop) gathered = asyncio.gather(*pending, loop=red.loop, return_exceptions=True) gathered.cancel() + try: + red.rpc.server.close() + except AttributeError: + pass sys.exit(red._shutdown_mode.value) diff --git a/redbot/core/bot.py b/redbot/core/bot.py index 4e3e5fae5..4e29659e3 100644 --- a/redbot/core/bot.py +++ b/redbot/core/bot.py @@ -18,12 +18,13 @@ from discord.voice_client import VoiceClient VoiceClient.warn_nacl = False from .cog_manager import CogManager -from . import Config, i18n, commands, rpc +from . import Config, i18n, commands +from .rpc import RPCMixin from .help_formatter import Help, help as help_ from .sentry import SentryManager -class RedBase(BotBase): +class RedBase(BotBase, RPCMixin): """Mixin for the main bot class. This exists because `Red` inherits from `discord.AutoShardedClient`, which @@ -33,7 +34,7 @@ class RedBase(BotBase): Selfbots should inherit from this mixin along with `discord.Client`. """ - def __init__(self, cli_flags, bot_dir: Path = Path.cwd(), **kwargs): + def __init__(self, *args, cli_flags=None, bot_dir: Path = Path.cwd(), **kwargs): self._shutdown_mode = ExitCodes.CRITICAL self.db = Config.get_core_conf(force_registration=True) self._co_owners = cli_flags.co_owner @@ -107,10 +108,7 @@ class RedBase(BotBase): self.cog_mgr = CogManager(paths=(str(self.main_dir / "cogs"),)) - super().__init__(formatter=Help(), **kwargs) - - if self.rpc_enabled: - self.rpc = rpc.RPC(self) + super().__init__(*args, formatter=Help(), **kwargs) self.remove_command("help") @@ -235,12 +233,24 @@ class RedBase(BotBase): lib_name = lib.__name__ # Thank you # find all references to the module + cog_names = [] # remove the cogs registered from the module for cogname, cog in self.cogs.copy().items(): if cog.__module__.startswith(lib_name): self.remove_cog(cogname) + cog_names.append(cogname) + + # remove all rpc handlers + for cogname in cog_names: + if cogname.upper() in self.rpc_handlers: + methods = self.rpc_handlers[cogname] + for meth in methods: + self.unregister_rpc_handler(meth) + + del self.rpc_handlers[cogname] + # first remove all the commands from the module for cmd in self.all_commands.copy().values(): if cmd.module.startswith(lib_name): diff --git a/redbot/core/core_commands.py b/redbot/core/core_commands.py index b389dc2cf..49919e50e 100644 --- a/redbot/core/core_commands.py +++ b/redbot/core/core_commands.py @@ -46,6 +46,13 @@ _ = i18n.Translator("Core", __file__) class CoreLogic: def __init__(self, bot: "Red"): self.bot = bot + self.bot.register_rpc_handler(self._load) + self.bot.register_rpc_handler(self._unload) + self.bot.register_rpc_handler(self._reload) + self.bot.register_rpc_handler(self._name) + self.bot.register_rpc_handler(self._prefixes) + self.bot.register_rpc_handler(self._version_info) + self.bot.register_rpc_handler(self._invite_url) async def _load(self, cog_names: list): """ diff --git a/redbot/core/events.py b/redbot/core/events.py index 4acdf9075..bf5274b21 100644 --- a/redbot/core/events.py +++ b/redbot/core/events.py @@ -18,6 +18,7 @@ from .data_manager import storage_type from .utils.chat_formatting import inline, bordered, pagify, box from .utils import fuzzy_command_search from colorama import Fore, Style, init +from . import rpc log = logging.getLogger("red") sentry_log = logging.getLogger("red.sentry") @@ -84,6 +85,9 @@ def init_events(bot, cli_flags): if packages: print("Loaded packages: " + ", ".join(packages)) + if bot.rpc_enabled: + await bot.rpc.initialize() + guilds = len(bot.guilds) users = len(set([m for m in bot.get_all_members()])) @@ -172,8 +176,6 @@ def init_events(bot, cli_flags): print("\nInvite URL: {}\n".format(invite_url)) bot.color = discord.Colour(await bot.db.color()) - if bot.rpc_enabled: - await bot.rpc.initialize() @bot.event async def on_error(event_method, *args, **kwargs): diff --git a/redbot/core/rpc.py b/redbot/core/rpc.py index 692833e68..dfe77542d 100644 --- a/redbot/core/rpc.py +++ b/redbot/core/rpc.py @@ -1,116 +1,74 @@ -import weakref +import asyncio from aiohttp import web -import jsonrpcserver.aio +from aiohttp_json_rpc import JsonRpc +from aiohttp_json_rpc.rpc import unpack_request_args -import inspect import logging -__all__ = ["methods", "RPC", "Methods"] - log = logging.getLogger("red.rpc") - -class Methods(jsonrpcserver.aio.AsyncMethods): - """ - Container class for all registered RPC methods, please use the existing `methods` - attribute rather than creating a new instance of this class. - - .. warning:: - - **NEVER** create a new instance of this class! - """ - - def __init__(self): - super().__init__() - - self._items = weakref.WeakValueDictionary() - - def add(self, method, name: str = None): - """ - Registers a method to the internal RPC server making it available for - RPC users to call. - - .. important:: - - Any method added here must take ONLY JSON serializable parameters and - MUST return a JSON serializable object. - - Parameters - ---------- - method : function - A reference to the function to register. - - name : str - Name of the function as seen by the RPC clients. - """ - if not inspect.iscoroutinefunction(method): - raise TypeError("Method must be a coroutine.") - - if name is None: - name = method.__qualname__ - - self._items[str(name)] = method - - def remove(self, *, name: str = None, method=None): - """ - Unregisters an RPC method. Either a name or reference to the method must - be provided and name will take priority. - - Parameters - ---------- - name : str - method : function - """ - if name and name in self._items: - del self._items[name] - - elif method and method in self._items.values(): - to_remove = [] - for name, val in self._items.items(): - if method == val: - to_remove.append(name) - - for name in to_remove: - del self._items[name] - - def all_methods(self): - """ - Lists all available method names. - - Returns - ------- - list of str - """ - return self._items.keys() +__all__ = ["RPC", "RPCMixin", "get_name"] -methods = Methods() +def get_name(func, prefix=None): + class_name = prefix or func.__self__.__class__.__name__.lower() + func_name = func.__name__.strip("_") + if class_name == "redrpc": + return func_name.upper() + return f"{class_name}__{func_name}".upper() -class BaseRPCMethodMixin: - def __init__(self): - methods.add(self.all_methods, name="all_methods") +class RedRpc(JsonRpc): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.add_methods(("", self.get_method_info)) - async def all_methods(self): - return list(methods.all_methods()) + def _add_method(self, method, prefix=""): + if not asyncio.iscoroutinefunction(method): + return + + name = get_name(method, prefix) + + self.methods[name] = method + + def remove_method(self, method): + meth_name = get_name(method) + new_methods = {} + for name, meth in self.methods.items(): + if name != meth_name: + new_methods[name] = meth + self.methods = new_methods + + def remove_methods(self, prefix: str): + new_methods = {} + for name, meth in self.methods.items(): + splitted = name.split("__") + if len(splitted) < 2 or splitted[0] != prefix: + new_methods[name] = meth + self.methods = new_methods + + async def get_method_info(self, request): + method_name = request.params[0] + if method_name in self.methods: + return self.methods[method_name].__doc__ + return "No docstring available." -class RPC(BaseRPCMethodMixin): +class RPC: """ RPC server manager. """ - def __init__(self, bot): - self.app = web.Application(loop=bot.loop) - self.app.router.add_post("/rpc", self.handle) + def __init__(self): + self.app = web.Application() + self._rpc = RedRpc() + self.app.router.add_route("*", "/", self._rpc) self.app_handler = self.app.make_handler() self.server = None - super().__init__() - async def initialize(self): """ Finalizes the initialization of the RPC server and allows it to begin @@ -125,10 +83,79 @@ class RPC(BaseRPCMethodMixin): """ self.server.close() - async def handle(self, request): - request = await request.text() - response = await methods.dispatch(request) - if response.is_notification: - return web.Response() - else: - return web.json_response(response, status=response.http_status) + def add_method(self, method, prefix: str = None): + if prefix is None: + prefix = method.__self__.__class__.__name__.lower() + + if not asyncio.iscoroutinefunction(method): + raise TypeError("RPC methods must be coroutines.") + + self._rpc.add_methods((prefix, unpack_request_args(method))) + + def add_multi_method(self, *methods, prefix: str = None): + if not all(asyncio.iscoroutinefunction(m) for m in methods): + raise TypeError("RPC methods must be coroutines.") + + for method in methods: + self.add_method(method, prefix=prefix) + + def remove_method(self, method): + self._rpc.remove_method(method) + + def remove_methods(self, prefix: str): + self._rpc.remove_methods(prefix) + + +class RPCMixin: + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.rpc = RPC() + + self.rpc_handlers = {} # Lowered cog name to method + + def register_rpc_handler(self, method): + """ + Registers a method to act as an RPC handler if the internal RPC server is active. + + When calling this method through the RPC server, use the naming scheme "cogname__methodname". + + .. important:: + + All parameters to RPC handler methods must be JSON serializable objects. + The return value of handler methods must also be JSON serializable. + + Parameters + ---------- + method : coroutine + The method to register with the internal RPC server. + """ + self.rpc.add_method(method) + + cog_name = method.__self__.__class__.__name__.upper() + if cog_name not in self.rpc_handlers: + self.rpc_handlers[cog_name] = [] + + self.rpc_handlers[cog_name].append(method) + + def unregister_rpc_handler(self, method): + """ + Unregisters an RPC method handler. + + This will be called automatically for you on cog unload and will pass silently if the + method is not previously registered. + + Parameters + ---------- + method : coroutine + The method to unregister from the internal RPC server. + """ + self.rpc.remove_method(method) + + name = get_name(method) + cog_name = name.split("__")[0] + + if cog_name in self.rpc_handlers: + try: + self.rpc_handlers[cog_name].remove(method) + except ValueError: + pass diff --git a/requirements.txt b/requirements.txt index 40efb6ee9..24280b82e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ aiohttp>=2.0.0,<2.3.0 appdirs==1.4.3 raven==6.5.0 colorama==0.3.9 -jsonrpcserver +aiohttp-json-rpc==0.8.7 pyyaml==3.12 fuzzywuzzy[speedup]<=0.16.0 Red-Trivia>=1.1.1 diff --git a/tests/conftest.py b/tests/conftest.py index 4ecfc2048..fb715820a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -154,7 +154,7 @@ def red(config_fr): Config.get_core_conf = lambda *args, **kwargs: config_fr - red = Red(cli_flags, description=description, pm_help=None) + red = Red(cli_flags=cli_flags, description=description, pm_help=None) yield red diff --git a/tests/core/test_rpc.py b/tests/core/test_rpc.py new file mode 100644 index 000000000..7fa1f6564 --- /dev/null +++ b/tests/core/test_rpc.py @@ -0,0 +1,145 @@ +import pytest +from redbot.core.rpc import RPC, RPCMixin, get_name + +from unittest.mock import MagicMock + + +@pytest.fixture() +def rpc(): + return RPC() + + +@pytest.fixture() +def rpcmixin(): + r = RPCMixin() + r.rpc = MagicMock(spec=RPC) + return r + + +@pytest.fixture() +def cog(): + class Cog: + async def cofunc(*args, **kwargs): + pass + + async def cofunc2(*args, **kwargs): + pass + + async def cofunc3(*args, **kwargs): + pass + + def func(*args, **kwargs): + pass + + return Cog() + + +@pytest.fixture() +def existing_func(rpc, cog): + rpc.add_method(cog.cofunc) + + return cog.cofunc + + +@pytest.fixture() +def existing_multi_func(rpc, cog): + funcs = [cog.cofunc, cog.cofunc2, cog.cofunc3] + rpc.add_multi_method(*funcs) + + return funcs + + +def test_get_name(cog): + assert get_name(cog.cofunc) == "COG__COFUNC" + assert get_name(cog.cofunc2) == "COG__COFUNC2" + assert get_name(cog.func) == "COG__FUNC" + + +def test_internal_methods_exist(rpc): + assert "GET_METHODS" in rpc._rpc.methods + + +def test_add_method(rpc, cog): + rpc.add_method(cog.cofunc) + + assert get_name(cog.cofunc) in rpc._rpc.methods + + +def test_double_add(rpc, cog): + rpc.add_method(cog.cofunc) + count = len(rpc._rpc.methods) + + rpc.add_method(cog.cofunc) + + assert count == len(rpc._rpc.methods) + + +def test_add_notcoro_method(rpc, cog): + with pytest.raises(TypeError): + rpc.add_method(cog.func) + + +def test_add_multi(rpc, cog): + funcs = [cog.cofunc, cog.cofunc2, cog.cofunc3] + rpc.add_multi_method(*funcs) + + names = [get_name(f) for f in funcs] + + assert all(n in rpc._rpc.methods for n in names) + + +def test_add_multi_bad(rpc, cog): + funcs = [cog.cofunc, cog.cofunc2, cog.cofunc3, cog.func] + + with pytest.raises(TypeError): + rpc.add_multi_method(*funcs) + + names = [get_name(f) for f in funcs] + + assert not any(n in rpc._rpc.methods for n in names) + + +def test_remove_method(rpc, existing_func): + before_count = len(rpc._rpc.methods) + rpc.remove_method(existing_func) + + assert get_name(existing_func) not in rpc._rpc.methods + assert before_count - 1 == len(rpc._rpc.methods) + + +def test_remove_multi_method(rpc, existing_multi_func): + before_count = len(rpc._rpc.methods) + name = get_name(existing_multi_func[0]) + prefix = name.split("__")[0] + + rpc.remove_methods(prefix) + + assert before_count - len(existing_multi_func) == len(rpc._rpc.methods) + + names = [get_name(f) for f in existing_multi_func] + + assert not any(n in rpc._rpc.methods for n in names) + + +def test_rpcmixin_register(rpcmixin, cog): + rpcmixin.register_rpc_handler(cog.cofunc) + + assert rpcmixin.rpc.add_method.called_once_with(cog.cofunc) + + name = get_name(cog.cofunc) + cogname = name.split("__")[0] + + assert cogname in rpcmixin.rpc_handlers + + +def test_rpcmixin_unregister(rpcmixin, cog): + rpcmixin.register_rpc_handler(cog.cofunc) + rpcmixin.unregister_rpc_handler(cog.cofunc) + + assert rpcmixin.rpc.remove_method.called_once_with(cog.cofunc) + + name = get_name(cog.cofunc) + cogname = name.split("__")[0] + + if cogname in rpcmixin.rpc_handlers: + assert cog.cofunc not in rpcmixin.rpc_handlers[cogname] diff --git a/tests/rpc_test.html b/tests/rpc_test.html new file mode 100644 index 000000000..561e378e8 --- /dev/null +++ b/tests/rpc_test.html @@ -0,0 +1,21 @@ + +