Repository: abhinavsingh/proxy.py
Pull request #453 (2020-10-14 14:17:49)

Separate basic auth plugin outside of core server

codecov[bot]: # [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453?src=pr&el=h1) Report
> Merging [#453](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453?src=pr&el=desc) into [develop](https://codecov.io/gh/abhinavsingh/proxy.py/commit/0744cd8e7fc6985f5ab5656685833b85a096ddf0?el=desc) will **decrease** coverage by `0.08%`.
> The diff coverage is `76.66%`.

[![Impacted file tree graph](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453/graphs/tree.svg?width=650&height=150&src=pr&token=Zh9J7b4la2)](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453?src=pr&el=tree)

```diff
@@             Coverage Diff             @@
##           develop     #453      +/-   ##
===========================================
- Coverage    81.26%   81.18%   -0.09%     
===========================================
  Files           79       80       +1     
  Lines         2995     3014      +19     
===========================================
+ Hits          2434     2447      +13     
- Misses         561      567       +6     
```


| [Impacted Files](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [proxy/http/proxy/auth.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453/diff?src=pr&el=tree#diff-cHJveHkvaHR0cC9wcm94eS9hdXRoLnB5) | `69.56% <69.56%> (ø)` | |
| [proxy/common/constants.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453/diff?src=pr&el=tree#diff-cHJveHkvY29tbW9uL2NvbnN0YW50cy5weQ==) | `100.00% <100.00%> (ø)` | |
| [proxy/core/connection/connection.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453/diff?src=pr&el=tree#diff-cHJveHkvY29yZS9jb25uZWN0aW9uL2Nvbm5lY3Rpb24ucHk=) | `100.00% <100.00%> (ø)` | |
| [proxy/http/proxy/\_\_init\_\_.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453/diff?src=pr&el=tree#diff-cHJveHkvaHR0cC9wcm94eS9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [proxy/http/proxy/server.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453/diff?src=pr&el=tree#diff-cHJveHkvaHR0cC9wcm94eS9zZXJ2ZXIucHk=) | `77.59% <100.00%> (-0.26%)` | :arrow_down: |
| [proxy/proxy.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453/diff?src=pr&el=tree#diff-cHJveHkvcHJveHkucHk=) | `94.76% <100.00%> (+0.02%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453?src=pr&el=footer). Last update [0744cd8...b0c5934](https://codecov.io/gh/abhinavsingh/proxy.py/pull/453?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/Makefile b/Makefile
index 92ae5d0..e0331a5 100644
--- a/Makefile
+++ b/Makefile
@@ -110,7 +110,7 @@ lib-coverage:
 	open htmlcov/index.html
 
 lib-profile:
-	sudo py-spy -F -f profile.svg -d 3600 proxy.py
+	sudo py-spy record -o profile.svg -t -F -s -- python -m proxy
 
 dashboard:
 	pushd dashboard && npm run build && popd
diff --git a/proxy/common/constants.py b/proxy/common/constants.py
index a65b087..21e3955 100644
--- a/proxy/common/constants.py
+++ b/proxy/common/constants.py
@@ -86,6 +86,7 @@ PLUGIN_PAC_FILE = 'proxy.http.server.HttpWebServerPacFilePlugin'
 PLUGIN_DEVTOOLS_PROTOCOL = 'proxy.http.inspector.DevtoolsProtocolPlugin'
 PLUGIN_DASHBOARD = 'proxy.dashboard.dashboard.ProxyDashboard'
 PLUGIN_INSPECT_TRAFFIC = 'proxy.dashboard.inspect_traffic.InspectTrafficPlugin'
+PLUGIN_PROXY_AUTH = 'proxy.http.proxy.AuthPlugin'
 
 PY2_DEPRECATION_MESSAGE = '''DEPRECATION: proxy.py no longer supports Python 2.7.  Kindly upgrade to Python 3+. '
                 'If for some reasons you cannot upgrade, use'
diff --git a/proxy/core/connection/connection.py b/proxy/core/connection/connection.py
index 3aa72ee..73eabe2 100644
--- a/proxy/core/connection/connection.py
+++ b/proxy/core/connection/connection.py
@@ -88,5 +88,6 @@ class TcpConnection(ABC):
             self.buffer.pop(0)
         else:
             self.buffer[0] = memoryview(mv[sent:])
+        del mv
         logger.debug('flushed %d bytes to %s' % (sent, self.tag))
         return sent
diff --git a/proxy/http/proxy/__init__.py b/proxy/http/proxy/__init__.py
index afd3527..4e18002 100644
--- a/proxy/http/proxy/__init__.py
+++ b/proxy/http/proxy/__init__.py
@@ -10,8 +10,10 @@
 """
 from .plugin import HttpProxyBasePlugin
 from .server import HttpProxyPlugin
+from .auth import AuthPlugin
 
 __all__ = [
     'HttpProxyBasePlugin',
     'HttpProxyPlugin',
+    'AuthPlugin',
 ]
diff --git a/proxy/http/proxy/auth.py b/proxy/http/proxy/auth.py
new file mode 100644
index 0000000..caa7cf4
--- /dev/null
+++ b/proxy/http/proxy/auth.py
@@ -0,0 +1,50 @@
+# -*- 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.
+"""
+from typing import Optional
+
+from ..exception import ProxyAuthenticationFailed
+from ...common.flag import flags
+from ...common.constants import DEFAULT_BASIC_AUTH
+from ...http.parser import HttpParser
+from ...http.proxy import HttpProxyBasePlugin
+
+
+flags.add_argument(
+    '--basic-auth',
+    type=str,
+    default=DEFAULT_BASIC_AUTH,
+    help='Default: No authentication. Specify colon separated user:password '
+    'to enable basic authentication.')
+
+
+class AuthPlugin(HttpProxyBasePlugin):
+    """Performs proxy authentication."""
+
+    def before_upstream_connection(
+            self, request: HttpParser) -> Optional[HttpParser]:
+        if self.flags.auth_code:
+            if b'proxy-authorization' not in request.headers:
+                raise ProxyAuthenticationFailed()
+            parts = request.headers[b'proxy-authorization'][1].split()
+            if len(parts) != 2 and parts[0].lower(
+            ) != b'basic' and parts[1] != self.flags.auth_code:
+                raise ProxyAuthenticationFailed()
+        return request
+
+    def handle_client_request(
+            self, request: HttpParser) -> Optional[HttpParser]:
+        return request
+
+    def handle_upstream_chunk(self, chunk: memoryview) -> memoryview:
+        return chunk
+
+    def on_upstream_connection_close(self) -> None:
+        pass
diff --git a/proxy/http/proxy/server.py b/proxy/http/proxy/server.py
index 966b210..f60ff54 100644
--- a/proxy/http/proxy/server.py
+++ b/proxy/http/proxy/server.py
@@ -20,7 +20,7 @@ from typing import Optional, List, Union, Dict, cast, Any, Tuple
 
 from .plugin import HttpProxyBasePlugin
 from ..plugin import HttpProtocolHandlerPlugin
-from ..exception import HttpProtocolException, ProxyConnectionFailed, ProxyAuthenticationFailed
+from ..exception import HttpProtocolException, ProxyConnectionFailed
 from ..codes import httpStatusCodes
 from ..parser import HttpParser, httpParserStates, httpParserTypes
 from ..methods import httpMethods
@@ -105,8 +105,7 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
         reason=b'Connection established'
     ))
 
-    # Used to synchronize with other HttpProxyPlugin instances while
-    # generating certificates
+    # Used to synchronization during certificate generation.
     lock = threading.Lock()
 
     def __init__(
@@ -322,8 +321,6 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
 
         self.emit_request_complete()
 
-        self.authenticate()
-
         # Note: can raise HttpRequestRejected exception
         # Invoke plugin.before_upstream_connection
         do_connect = True
@@ -533,15 +530,6 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
         logger.debug(
             'TLS interception using %s', generated_cert)
 
-    def authenticate(self) -> None:
-        if self.flags.auth_code:
-            if b'proxy-authorization' not in self.request.headers:
-                raise ProxyAuthenticationFailed()
-            parts = self.request.headers[b'proxy-authorization'][1].split()
-            if len(parts) != 2 and parts[0].lower(
-            ) != b'basic' and parts[1] != self.flags.auth_code:
-                raise ProxyAuthenticationFailed()
-
     def connect_upstream(self) -> None:
         host, port = self.request.host, self.request.port
         if host and port:
diff --git a/proxy/proxy.py b/proxy/proxy.py
index 8a64408..023c43c 100644
--- a/proxy/proxy.py
+++ b/proxy/proxy.py
@@ -32,7 +32,7 @@ from .common.version import __version__
 from .core.acceptor import AcceptorPool
 from .http.handler import HttpProtocolHandler
 from .common.flag import flags
-from .common.constants import COMMA, DEFAULT_BASIC_AUTH, DEFAULT_DATA_DIRECTORY_PATH
+from .common.constants import COMMA, DEFAULT_DATA_DIRECTORY_PATH, PLUGIN_PROXY_AUTH
 from .common.constants import DEFAULT_DEVTOOLS_WS_PATH, DEFAULT_DISABLE_HEADERS
 from .common.constants import DEFAULT_DISABLE_HTTP_PROXY, DEFAULT_NUM_WORKERS
 from .common.constants import DEFAULT_ENABLE_DASHBOARD, DEFAULT_ENABLE_DEVTOOLS
@@ -54,12 +54,6 @@ flags.add_argument(
     type=str,
     default=DEFAULT_PID_FILE,
     help='Default: None. Save parent process ID to a file.')
-flags.add_argument(
-    '--basic-auth',
-    type=str,
-    default=DEFAULT_BASIC_AUTH,
-    help='Default: No authentication. Specify colon separated user:password '
-    'to enable basic authentication.')
 flags.add_argument(
     '--version',
     '-v',
@@ -360,9 +354,11 @@ class Proxy:
     @staticmethod
     def get_default_plugins(
             args: argparse.Namespace) -> List[Tuple[str, bool]]:
-        # Prepare list of plugins to load based upon --enable-* and --disable-*
-        # flags
+        # Prepare list of plugins to load based upon
+        # --enable-*, --disable-* and --basic-auth flags.
         default_plugins: List[Tuple[str, bool]] = []
+        if args.basic_auth is not None:
+            default_plugins.append((PLUGIN_PROXY_AUTH, True))
         if args.enable_dashboard:
             default_plugins.append((PLUGIN_WEB_SERVER, True))
             args.enable_static_server = True


Test patch:
diff --git a/tests/common/test_flags.py b/tests/common/test_flags.py
index 50adb83..c36f8e6 100644
--- a/tests/common/test_flags.py
+++ b/tests/common/test_flags.py
@@ -8,6 +8,8 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
+from proxy.common.utils import bytes_
+from proxy.common.constants import PLUGIN_HTTP_PROXY
 import unittest
 
 from typing import List, Dict
@@ -86,7 +88,7 @@ class TestFlags(unittest.TestCase):
 
     def test_unique_plugin_from_bytes(self) -> None:
         self.flags = Proxy.initialize([], plugins=[
-            b'proxy.http.proxy.HttpProxyPlugin',
+            bytes_(PLUGIN_HTTP_PROXY),
         ])
         self.assert_plugins({'HttpProtocolHandlerPlugin': [
             HttpProxyPlugin,
@@ -94,7 +96,7 @@ class TestFlags(unittest.TestCase):
 
     def test_unique_plugin_from_args(self) -> None:
         self.flags = Proxy.initialize([
-            '--plugins', 'proxy.http.proxy.HttpProxyPlugin',
+            '--plugins', PLUGIN_HTTP_PROXY,
         ])
         self.assert_plugins({'HttpProtocolHandlerPlugin': [
             HttpProxyPlugin,
diff --git a/tests/http/test_protocol_handler.py b/tests/http/test_protocol_handler.py
index a96d43f..a79e6b4 100644
--- a/tests/http/test_protocol_handler.py
+++ b/tests/http/test_protocol_handler.py
@@ -18,7 +18,7 @@ from unittest import mock
 from proxy.proxy import Proxy
 from proxy.common.version import __version__
 from proxy.common.utils import bytes_
-from proxy.common.constants import CRLF
+from proxy.common.constants import CRLF, PLUGIN_HTTP_PROXY, PLUGIN_PROXY_AUTH, PLUGIN_WEB_SERVER
 from proxy.core.connection import TcpClientConnection
 from proxy.http.parser import HttpParser
 from proxy.http.proxy import HttpProxyPlugin
@@ -41,8 +41,8 @@ class TestHttpProtocolHandler(unittest.TestCase):
         self.http_server_port = 65535
         self.flags = Proxy.initialize()
         self.flags.plugins = Proxy.load_plugins([
-            b'proxy.http.proxy.HttpProxyPlugin',
-            b'proxy.http.server.HttpWebServerPlugin',
+            bytes_(PLUGIN_HTTP_PROXY),
+            bytes_(PLUGIN_WEB_SERVER),
         ])
 
         self.mock_selector = mock_selector
@@ -175,8 +175,9 @@ class TestHttpProtocolHandler(unittest.TestCase):
         flags = Proxy.initialize(
             auth_code=base64.b64encode(b'user:pass'))
         flags.plugins = Proxy.load_plugins([
-            b'proxy.http.proxy.HttpProxyPlugin',
-            b'proxy.http.server.HttpWebServerPlugin',
+            bytes_(PLUGIN_HTTP_PROXY),
+            bytes_(PLUGIN_WEB_SERVER),
+            bytes_(PLUGIN_PROXY_AUTH),
         ])
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr), flags=flags)
@@ -208,8 +209,8 @@ class TestHttpProtocolHandler(unittest.TestCase):
         flags = Proxy.initialize(
             auth_code=base64.b64encode(b'user:pass'))
         flags.plugins = Proxy.load_plugins([
-            b'proxy.http.proxy.HttpProxyPlugin',
-            b'proxy.http.server.HttpWebServerPlugin',
+            bytes_(PLUGIN_HTTP_PROXY),
+            bytes_(PLUGIN_WEB_SERVER),
         ])
 
         self.protocol_handler = HttpProtocolHandler(
@@ -257,8 +258,8 @@ class TestHttpProtocolHandler(unittest.TestCase):
         flags = Proxy.initialize(
             auth_code=base64.b64encode(b'user:pass'))
         flags.plugins = Proxy.load_plugins([
-            b'proxy.http.proxy.HttpProxyPlugin',
-            b'proxy.http.server.HttpWebServerPlugin'
+            bytes_(PLUGIN_HTTP_PROXY),
+            bytes_(PLUGIN_WEB_SERVER)
         ])
 
         self.protocol_handler = HttpProtocolHandler(
diff --git a/tests/http/test_web_server.py b/tests/http/test_web_server.py
index 43ff86c..cdb3592 100644
--- a/tests/http/test_web_server.py
+++ b/tests/http/test_web_server.py
@@ -20,7 +20,7 @@ from proxy.core.connection import TcpClientConnection
 from proxy.http.handler import HttpProtocolHandler
 from proxy.http.parser import httpParserStates
 from proxy.common.utils import build_http_response, build_http_request, bytes_, text_
-from proxy.common.constants import CRLF, PROXY_PY_DIR
+from proxy.common.constants import CRLF, PLUGIN_HTTP_PROXY, PLUGIN_PAC_FILE, PLUGIN_WEB_SERVER, PROXY_PY_DIR
 from proxy.http.server import HttpWebServerPlugin
 
 
@@ -35,8 +35,8 @@ class TestWebServerPlugin(unittest.TestCase):
         self.mock_selector = mock_selector
         self.flags = Proxy.initialize()
         self.flags.plugins = Proxy.load_plugins([
-            b'proxy.http.proxy.HttpProxyPlugin',
-            b'proxy.http.server.HttpWebServerPlugin',
+            bytes_(PLUGIN_HTTP_PROXY),
+            bytes_(PLUGIN_WEB_SERVER),
         ])
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr),
@@ -98,8 +98,8 @@ class TestWebServerPlugin(unittest.TestCase):
                 data=None), selectors.EVENT_READ), ]
         flags = Proxy.initialize()
         flags.plugins = Proxy.load_plugins([
-            b'proxy.http.proxy.HttpProxyPlugin',
-            b'proxy.http.server.HttpWebServerPlugin',
+            bytes_(PLUGIN_HTTP_PROXY),
+            bytes_(PLUGIN_WEB_SERVER),
         ])
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr),
@@ -151,8 +151,8 @@ class TestWebServerPlugin(unittest.TestCase):
             enable_static_server=True,
             static_server_dir=static_server_dir)
         flags.plugins = Proxy.load_plugins([
-            b'proxy.http.proxy.HttpProxyPlugin',
-            b'proxy.http.server.HttpWebServerPlugin',
+            bytes_(PLUGIN_HTTP_PROXY),
+            bytes_(PLUGIN_WEB_SERVER),
         ])
 
         self.protocol_handler = HttpProtocolHandler(
@@ -201,8 +201,8 @@ class TestWebServerPlugin(unittest.TestCase):
 
         flags = Proxy.initialize(enable_static_server=True)
         flags.plugins = Proxy.load_plugins([
-            b'proxy.http.proxy.HttpProxyPlugin',
-            b'proxy.http.server.HttpWebServerPlugin',
+            bytes_(PLUGIN_HTTP_PROXY),
+            bytes_(PLUGIN_WEB_SERVER),
         ])
 
         self.protocol_handler = HttpProtocolHandler(
@@ -239,9 +239,9 @@ class TestWebServerPlugin(unittest.TestCase):
     def init_and_make_pac_file_request(self, pac_file: str) -> None:
         flags = Proxy.initialize(pac_file=pac_file)
         flags.plugins = Proxy.load_plugins([
-            b'proxy.http.proxy.HttpProxyPlugin',
-            b'proxy.http.server.HttpWebServerPlugin',
-            b'proxy.http.server.HttpWebServerPacFilePlugin',
+            bytes_(PLUGIN_HTTP_PROXY),
+            bytes_(PLUGIN_WEB_SERVER),
+            bytes_(PLUGIN_PAC_FILE),
         ])
         self.protocol_handler = HttpProtocolHandler(
             TcpClientConnection(self._conn, self._addr),