diff --git a/.github/workflows/test-library.yml b/.github/workflows/test-library.yml index c344b67610..e6fc703fa9 100644 --- a/.github/workflows/test-library.yml +++ b/.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: diff --git a/examples/https_connect_tunnel.py b/examples/https_connect_tunnel.py index a20a4b3433..21ef67e3e3 100644 --- a/examples/https_connect_tunnel.py +++ b/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 diff --git a/proxy/dashboard/dashboard.py b/proxy/dashboard/dashboard.py index cfdb5c8c65..196ff16ac8 100644 --- a/proxy/dashboard/dashboard.py +++ b/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') diff --git a/proxy/http/exception/proxy_auth_failed.py b/proxy/http/exception/proxy_auth_failed.py index 60782e0c23..76bdcf227c 100644 --- a/proxy/http/exception/proxy_auth_failed.py +++ b/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 diff --git a/proxy/http/exception/proxy_conn_failed.py b/proxy/http/exception/proxy_conn_failed.py index d0eae48192..d6af42131e 100644 --- a/proxy/http/exception/proxy_conn_failed.py +++ b/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,18 +25,6 @@ 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 @@ -47,4 +32,4 @@ def __init__(self, host: str, port: int, reason: str, **kwargs: Any): super().__init__('%s %s' % (self.__class__.__name__, reason), **kwargs) def response(self, _request: 'HttpParser') -> memoryview: - return self.RESPONSE_PKT + return BAD_GATEWAY_RESPONSE_PKT diff --git a/proxy/http/proxy/server.py b/proxy/http/proxy/server.py index 462cf0fc83..7c2c41083c 100644 --- a/proxy/http/proxy/server.py +++ b/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, diff --git a/proxy/http/responses.py b/proxy/http/responses.py new file mode 100644 index 0000000000..9e865cd274 --- /dev/null +++ b/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, + ), + ) diff --git a/proxy/http/server/pac_plugin.py b/proxy/http/server/pac_plugin.py index f6858cf452..2e6e233872 100644 --- a/proxy/http/server/pac_plugin.py +++ b/proxy/http/server/pac_plugin.py @@ -12,16 +12,15 @@ pac """ -import gzip - from typing import List, Tuple, Optional, Any from .plugin import HttpWebServerBasePlugin from .protocols import httpProtocolTypes from ..parser import HttpParser +from ..responses import okResponse -from ...common.utils import bytes_, text_, build_http_response +from ...common.utils import bytes_, text_ from ...common.flag import flags from ...common.constants import DEFAULT_PAC_FILE, DEFAULT_PAC_FILE_URL_PATH @@ -69,11 +68,11 @@ def cache_pac_file_response(self) -> None: content = f.read() except IOError: content = bytes_(self.flags.pac_file) - self.pac_file_response = memoryview( - build_http_response( - 200, reason=b'OK', headers={ - b'Content-Type': b'application/x-ns-proxy-autoconfig', - b'Content-Encoding': b'gzip', - }, body=gzip.compress(content), - ), + self.pac_file_response = okResponse( + content=content, + headers={ + b'Content-Type': b'application/x-ns-proxy-autoconfig', + b'Content-Encoding': b'gzip', + }, + conn_close=True, ) diff --git a/proxy/http/server/web.py b/proxy/http/server/web.py index 520f170d53..89ed43c618 100644 --- a/proxy/http/server/web.py +++ b/proxy/http/server/web.py @@ -9,7 +9,6 @@ :license: BSD, see LICENSE for more details. """ import re -import gzip import time import socket import logging @@ -17,18 +16,18 @@ from typing import List, Tuple, Optional, Dict, Union, Any, Pattern -from ...common.constants import DEFAULT_STATIC_SERVER_DIR, PROXY_AGENT_HEADER_VALUE +from ...common.constants import DEFAULT_STATIC_SERVER_DIR from ...common.constants import DEFAULT_ENABLE_STATIC_SERVER, DEFAULT_ENABLE_WEB_SERVER from ...common.constants import DEFAULT_MIN_COMPRESSION_LIMIT, DEFAULT_WEB_ACCESS_LOG_FORMAT -from ...common.utils import bytes_, text_, build_http_response, build_websocket_handshake_response +from ...common.utils import bytes_, text_, build_websocket_handshake_response from ...common.types import Readables, Writables from ...common.flag import flags -from ..codes import httpStatusCodes from ..exception import HttpProtocolException from ..plugin import HttpProtocolHandlerPlugin from ..websocket import WebsocketFrame, websocketOpcodes from ..parser import HttpParser, httpParserTypes +from ..responses import NOT_FOUND_RESPONSE_PKT, NOT_IMPLEMENTED_RESPONSE_PKT, okResponse from .plugin import HttpWebServerBasePlugin from .protocols import httpProtocolTypes @@ -74,30 +73,6 @@ class HttpWebServerPlugin(HttpProtocolHandlerPlugin): """HttpProtocolHandler plugin which handles incoming requests to local web server.""" - DEFAULT_404_RESPONSE = 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, - ), - ) - - DEFAULT_501_RESPONSE = 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, - ), - ) - def __init__( self, *args: Any, **kwargs: Any, @@ -137,7 +112,7 @@ def encryption_enabled(self) -> bool: self.flags.certfile is not None @staticmethod - def read_and_build_static_file_response(path: str, min_compression_limit: int) -> memoryview: + def read_and_build_static_file_response(path: str) -> memoryview: try: with open(path, 'rb') as f: content = f.read() @@ -148,22 +123,14 @@ def read_and_build_static_file_response(path: str, min_compression_limit: int) - b'Content-Type': bytes_(content_type), b'Cache-Control': b'max-age=86400', } - do_compress = len(content) > min_compression_limit - if do_compress: - headers.update({ - b'Content-Encoding': b'gzip', - }) - return memoryview( - build_http_response( - httpStatusCodes.OK, - reason=b'OK', - headers=headers, - body=gzip.compress(content) if do_compress else content, - conn_close=True, - ), + return okResponse( + content=content, + headers=headers, + # TODO: Should we really close or take advantage of keep-alive? + conn_close=True, ) except FileNotFoundError: - return HttpWebServerPlugin.DEFAULT_404_RESPONSE + return NOT_FOUND_RESPONSE_PKT def try_upgrade(self) -> bool: if self.request.has_header(b'connection') and \ @@ -181,7 +148,7 @@ def try_upgrade(self) -> bool: ) self.switched_protocol = httpProtocolTypes.WEBSOCKET else: - self.client.queue(self.DEFAULT_501_RESPONSE) + self.client.queue(NOT_IMPLEMENTED_RESPONSE_PKT) return True return False @@ -223,12 +190,11 @@ def on_request_complete(self) -> Union[socket.socket, bool]: self.client.queue( self.read_and_build_static_file_response( self.flags.static_server_dir + path, - self.flags.min_compression_limit, ), ) return True # Catch all unhandled web server requests, return 404 - self.client.queue(self.DEFAULT_404_RESPONSE) + self.client.queue(NOT_FOUND_RESPONSE_PKT) return True def get_descriptors(self) -> Tuple[List[int], List[int]]: diff --git a/proxy/plugin/man_in_the_middle.py b/proxy/plugin/man_in_the_middle.py index 970547e6f2..a6d6220dfd 100644 --- a/proxy/plugin/man_in_the_middle.py +++ b/proxy/plugin/man_in_the_middle.py @@ -8,8 +8,7 @@ :copyright: (c) 2013-present by Abhinav Singh and contributors. :license: BSD, see LICENSE for more details. """ -from ..common.utils import build_http_response -from ..http import httpStatusCodes +from ..http.responses import okResponse from ..http.proxy import HttpProxyBasePlugin @@ -17,10 +16,4 @@ class ManInTheMiddlePlugin(HttpProxyBasePlugin): """Modifies upstream server responses.""" def handle_upstream_chunk(self, chunk: memoryview) -> memoryview: - return memoryview( - build_http_response( - httpStatusCodes.OK, - reason=b'OK', - body=b'Hello from man in the middle', - ), - ) + return okResponse(content=b'Hello from man in the middle') diff --git a/proxy/plugin/mock_rest_api.py b/proxy/plugin/mock_rest_api.py index 12e3543deb..56cac4cb27 100644 --- a/proxy/plugin/mock_rest_api.py +++ b/proxy/plugin/mock_rest_api.py @@ -15,9 +15,9 @@ import json from typing import Optional -from ..common.utils import bytes_, build_http_response, text_ +from ..common.utils import bytes_, text_ -from ..http import httpStatusCodes +from ..http.responses import okResponse, NOT_FOUND_RESPONSE_PKT from ..http.parser import HttpParser from ..http.proxy import HttpProxyBasePlugin @@ -75,26 +75,15 @@ def handle_client_request( assert request.path if request.path in self.REST_API_SPEC: self.client.queue( - memoryview( - build_http_response( - httpStatusCodes.OK, - reason=b'OK', - headers={b'Content-Type': b'application/json'}, - body=bytes_( - json.dumps( - self.REST_API_SPEC[request.path], - ), + okResponse( + headers={b'Content-Type': b'application/json'}, + content=bytes_( + json.dumps( + self.REST_API_SPEC[request.path], ), ), ), ) else: - self.client.queue( - memoryview( - build_http_response( - httpStatusCodes.NOT_FOUND, - reason=b'NOT FOUND', body=b'Not Found', - ), - ), - ) + self.client.queue(NOT_FOUND_RESPONSE_PKT) return None diff --git a/proxy/plugin/shortlink.py b/proxy/plugin/shortlink.py index 43e2534ea6..11b7a7d823 100644 --- a/proxy/plugin/shortlink.py +++ b/proxy/plugin/shortlink.py @@ -15,9 +15,8 @@ from typing import Optional from ..common.constants import DOT, SLASH -from ..common.utils import build_http_response -from ..http import httpStatusCodes +from ..http.responses import NOT_FOUND_RESPONSE_PKT, seeOthersResponse from ..http.parser import HttpParser from ..http.proxy import HttpProxyBasePlugin @@ -66,28 +65,11 @@ def handle_client_request( if request.host in self.SHORT_LINKS: path = SLASH if not request.path else request.path self.client.queue( - memoryview( - build_http_response( - httpStatusCodes.SEE_OTHER, reason=b'See Other', - headers={ - b'Location': b'http://' + self.SHORT_LINKS[request.host] + path, - b'Content-Length': b'0', - }, - conn_close=True, - ), + seeOthersResponse( + b'http://' + self.SHORT_LINKS[request.host] + path, ), ) else: - self.client.queue( - memoryview( - build_http_response( - httpStatusCodes.NOT_FOUND, reason=b'NOT FOUND', - headers={ - b'Content-Length': b'0', - }, - conn_close=True, - ), - ), - ) + self.client.queue(NOT_FOUND_RESPONSE_PKT) return None return request diff --git a/proxy/plugin/web_server_route.py b/proxy/plugin/web_server_route.py index bfbe4bd345..fe5e95707c 100644 --- a/proxy/plugin/web_server_route.py +++ b/proxy/plugin/web_server_route.py @@ -11,27 +11,16 @@ import logging from typing import List, Tuple -from ..common.utils import build_http_response - -from ..http import httpStatusCodes +from ..http.responses import okResponse from ..http.parser import HttpParser from ..http.websocket import WebsocketFrame from ..http.server import HttpWebServerBasePlugin, httpProtocolTypes -HTTP_RESPONSE = memoryview( - build_http_response( - httpStatusCodes.OK, body=b'HTTP route response', - ), -) - -HTTPS_RESPONSE = memoryview( - build_http_response( - httpStatusCodes.OK, body=b'HTTPS route response', - ), -) - logger = logging.getLogger(__name__) +HTTP_RESPONSE = okResponse(content=b'HTTP route response') +HTTPS_RESPONSE = okResponse(content=b'HTTP route response') + class WebServerPlugin(HttpWebServerBasePlugin): """Demonstrates inbuilt web server routing using plugin.""" diff --git a/tests/http/exceptions/test_http_proxy_auth_failed.py b/tests/http/exceptions/test_http_proxy_auth_failed.py index fdf1b9d6da..23a3b4dcd5 100644 --- a/tests/http/exceptions/test_http_proxy_auth_failed.py +++ b/tests/http/exceptions/test_http_proxy_auth_failed.py @@ -14,8 +14,8 @@ from pytest_mock import MockerFixture from proxy.common.flag import FlagParser -from proxy.http.exception.proxy_auth_failed import ProxyAuthenticationFailed from proxy.http import HttpProtocolHandler, httpHeaders +from proxy.http.responses import PROXY_AUTH_FAILED_RESPONSE_PKT from proxy.core.connection import TcpClientConnection from proxy.common.utils import build_http_request @@ -67,7 +67,8 @@ async def test_proxy_auth_fails_without_cred(self) -> None: self.mock_server_conn.assert_not_called() self.assertEqual(self.protocol_handler.work.has_buffer(), True) self.assertEqual( - self.protocol_handler.work.buffer[0], ProxyAuthenticationFailed.RESPONSE_PKT, + self.protocol_handler.work.buffer[0], + PROXY_AUTH_FAILED_RESPONSE_PKT, ) self._conn.send.assert_not_called() @@ -95,7 +96,8 @@ async def test_proxy_auth_fails_with_invalid_cred(self) -> None: self.mock_server_conn.assert_not_called() self.assertEqual(self.protocol_handler.work.has_buffer(), True) self.assertEqual( - self.protocol_handler.work.buffer[0], ProxyAuthenticationFailed.RESPONSE_PKT, + self.protocol_handler.work.buffer[0], + PROXY_AUTH_FAILED_RESPONSE_PKT, ) self._conn.send.assert_not_called() diff --git a/tests/http/test_http_parser.py b/tests/http/test_http_parser.py index 5f771bd708..0489b85741 100644 --- a/tests/http/test_http_parser.py +++ b/tests/http/test_http_parser.py @@ -11,12 +11,13 @@ import unittest from proxy.common.constants import CRLF, HTTP_1_0 -from proxy.common.utils import build_http_request, build_http_response, build_http_header +from proxy.common.utils import build_http_request, build_http_header from proxy.common.utils import find_http_line, bytes_ -from proxy.http import httpStatusCodes, httpMethods +from proxy.http import httpMethods from proxy.http.exception import HttpProtocolException from proxy.http.parser import HttpParser, httpParserTypes, httpParserStates +from proxy.http.responses import okResponse class TestHttpParser(unittest.TestCase): @@ -141,18 +142,16 @@ def test_build_request(self) -> None: def test_build_response(self) -> None: self.assertEqual( - build_http_response( - 200, reason=b'OK', protocol_version=b'HTTP/1.1', - ), + okResponse(protocol_version=b'HTTP/1.1'), CRLF.join([ b'HTTP/1.1 200 OK', CRLF, ]), ) self.assertEqual( - build_http_response( - 200, reason=b'OK', protocol_version=b'HTTP/1.1', + okResponse( headers={b'key': b'value'}, + protocol_version=b'HTTP/1.1', ), CRLF.join([ b'HTTP/1.1 200 OK', @@ -164,10 +163,10 @@ def test_build_response(self) -> None: def test_build_response_adds_content_length_header(self) -> None: body = b'Hello world!!!' self.assertEqual( - build_http_response( - 200, reason=b'OK', protocol_version=b'HTTP/1.1', + okResponse( headers={b'key': b'value'}, - body=body, + content=body, + protocol_version=b'HTTP/1.1', ), CRLF.join([ b'HTTP/1.1 200 OK', @@ -610,29 +609,30 @@ def test_chunked_response_parse(self) -> None: self.assertEqual(self.parser.state, httpParserStates.COMPLETE) def test_pipelined_response_parse(self) -> None: - response = build_http_response( - httpStatusCodes.OK, reason=b'OK', - headers={ - b'Content-Length': b'15', - }, - body=b'{"key":"value"}', + self.assert_pipeline_response( + okResponse( + headers={ + b'Content-Length': b'15', + }, + content=b'{"key":"value"}', + ), ) - self.assert_pipeline_response(response) def test_pipelined_chunked_response_parse(self) -> None: - response = build_http_response( - httpStatusCodes.OK, reason=b'OK', - headers={ - b'Transfer-Encoding': b'chunked', - b'Content-Type': b'application/json', - }, - body=b'f\r\n{"key":"value"}\r\n0\r\n\r\n', + self.assert_pipeline_response( + okResponse( + headers={ + b'Transfer-Encoding': b'chunked', + b'Content-Type': b'application/json', + }, + content=b'f\r\n{"key":"value"}\r\n0\r\n\r\n', + compress=False, + ), ) - self.assert_pipeline_response(response) - def assert_pipeline_response(self, response: bytes) -> None: + def assert_pipeline_response(self, response: memoryview) -> None: self.parser = HttpParser(httpParserTypes.RESPONSE_PARSER) - self.parser.parse(response + response) + self.parser.parse(response.tobytes() + response.tobytes()) self.assertEqual(self.parser.state, httpParserStates.COMPLETE) self.assertEqual(self.parser.body, b'{"key":"value"}') self.assertEqual(self.parser.buffer, response) diff --git a/tests/http/test_http_proxy_tls_interception.py b/tests/http/test_http_proxy_tls_interception.py index f569afbdc3..23ae05faa9 100644 --- a/tests/http/test_http_proxy_tls_interception.py +++ b/tests/http/test_http_proxy_tls_interception.py @@ -19,11 +19,12 @@ from unittest import mock from proxy.common.constants import DEFAULT_CA_FILE -from proxy.core.connection import TcpClientConnection, TcpServerConnection -from proxy.http import HttpProtocolHandler, httpMethods -from proxy.http.proxy import HttpProxyPlugin from proxy.common.utils import build_http_request, bytes_ from proxy.common.flag import FlagParser +from proxy.http import HttpProtocolHandler, httpMethods +from proxy.http.proxy import HttpProxyPlugin +from proxy.http.responses import PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT +from proxy.core.connection import TcpClientConnection, TcpServerConnection from ..test_assertions import Assertions @@ -178,7 +179,7 @@ async def asyncReturnBool(val: bool) -> bool: ssl_connection, ) self._conn.send.assert_called_with( - HttpProxyPlugin.PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, + PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, ) assert self.flags.ca_cert_dir is not None self.mock_ssl_wrap.assert_called_with( diff --git a/tests/http/test_protocol_handler.py b/tests/http/test_protocol_handler.py index c5f1ff1163..766208c959 100644 --- a/tests/http/test_protocol_handler.py +++ b/tests/http/test_protocol_handler.py @@ -24,8 +24,12 @@ from proxy.core.connection import TcpClientConnection from proxy.http.parser import HttpParser from proxy.http.proxy import HttpProxyPlugin +from proxy.http.responses import ( + BAD_GATEWAY_RESPONSE_PKT, + PROXY_AUTH_FAILED_RESPONSE_PKT, + PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, +) from proxy.http.parser import httpParserStates, httpParserTypes -from proxy.http.exception import ProxyAuthenticationFailed, ProxyConnectionFailed from proxy.http import HttpProtocolHandler, httpHeaders from ..test_assertions import Assertions @@ -80,7 +84,7 @@ async def test_proxy_connection_failed(self) -> None: await self.protocol_handler._run_once() self.assertEqual( self.protocol_handler.work.buffer[0], - ProxyConnectionFailed.RESPONSE_PKT, + BAD_GATEWAY_RESPONSE_PKT, ) @pytest.mark.asyncio # type: ignore[misc] @@ -108,7 +112,7 @@ async def test_proxy_authentication_failed(self) -> None: await self.protocol_handler._run_once() self.assertEqual( self.protocol_handler.work.buffer[0], - ProxyAuthenticationFailed.RESPONSE_PKT, + PROXY_AUTH_FAILED_RESPONSE_PKT, ) @@ -191,7 +195,7 @@ async def assert_tunnel_response( ) self.assertEqual( self.protocol_handler.work.buffer[0], - HttpProxyPlugin.PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, + PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, ) self.mock_server_connection.assert_called_once() server.connect.assert_called_once() @@ -432,7 +436,7 @@ async def assert_data_queued_to_server(self, server: mock.Mock) -> None: assert self.http_server_port is not None self.assertEqual( self._conn.send.call_args[0][0], - HttpProxyPlugin.PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, + PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, ) pkt = CRLF.join([ diff --git a/tests/http/test_web_server.py b/tests/http/test_web_server.py index 3a59fa63af..96083356f8 100644 --- a/tests/http/test_web_server.py +++ b/tests/http/test_web_server.py @@ -25,7 +25,7 @@ from proxy.http.parser import HttpParser, httpParserStates, httpParserTypes from proxy.common.utils import build_http_response, build_http_request, bytes_ from proxy.common.constants import CRLF, PLUGIN_HTTP_PROXY, PLUGIN_PAC_FILE, PLUGIN_WEB_SERVER, PROXY_PY_DIR -from proxy.http.server import HttpWebServerPlugin +from proxy.http.responses import NOT_FOUND_RESPONSE_PKT from ..test_assertions import Assertions @@ -309,7 +309,7 @@ async def test_static_web_server_serves_404(self) -> None: self.assertEqual(self._conn.send.call_count, 1) self.assertEqual( self._conn.send.call_args[0][0], - HttpWebServerPlugin.DEFAULT_404_RESPONSE, + NOT_FOUND_RESPONSE_PKT, ) @@ -368,5 +368,5 @@ async def test_default_web_server_returns_404(self) -> None: ) self.assertEqual( self.protocol_handler.work.buffer[0], - HttpWebServerPlugin.DEFAULT_404_RESPONSE, + NOT_FOUND_RESPONSE_PKT, ) diff --git a/tests/plugin/test_http_proxy_plugins.py b/tests/plugin/test_http_proxy_plugins.py index 03a32527ed..50d2055888 100644 --- a/tests/plugin/test_http_proxy_plugins.py +++ b/tests/plugin/test_http_proxy_plugins.py @@ -8,6 +8,7 @@ :copyright: (c) 2013-present by Abhinav Singh and contributors. :license: BSD, see LICENSE for more details. """ +import gzip import json import pytest import selectors @@ -23,6 +24,7 @@ from proxy.http import HttpProtocolHandler from proxy.http import httpStatusCodes from proxy.http.proxy import HttpProxyPlugin +from proxy.http.parser import HttpParser, httpParserTypes from proxy.common.utils import build_http_request, bytes_, build_http_response from proxy.common.constants import PROXY_AGENT_HEADER_VALUE, DEFAULT_HTTP_PORT from proxy.plugin import ProposedRestApiPlugin, RedirectToCustomServerPlugin @@ -151,15 +153,22 @@ async def test_proposed_rest_api_plugin(self) -> None: await self.protocol_handler._run_once() self.mock_server_conn.assert_not_called() + response = HttpParser(httpParserTypes.RESPONSE_PARSER) + response.parse(self.protocol_handler.work.buffer[0].tobytes()) + assert response.body self.assertEqual( - self.protocol_handler.work.buffer[0].tobytes(), - build_http_response( - httpStatusCodes.OK, reason=b'OK', - headers={b'Content-Type': b'application/json'}, - body=bytes_( - json.dumps( - ProposedRestApiPlugin.REST_API_SPEC[path], - ), + response.header(b'content-type'), + b'application/json', + ) + self.assertEqual( + response.header(b'content-encoding'), + b'gzip', + ) + self.assertEqual( + gzip.decompress(response.body), + bytes_( + json.dumps( + ProposedRestApiPlugin.REST_API_SPEC[path], ), ), ) @@ -329,15 +338,16 @@ def closed() -> bool: server.recv.return_value = \ build_http_response( httpStatusCodes.OK, - reason=b'OK', body=b'Original Response From Upstream', + reason=b'OK', + body=b'Original Response From Upstream', ) await self.protocol_handler._run_once() + response = HttpParser(httpParserTypes.RESPONSE_PARSER) + response.parse(self.protocol_handler.work.buffer[0].tobytes()) + assert response.body self.assertEqual( - self.protocol_handler.work.buffer[0].tobytes(), - build_http_response( - httpStatusCodes.OK, - reason=b'OK', body=b'Hello from man in the middle', - ), + gzip.decompress(response.body), + b'Hello from man in the middle', ) @pytest.mark.asyncio # type: ignore[misc] diff --git a/tests/plugin/test_http_proxy_plugins_with_tls_interception.py b/tests/plugin/test_http_proxy_plugins_with_tls_interception.py index 232b0dd954..9845980201 100644 --- a/tests/plugin/test_http_proxy_plugins_with_tls_interception.py +++ b/tests/plugin/test_http_proxy_plugins_with_tls_interception.py @@ -9,6 +9,7 @@ :license: BSD, see LICENSE for more details. """ import ssl +import gzip import socket import pytest import selectors @@ -17,12 +18,13 @@ from typing import Any, cast from proxy.common.flag import FlagParser -from proxy.common.utils import bytes_, build_http_request, build_http_response +from proxy.common.utils import bytes_, build_http_request from proxy.core.connection import TcpClientConnection, TcpServerConnection -from proxy.http import httpMethods, httpStatusCodes, HttpProtocolHandler +from proxy.http import httpMethods, HttpProtocolHandler from proxy.http.proxy import HttpProxyPlugin -from proxy.http.parser import HttpParser +from proxy.http.parser import HttpParser, httpParserTypes +from proxy.http.responses import PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, okResponse from .utils import get_plugin_by_test_name @@ -174,7 +176,7 @@ async def test_modify_post_data_plugin(self) -> None: ) self.assertEqual(self.server.connection, self.server_ssl_connection) self._conn.send.assert_called_with( - HttpProxyPlugin.PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, + PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, ) self.assertFalse(self.protocol_handler.work.has_buffer()) @@ -229,7 +231,7 @@ async def test_man_in_the_middle_plugin(self) -> None: ) self.assertEqual(self.server.connection, self.server_ssl_connection) self._conn.send.assert_called_with( - HttpProxyPlugin.PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, + PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, ) self.assertFalse(self.protocol_handler.work.has_buffer()) # @@ -250,17 +252,14 @@ async def test_man_in_the_middle_plugin(self) -> None: self.server.flush.assert_called_once() # Server read - self.server.recv.return_value = memoryview( - build_http_response( - httpStatusCodes.OK, - reason=b'OK', body=b'Original Response From Upstream', - ), + self.server.recv.return_value = okResponse( + content=b'Original Response From Upstream', ) await self.protocol_handler._run_once() + response = HttpParser(httpParserTypes.RESPONSE_PARSER) + response.parse(self.protocol_handler.work.buffer[0].tobytes()) + assert response.body self.assertEqual( - self.protocol_handler.work.buffer[0].tobytes(), - build_http_response( - httpStatusCodes.OK, - reason=b'OK', body=b'Hello from man in the middle', - ), + gzip.decompress(response.body), + b'Hello from man in the middle', ) diff --git a/tests/testing/test_embed.py b/tests/testing/test_embed.py index 6db0458e81..1ab6f404e0 100644 --- a/tests/testing/test_embed.py +++ b/tests/testing/test_embed.py @@ -19,7 +19,7 @@ from proxy.common.constants import DEFAULT_CLIENT_RECVBUF_SIZE, PROXY_AGENT_HEADER_VALUE from proxy.common.utils import socket_connection, build_http_request from proxy.http import httpMethods -from proxy.http.server import HttpWebServerPlugin +from proxy.http.responses import NOT_FOUND_RESPONSE_PKT @pytest.mark.skipif( @@ -49,7 +49,7 @@ def test_with_proxy(self) -> None: response = conn.recv(DEFAULT_CLIENT_RECVBUF_SIZE) self.assertEqual( response, - HttpWebServerPlugin.DEFAULT_404_RESPONSE.tobytes(), + NOT_FOUND_RESPONSE_PKT.tobytes(), ) def test_proxy_vcr(self) -> None: