Response Packet Utilities (#903)

* Add response pkt utility

* Unused import

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix tests as some content is now by default gzipped based upon min compression config

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused

* Update necessary tests to use `okResponse` utility

* Add option to explicitly disable compression

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: abhinavsingh

.github/workflows/test-library.yml

@@ -247,8 +247,10 @@ jobs:
         path: ${{ steps.pip-cache.outputs.dir }}
         key: >-
           ${{ runner.os }}-pip-${{
-          steps.calc-cache-key-py.outputs.py-hash-key }}-${{
-          hashFiles('tox.ini') }}
+            steps.calc-cache-key-py.outputs.py-hash-key
+          }}-${{
+            hashFiles('tox.ini')
+          }}
         restore-keys: |
           ${{ runner.os }}-pip-${{
               steps.calc-cache-key-py.outputs.py-hash-key
@@ -370,8 +372,10 @@ jobs:
         path: ${{ steps.pip-cache.outputs.dir }}
         key: >-
           ${{ runner.os }}-pip-${{
-          steps.calc-cache-key-py.outputs.py-hash-key }}-${{
-          hashFiles('tox.ini') }}
+            steps.calc-cache-key-py.outputs.py-hash-key
+          }}-${{
+            hashFiles('tox.ini')
+          }}
         restore-keys: |
           ${{ runner.os }}-pip-${{
               steps.calc-cache-key-py.outputs.py-hash-key
@@ -701,6 +705,8 @@ jobs:
     steps:
     - name: Checkout
       uses: actions/checkout@v2
+      with:
+        ref: ${{ github.event.inputs.release-commitish }}
     - name: Download all the dists
       uses: actions/download-artifact@v2
       with:

examples/https_connect_tunnel.py

@@ -13,29 +13,18 @@
 from typing import Any, Optional
 
 from proxy import Proxy
-from proxy.common.utils import build_http_response
-from proxy.http import httpStatusCodes
+
+from proxy.http.responses import (
+    PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT,
+    PROXY_TUNNEL_UNSUPPORTED_SCHEME,
+)
+
 from proxy.core.base import BaseTcpTunnelHandler
 
 
 class HttpsConnectTunnelHandler(BaseTcpTunnelHandler):
     """A https CONNECT tunnel."""
 
-    PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT = memoryview(
-        build_http_response(
-            httpStatusCodes.OK,
-            reason=b'Connection established',
-        ),
-    )
-
-    PROXY_TUNNEL_UNSUPPORTED_SCHEME = memoryview(
-        build_http_response(
-            httpStatusCodes.BAD_REQUEST,
-            reason=b'Unsupported protocol scheme',
-            conn_close=True,
-        ),
-    )
-
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
 
@@ -50,9 +39,7 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
 
         # Drop the request if not a CONNECT request
         if not self.request.is_https_tunnel:
-            self.work.queue(
-                HttpsConnectTunnelHandler.PROXY_TUNNEL_UNSUPPORTED_SCHEME,
-            )
+            self.work.queue(PROXY_TUNNEL_UNSUPPORTED_SCHEME)
             return True
 
         # CONNECT requests are short and we need not worry about
@@ -63,9 +50,7 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
         self.connect_upstream()
 
         # Queue tunnel established response to client
-        self.work.queue(
-            HttpsConnectTunnelHandler.PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT,
-        )
+        self.work.queue(PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT)
 
         return None
 

proxy/dashboard/dashboard.py

@@ -15,9 +15,9 @@
 
 from .plugin import ProxyDashboardWebsocketPlugin
 
-from ..common.utils import build_http_response, bytes_
+from ..common.utils import bytes_
 
-from ..http import httpStatusCodes
+from ..http.responses import permanentRedirectResponse
 from ..http.parser import HttpParser
 from ..http.websocket import WebsocketFrame
 from ..http.server import HttpWebServerPlugin, HttpWebServerBasePlugin, httpProtocolTypes
@@ -69,25 +69,13 @@ def handle_request(self, request: HttpParser) -> None:
                         self.flags.static_server_dir,
                         'dashboard', 'proxy.html',
                     ),
-                    self.flags.min_compression_limit,
                 ),
             )
         elif request.path in (
                 b'/dashboard',
                 b'/dashboard/proxy.html',
         ):
-            self.client.queue(
-                memoryview(
-                    build_http_response(
-                        httpStatusCodes.PERMANENT_REDIRECT, reason=b'Permanent Redirect',
-                        headers={
-                            b'Location': b'/dashboard/',
-                            b'Content-Length': b'0',
-                        },
-                        conn_close=True,
-                    ),
-                ),
-            )
+            self.client.queue(permanentRedirectResponse(b'/dashboard/'))
 
     def on_websocket_open(self) -> None:
         logger.info('app ws opened')

