From 7525af0ca1758e0cacf3b691aff9e4a1846a4e41 Mon Sep 17 00:00:00 2001
From: "Colton Wolkins (Indicio work address)" <colton@indicio.tech>
Date: Tue, 19 Apr 2022 15:53:35 -0600
Subject: [PATCH] feat: Add the ability to deny specific plugins from loading

As part of #1734, I found that the Revocation Notification V2 code
failed to properly run due to the fact that the V1 code was deleting the
DB objects before V2 could get to it. @dbluhm and I decided that adding
an option to deny specific plugins from loading was one of the better
solutions to the problem.

Before the plugin has been initialized, any plugins listed within
`--deny-plugins` will be unregistered to prevent them from initializing
and potentially interferring with other plugins that have conflicts.

Signed-off-by: Colton Wolkins (Indicio work address) <colton@indicio.tech>
---
 aries_cloudagent/config/argparse.py           | 17 +++++++++++++++++
 aries_cloudagent/config/default_context.py    |  3 +++
 aries_cloudagent/core/plugin_registry.py      |  7 +++++++
 .../core/tests/test_plugin_registry.py        | 19 +++++++++++++++++++
 4 files changed, 46 insertions(+)

diff --git a/aries_cloudagent/config/argparse.py b/aries_cloudagent/config/argparse.py
index 4b91b22aad..0216ebeaa1 100644
--- a/aries_cloudagent/config/argparse.py
+++ b/aries_cloudagent/config/argparse.py
@@ -534,6 +534,20 @@ def add_arguments(self, parser: ArgumentParser):
             ),
         )
 
+        parser.add_argument(
+            "--deny-plugin",
+            dest="denied_plugins",
+            type=str,
+            action="append",
+            required=False,
+            metavar="<module>",
+            env_var="ACAPY_DENY_PLUGIN",
+            help=(
+                "Deny <module> plugin module from loading. Multiple "
+                "instances of this parameter can be specified."
+            ),
+        )
+
         parser.add_argument(
             "--plugin-config",
             dest="plugin_config",
@@ -611,6 +625,9 @@ def get_settings(self, args: Namespace) -> dict:
         if args.external_plugins:
             settings["external_plugins"] = args.external_plugins
 
+        if args.denied_plugins:
+            settings["denied_plugins"] = args.denied_plugins
+
         if args.plugin_config:
             with open(args.plugin_config, "r") as stream:
                 settings["plugin_config"] = yaml.safe_load(stream)
diff --git a/aries_cloudagent/config/default_context.py b/aries_cloudagent/config/default_context.py
index 4e9b0eddef..d92531f1ab 100644
--- a/aries_cloudagent/config/default_context.py
+++ b/aries_cloudagent/config/default_context.py
@@ -136,5 +136,8 @@ async def load_plugins(self, context: InjectionContext):
         for plugin_path in self.settings.get("external_plugins", []):
             plugin_registry.register_plugin(plugin_path)
 
+        for plugin_path in self.settings.get("deny_plugins", []):
+            plugin_registry.unregister_plugin(plugin_path)
+
         # Register message protocols
         await plugin_registry.init_context(context)
diff --git a/aries_cloudagent/core/plugin_registry.py b/aries_cloudagent/core/plugin_registry.py
index 193c695241..1a4aa62fc5 100644
--- a/aries_cloudagent/core/plugin_registry.py
+++ b/aries_cloudagent/core/plugin_registry.py
@@ -177,6 +177,13 @@ def register_plugin(self, module_name: str) -> ModuleType:
         #     self._plugins[module_name] = mod
         #     return mod
 
+    def unregister_plugin(self, module_name: str):
+        """Unregister a plugin module. MUST BE CALLED BEFORE INITIALIZATION."""
+        if module_name not in self._plugins:
+            return
+
+        del self._plugins[module_name]
+
     def register_package(self, package_name: str) -> Sequence[ModuleType]:
         """Register all modules (sub-packages) under a given package name."""
         try:
diff --git a/aries_cloudagent/core/tests/test_plugin_registry.py b/aries_cloudagent/core/tests/test_plugin_registry.py
index 7213e96b7e..2b16a644e4 100644
--- a/aries_cloudagent/core/tests/test_plugin_registry.py
+++ b/aries_cloudagent/core/tests/test_plugin_registry.py
@@ -478,6 +478,25 @@ class MODULE:
             ]
             assert self.registry.register_plugin("dummy") == obj
 
+    async def test_unregister_plugin_has_setup(self):
+        class MODULE:
+            setup = "present"
+
+        obj = MODULE()
+        with async_mock.patch.object(
+            ClassLoader, "load_module", async_mock.MagicMock()
+        ) as load_module:
+            load_module.side_effect = [
+                obj,  # module
+                None,  # routes
+                None,  # message types
+                None,  # definition without versions attr
+            ]
+            assert self.registry.register_plugin("dummy") == obj
+            assert "dummy" in self.registry._plugins.keys()
+            self.registry.unregister_plugin("dummy")
+            assert "dummy" not in self.registry._plugins.keys()
+
     async def test_register_definitions_malformed(self):
         class MODULE:
             no_setup = "no setup attr"