Repository: abhinavsingh/proxy.py
Pull request #420 (2020-08-11 20:38:14)

Allow to use types when embeding Proxy
When embedding `Proxy` in another application (e.g. automatic tests), it may be simpler to directly give the `type` of the plugins to load via the `plugins` keyword argument. 
This requires to convert the `plugins` keyword argument into a `list` of `bytes` or `type` and modify `Flags.load_plugin`. I corrected the test so that they comply to the new type of the `plugin` keyword argument.
I also included unit tests for `Flags.load_plugin` as well as a correction ensuring plugins are only loaded once (see the three last unit tests).




Patch:
diff --git a/proxy/common/flags.py b/proxy/common/flags.py
index e22ae5d..12318c9 100644
--- a/proxy/common/flags.py
+++ b/proxy/common/flags.py
@@ -21,7 +21,7 @@ import multiprocessing
 import sys
 import inspect
 
-from typing import Optional, Dict, List, TypeVar, Type, cast, Any, Tuple
+from typing import Optional, Dict, List, TypeVar, Type, cast, Any, Tuple, Union
 
 from .types import IpAddress
 from .utils import text_, bytes_
@@ -184,10 +184,9 @@ class Flags:
 
         # Load default plugins along with user provided --plugins
         plugins = Flags.load_plugins(
-            bytes_(
-                '%s,%s' %
-                (text_(COMMA).join(collections.OrderedDict(default_plugins).keys()),
-                    opts.get('plugins', args.plugins))))
+            [bytes_(p) for p in collections.OrderedDict(default_plugins).keys()] +
+            [p if isinstance(p, type) else bytes_(p) for p in opts.get('plugins', args.plugins.split(text_(COMMA)))]
+        )
 
         # proxy.py currently cannot serve over HTTPS and perform TLS interception
         # at the same time.  Check if user is trying to enable both feature
@@ -508,7 +507,7 @@ class Flags:
                     'Open file soft limit set to %d', soft_limit)
 
     @staticmethod
-    def load_plugins(plugins: bytes) -> Dict[bytes, List[type]]:
+    def load_plugins(plugins: List[Union[bytes, type]]) -> Dict[bytes, List[type]]:
         """Accepts a comma separated list of Python modules and returns
         a list of respective Python classes."""
         p: Dict[bytes, List[type]] = {
@@ -517,30 +516,29 @@ class Flags:
             b'HttpWebServerBasePlugin': [],
             b'ProxyDashboardWebsocketPlugin': []
         }
-        for plugin_ in plugins.split(COMMA):
-            plugin = text_(plugin_.strip())
-            if plugin == '':
-                continue
-            module_name, klass_name = plugin.rsplit(text_(DOT), 1)
-            klass = getattr(
-                importlib.import_module(
-                    module_name.replace(
-                        os.path.sep, text_(DOT))),
-                klass_name)
+        for plugin_ in plugins:
+            if isinstance(plugin_, type):
+                module_name = '__main__'
+                klass = plugin_
+            else:
+                plugin = text_(plugin_.strip())
+                if plugin == '':
+                    continue
+                module_name, klass_name = plugin.rsplit(text_(DOT), 1)
+                klass = getattr(
+                    importlib.import_module(
+                        module_name.replace(
+                            os.path.sep, text_(DOT))),
+                    klass_name)
             mro = list(inspect.getmro(klass))
             mro.reverse()
             iterator = iter(mro)
             while next(iterator) is not abc.ABC:
                 pass
             base_klass = next(iterator)
-            p[bytes_(base_klass.__name__)].append(klass)
-            logger.info(
-                'Loaded %s %s.%s',
-                'plugin' if klass.__name__ != 'HttpWebServerRouteHandler' else 'route',
-                module_name,
-                # HttpWebServerRouteHandler route decorator adds a special
-                # staticmethod to return decorated function name
-                klass.__name__ if klass.__name__ != 'HttpWebServerRouteHandler' else klass.name())
+            if klass not in p[bytes_(base_klass.__name__)]:
+                p[bytes_(base_klass.__name__)].append(klass)
+            logger.info('Loaded plugin %s.%s', module_name, klass.__name__)
         return p
 
     @staticmethod