proxy/http/exception/proxy_auth_failed.py

@@ -17,10 +17,7 @@
 
 from .base import HttpProtocolException
 
-from ..codes import httpStatusCodes
-
-from ...common.constants import PROXY_AGENT_HEADER_VALUE, PROXY_AGENT_HEADER_KEY
-from ...common.utils import build_http_response
+from ..responses import PROXY_AUTH_FAILED_RESPONSE_PKT
 
 if TYPE_CHECKING:
     from ..parser import HttpParser
@@ -30,21 +27,8 @@ class ProxyAuthenticationFailed(HttpProtocolException):
     """Exception raised when HTTP Proxy auth is enabled and
     incoming request doesn't present necessary credentials."""
 
-    RESPONSE_PKT = memoryview(
-        build_http_response(
-            httpStatusCodes.PROXY_AUTH_REQUIRED,
-            reason=b'Proxy Authentication Required',
-            headers={
-                PROXY_AGENT_HEADER_KEY: PROXY_AGENT_HEADER_VALUE,
-                b'Proxy-Authenticate': b'Basic',
-            },
-            body=b'Proxy Authentication Required',
-            conn_close=True,
-        ),
-    )
-
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(self.__class__.__name__, **kwargs)
 
     def response(self, _request: 'HttpParser') -> memoryview:
-        return self.RESPONSE_PKT
+        return PROXY_AUTH_FAILED_RESPONSE_PKT

proxy/http/exception/proxy_conn_failed.py

@@ -16,10 +16,7 @@
 
 from .base import HttpProtocolException
 
-from ..codes import httpStatusCodes
-
-from ...common.constants import PROXY_AGENT_HEADER_VALUE, PROXY_AGENT_HEADER_KEY
-from ...common.utils import build_http_response
+from ..responses import BAD_GATEWAY_RESPONSE_PKT
 
 if TYPE_CHECKING:
     from ..parser import HttpParser
@@ -28,23 +25,11 @@
 class ProxyConnectionFailed(HttpProtocolException):
     """Exception raised when ``HttpProxyPlugin`` is unable to establish connection to upstream server."""
 
-    RESPONSE_PKT = memoryview(
-        build_http_response(
-            httpStatusCodes.BAD_GATEWAY,
-            reason=b'Bad Gateway',
-            headers={
-                PROXY_AGENT_HEADER_KEY: PROXY_AGENT_HEADER_VALUE,
-            },
-            body=b'Bad Gateway',
-            conn_close=True,
-        ),
-    )
-
     def __init__(self, host: str, port: int, reason: str, **kwargs: Any):
         self.host: str = host
         self.port: int = port
         self.reason: str = reason
         super().__init__('%s %s' % (self.__class__.__name__, reason), **kwargs)
 
     def response(self, _request: 'HttpParser') -> memoryview:
-        return self.RESPONSE_PKT
+        return BAD_GATEWAY_RESPONSE_PKT

proxy/http/proxy/server.py

@@ -28,10 +28,10 @@
 
 from ..headers import httpHeaders
 from ..methods import httpMethods
-from ..codes import httpStatusCodes
 from ..plugin import HttpProtocolHandlerPlugin
 from ..exception import HttpProtocolException, ProxyConnectionFailed
 from ..parser import HttpParser, httpParserStates, httpParserTypes
+from ..responses import PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT
 
 from ...common.types import Readables, Writables
 from ...common.constants import DEFAULT_CA_CERT_DIR, DEFAULT_CA_CERT_FILE, DEFAULT_CA_FILE
@@ -40,7 +40,7 @@
 from ...common.constants import PROXY_AGENT_HEADER_VALUE, DEFAULT_DISABLE_HEADERS
 from ...common.constants import DEFAULT_HTTP_ACCESS_LOG_FORMAT, DEFAULT_HTTPS_ACCESS_LOG_FORMAT
 from ...common.constants import DEFAULT_DISABLE_HTTP_PROXY, PLUGIN_PROXY_AUTH
-from ...common.utils import build_http_response, text_
+from ...common.utils import text_
 from ...common.pki import gen_public_key, gen_csr, sign_csr
 
 from ...core.event import eventNames
@@ -136,13 +136,6 @@
 class HttpProxyPlugin(HttpProtocolHandlerPlugin):
     """HttpProtocolHandler plugin which implements HttpProxy specifications."""
 
-    PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT = memoryview(
-        build_http_response(
-            httpStatusCodes.OK,
-            reason=b'Connection established',
-        ),
-    )
-
     # Used to synchronization during certificate generation and
     # connection pool operations.
     lock = threading.Lock()
@@ -546,9 +539,7 @@ def on_request_complete(self) -> Union[socket.socket, bool]:
         # Optionally, setup interceptor if TLS interception is enabled.
         if self.upstream:
             if self.request.is_https_tunnel:
-                self.client.queue(
-                    HttpProxyPlugin.PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT,
-                )
+                self.client.queue(PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT)
                 if self.tls_interception_enabled():
                     return self.intercept()
             # If an upstream server connection was established for http request,

proxy/http/responses.py

@@ -0,0 +1,138 @@
+# -*- 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 gzip
+from typing import Any, Dict, Optional
+
+from ..common.flag import flags
+from ..common.utils import build_http_response
+from ..common.constants import PROXY_AGENT_HEADER_VALUE, PROXY_AGENT_HEADER_KEY
+
+from .codes import httpStatusCodes
+
+
+PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT = memoryview(
+    build_http_response(
+        httpStatusCodes.OK,
+        reason=b'Connection established',
+    ),
+)
+
+PROXY_TUNNEL_UNSUPPORTED_SCHEME = memoryview(
+    build_http_response(
+        httpStatusCodes.BAD_REQUEST,
+        reason=b'Unsupported protocol scheme',
+        conn_close=True,
+    ),
+)
+
+PROXY_AUTH_FAILED_RESPONSE_PKT = memoryview(
+    build_http_response(
+        httpStatusCodes.PROXY_AUTH_REQUIRED,
+        reason=b'Proxy Authentication Required',
+        headers={
+            PROXY_AGENT_HEADER_KEY: PROXY_AGENT_HEADER_VALUE,
+            b'Proxy-Authenticate': b'Basic',
+        },
+        body=b'Proxy Authentication Required',
+        conn_close=True,
+    ),
+)
+
+NOT_FOUND_RESPONSE_PKT = memoryview(
+    build_http_response(
+        httpStatusCodes.NOT_FOUND,
+        reason=b'NOT FOUND',
+        headers={
+            b'Server': PROXY_AGENT_HEADER_VALUE,
+            b'Content-Length': b'0',
+        },
+        conn_close=True,
+    ),
+)
+
+NOT_IMPLEMENTED_RESPONSE_PKT = memoryview(
+    build_http_response(
+        httpStatusCodes.NOT_IMPLEMENTED,
+        reason=b'NOT IMPLEMENTED',
+        headers={
+            b'Server': PROXY_AGENT_HEADER_VALUE,
+            b'Content-Length': b'0',
+        },
+        conn_close=True,
+    ),
+)
+
+BAD_GATEWAY_RESPONSE_PKT = memoryview(
+    build_http_response(
+        httpStatusCodes.BAD_GATEWAY,
+        reason=b'Bad Gateway',
+        headers={
+            PROXY_AGENT_HEADER_KEY: PROXY_AGENT_HEADER_VALUE,
+        },
+        body=b'Bad Gateway',
+        conn_close=True,
+    ),
+)
+
+
+def okResponse(
+        content: Optional[bytes] = None,
+        headers: Optional[Dict[bytes, bytes]] = None,
+        compress: bool = True,
+        **kwargs: Any,
+) -> memoryview:
+    do_compress: bool = False
+    if compress and flags.args and content and len(content) > flags.args.min_compression_limit:
+        do_compress = True
+        if not headers:
+            headers = {}
+        headers.update({
+            b'Content-Encoding': b'gzip',
+        })
+    return memoryview(
+        build_http_response(
+            200,
+            reason=b'OK',
+            headers=headers,
+            body=gzip.compress(content)
+            if do_compress and content
+            else content,
+            **kwargs,
+        ),
+    )
+
+
+def permanentRedirectResponse(location: bytes) -> memoryview:
+    return memoryview(
+        build_http_response(
+            httpStatusCodes.PERMANENT_REDIRECT,
+            reason=b'Permanent Redirect',
+            headers={
+                b'Location': location,
+                b'Content-Length': b'0',
+            },
+            conn_close=True,
+        ),
+    )
+
+
+def seeOthersResponse(location: bytes) -> memoryview:
+    return memoryview(
+        build_http_response(
+            httpStatusCodes.SEE_OTHER,
+            reason=b'See Other',
+            headers={
+                b'Location': location,
+                b'Content-Length': b'0',
+            },
+            conn_close=True,
+        ),
+    )