Test patch:
diff --git a/tests/common/test_flags.py b/tests/common/test_flags.py
new file mode 100644
index 0000000..93774fc
--- /dev/null
+++ b/tests/common/test_flags.py
@@ -0,0 +1,112 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+import unittest
+
+from typing import List, Dict
+
+from proxy.common.flags import Flags
+from proxy.http.proxy import HttpProxyPlugin
+from proxy.plugin import CacheResponsesPlugin
+from proxy.plugin import FilterByUpstreamHostPlugin
+
+
+class TestFlags(unittest.TestCase):
+    def assert_plugins(self, expected: Dict[str, List[type]]) -> None:
+        for k in expected:
+            self.assertIn(k.encode(), self.flags.plugins)
+            for p in expected[k]:
+                self.assertIn(p, self.flags.plugins[k.encode()])
+                self.assertEqual(len([o for o in self.flags.plugins[k.encode()] if o == p]), 1)
+
+    def test_load_plugin_from_bytes(self) -> None:
+        self.flags = Flags.initialize([], plugins=[
+            b'proxy.plugin.CacheResponsesPlugin',
+        ])
+        self.assert_plugins({'HttpProxyBasePlugin': [CacheResponsesPlugin]})
+
+    def test_load_plugins_from_bytes(self) -> None:
+        self.flags = Flags.initialize([], plugins=[
+            b'proxy.plugin.CacheResponsesPlugin',
+            b'proxy.plugin.FilterByUpstreamHostPlugin',
+        ])
+        self.assert_plugins({'HttpProxyBasePlugin': [
+            CacheResponsesPlugin,
+            FilterByUpstreamHostPlugin,
+        ]})
+
+    def test_load_plugin_from_args(self) -> None:
+        self.flags = Flags.initialize([
+            '--plugins', 'proxy.plugin.CacheResponsesPlugin',
+        ])
+        self.assert_plugins({'HttpProxyBasePlugin': [CacheResponsesPlugin]})
+
+    def test_load_plugins_from_args(self) -> None:
+        self.flags = Flags.initialize([
+            '--plugins', 'proxy.plugin.CacheResponsesPlugin,proxy.plugin.FilterByUpstreamHostPlugin',
+        ])
+        self.assert_plugins({'HttpProxyBasePlugin': [
+            CacheResponsesPlugin,
+            FilterByUpstreamHostPlugin,
+        ]})
+
+    def test_load_plugin_from_class(self) -> None:
+        self.flags = Flags.initialize([], plugins=[
+            CacheResponsesPlugin,
+        ])
+        self.assert_plugins({'HttpProxyBasePlugin': [CacheResponsesPlugin]})
+
+    def test_load_plugins_from_class(self) -> None:
+        self.flags = Flags.initialize([], plugins=[
+            CacheResponsesPlugin,
+            FilterByUpstreamHostPlugin,
+        ])
+        self.assert_plugins({'HttpProxyBasePlugin': [
+            CacheResponsesPlugin,
+            FilterByUpstreamHostPlugin,
+        ]})
+
+    def test_load_plugins_from_bytes_and_class(self) -> None:
+        self.flags = Flags.initialize([], plugins=[
+            CacheResponsesPlugin,
+            b'proxy.plugin.FilterByUpstreamHostPlugin',
+        ])
+        self.assert_plugins({'HttpProxyBasePlugin': [
+            CacheResponsesPlugin,
+            FilterByUpstreamHostPlugin,
+        ]})
+
+    def test_unique_plugin_from_bytes(self) -> None:
+        self.flags = Flags.initialize([], plugins=[
+            b'proxy.http.proxy.HttpProxyPlugin',
+        ])
+        self.assert_plugins({'HttpProtocolHandlerPlugin': [
+            HttpProxyPlugin,
+        ]})
+
+    def test_unique_plugin_from_args(self) -> None:
+        self.flags = Flags.initialize([
+            '--plugins', 'proxy.http.proxy.HttpProxyPlugin',
+        ])
+        self.assert_plugins({'HttpProtocolHandlerPlugin': [
+            HttpProxyPlugin,
+        ]})
+
+    def test_unique_plugin_from_class(self) -> None:
+        self.flags = Flags.initialize([], plugins=[
+            HttpProxyPlugin,
+        ])
+        self.assert_plugins({'HttpProtocolHandlerPlugin': [
+            HttpProxyPlugin,
+        ]})
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/http/test_protocol_handler.py b/tests/http/test_protocol_handler.py
index ca4dac0..c0b9688 100644
--- a/tests/http/test_protocol_handler.py
+++ b/tests/http/test_protocol_handler.py
@@ -40,8 +40,10 @@ class TestHttpProtocolHandler(unittest.TestCase):
 
         self.http_server_port = 65535
         self.flags = Flags()
-        self.flags.plugins = Flags.load_plugins(
-            b'proxy.http.proxy.HttpProxyPlugin,proxy.http.server.HttpWebServerPlugin')
+        self.flags.plugins = Flags.load_plugins([
+            b'proxy.http.proxy.HttpProxyPlugin',
+            b'proxy.http.server.HttpWebServerPlugin',
+        ])
 
         self.mock_selector = mock_selector
         self.protocol_handler = HttpProtocolHandler(
@@ -173,8 +175,10 @@ class TestHttpProtocolHandler(unittest.TestCase):
         flags = Flags(
             auth_code=b'Basic %s' %
                       base64.b64encode(b'user:pass'))
-        flags.plugins = Flags.load_plugins(
-            b'proxy.http.proxy.HttpProxyPlugin,proxy.http.server.HttpWebServerPlugin')
+        flags.plugins = Flags.load_plugins([
+            b'proxy.http.proxy.HttpProxyPlugin',
+            b'proxy.http.server.HttpWebServerPlugin',
+        ])
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr), flags=flags)
         self.protocol_handler.initialize()
@@ -205,8 +209,10 @@ class TestHttpProtocolHandler(unittest.TestCase):
         flags = Flags(
             auth_code=b'Basic %s' %
                       base64.b64encode(b'user:pass'))
-        flags.plugins = Flags.load_plugins(
-            b'proxy.http.proxy.HttpProxyPlugin,proxy.http.server.HttpWebServerPlugin')
+        flags.plugins = Flags.load_plugins([
+            b'proxy.http.proxy.HttpProxyPlugin',
+            b'proxy.http.server.HttpWebServerPlugin',
+        ])
 
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr), flags=flags)
@@ -253,8 +259,10 @@ class TestHttpProtocolHandler(unittest.TestCase):
         flags = Flags(
             auth_code=b'Basic %s' %
                       base64.b64encode(b'user:pass'))
-        flags.plugins = Flags.load_plugins(
-            b'proxy.http.proxy.HttpProxyPlugin,proxy.http.server.HttpWebServerPlugin')
+        flags.plugins = Flags.load_plugins([
+            b'proxy.http.proxy.HttpProxyPlugin',
+            b'proxy.http.server.HttpWebServerPlugin'
+        ])
 
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr), flags=flags)
diff --git a/tests/http/test_web_server.py b/tests/http/test_web_server.py
index 3a21fc9..6ab15a9 100644
--- a/tests/http/test_web_server.py
+++ b/tests/http/test_web_server.py
@@ -34,8 +34,10 @@ class TestWebServerPlugin(unittest.TestCase):
         self._conn = mock_fromfd.return_value
         self.mock_selector = mock_selector
         self.flags = Flags()
-        self.flags.plugins = Flags.load_plugins(
-            b'proxy.http.proxy.HttpProxyPlugin,proxy.http.server.HttpWebServerPlugin')
+        self.flags.plugins = Flags.load_plugins([
+            b'proxy.http.proxy.HttpProxyPlugin',
+            b'proxy.http.server.HttpWebServerPlugin',
+        ])
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr),
             flags=self.flags)
@@ -95,8 +97,10 @@ class TestWebServerPlugin(unittest.TestCase):
                 events=selectors.EVENT_READ,
                 data=None), selectors.EVENT_READ), ]
         flags = Flags()
-        flags.plugins = Flags.load_plugins(
-            b'proxy.http.proxy.HttpProxyPlugin,proxy.http.server.HttpWebServerPlugin')
+        flags.plugins = Flags.load_plugins([
+            b'proxy.http.proxy.HttpProxyPlugin',
+            b'proxy.http.server.HttpWebServerPlugin',
+        ])
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr),
             flags=flags)
@@ -146,8 +150,10 @@ class TestWebServerPlugin(unittest.TestCase):
         flags = Flags(
             enable_static_server=True,
             static_server_dir=static_server_dir)
-        flags.plugins = Flags.load_plugins(
-            b'proxy.http.proxy.HttpProxyPlugin,proxy.http.server.HttpWebServerPlugin')
+        flags.plugins = Flags.load_plugins([
+            b'proxy.http.proxy.HttpProxyPlugin',
+            b'proxy.http.server.HttpWebServerPlugin',
+        ])
 
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr),
@@ -194,8 +200,10 @@ class TestWebServerPlugin(unittest.TestCase):
                 data=None), selectors.EVENT_WRITE)], ]
 
         flags = Flags(enable_static_server=True)
-        flags.plugins = Flags.load_plugins(
-            b'proxy.http.proxy.HttpProxyPlugin,proxy.http.server.HttpWebServerPlugin')
+        flags.plugins = Flags.load_plugins([
+            b'proxy.http.proxy.HttpProxyPlugin',
+            b'proxy.http.server.HttpWebServerPlugin',
+        ])
 
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr),
@@ -230,9 +238,11 @@ class TestWebServerPlugin(unittest.TestCase):
 
     def init_and_make_pac_file_request(self, pac_file: str) -> None:
         flags = Flags(pac_file=pac_file)
-        flags.plugins = Flags.load_plugins(
-            b'proxy.http.proxy.HttpProxyPlugin,proxy.http.server.HttpWebServerPlugin,'
-            b'proxy.http.server.HttpWebServerPacFilePlugin')
+        flags.plugins = Flags.load_plugins([
+            b'proxy.http.proxy.HttpProxyPlugin',
+            b'proxy.http.server.HttpWebServerPlugin',
+            b'proxy.http.server.HttpWebServerPacFilePlugin',
+        ])
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr),
             flags=flags)