Repository: abhinavsingh/proxy.py
Pull request #98 (2019-09-26 04:09:56)

Pluggable WebServer
1. Either implement `HttpWebServerRoutePlugin` base class
2. Or, decorate your method with `proxy.route`

See `plugin_examples.py` for samples or `HttpWebServerPacFilePlugin`


Patch:
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
index 0b537ce..13b07dc 100644
--- a/.github/workflows/python-package.yml
+++ b/.github/workflows/python-package.yml
@@ -7,13 +7,13 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [ubuntu-latest, macOS-latest]
+        os: [macOS-latest, ubuntu-latest, windows-latest]
         python: [3.6, 3.7]
       max-parallel: 4
       fail-fast: false
     steps:
       - uses: actions/checkout@v1
-      - name: Set up Python ${{ matrix.python }}
+      - name: Set up Python
         uses: actions/setup-python@v1
         with:
           python-version: ${{ matrix.python }}
diff --git a/plugin_examples.py b/plugin_examples.py
index f9dbd2f..02cdcfb 100644
--- a/plugin_examples.py
+++ b/plugin_examples.py
@@ -15,6 +15,11 @@ from urllib import parse as urlparse
 import proxy
 
 
+@proxy.route(b'/hello-world')
+def hello_world(_request: proxy.HttpParser) -> bytes:
+    return proxy.HttpParser.build_response(200, body=b'Hello World')
+
+
 class RedirectToCustomServerPlugin(proxy.HttpProxyBasePlugin):
     """Modifies client request to redirect all incoming requests to a fixed server address."""
 
diff --git a/proxy.py b/proxy.py
index e1dce4a..bc0af48 100755
--- a/proxy.py
+++ b/proxy.py
@@ -28,7 +28,7 @@ import time
 from abc import ABC, abstractmethod
 from multiprocessing import connection
 from multiprocessing.reduction import send_handle, recv_handle
-from typing import Any, Dict, List, Tuple, Optional, Union, NamedTuple
+from typing import Any, Dict, List, Tuple, Optional, Union, NamedTuple, Type, Callable
 from urllib import parse as urlparse
 
 import select
@@ -96,7 +96,9 @@ def bytes_(s: Any, encoding: str = 'utf-8', errors: str = 'strict') -> Any:
 
 version = bytes_(__version__)
 CRLF, COLON, WHITESPACE, COMMA, DOT = b'\r\n', b':', b' ', b',', b'.'
-PROXY_AGENT_HEADER = b'Proxy-agent: proxy.py v' + version
+PROXY_AGENT_HEADER_KEY = b'Proxy-agent'
+PROXY_AGENT_HEADER_VALUE = b'proxy.py v' + version
+PROXY_AGENT_HEADER = PROXY_AGENT_HEADER_KEY + COLON + WHITESPACE + PROXY_AGENT_HEADER_VALUE
 
 ##
 # Various NamedTuples
@@ -141,6 +143,12 @@ HttpParserTypes = NamedTuple('HttpParserTypes', [
 ])
 httpParserTypes = HttpParserTypes(1, 2)
 
+HttpProtocolTypes = NamedTuple('HttpProtocolTypes', [
+    ('HTTP', int),
+    ('HTTPS', int),
+])
+httpProtocolTypes = HttpProtocolTypes(1, 2)
+
 
 class TcpConnection:
     """TCP server/client connection abstraction."""
@@ -321,7 +329,7 @@ class HttpProtocolConfig:
             auth_code: Optional[bytes] = DEFAULT_BASIC_AUTH,
             server_recvbuf_size: int = DEFAULT_SERVER_RECVBUF_SIZE,
             client_recvbuf_size: int = DEFAULT_CLIENT_RECVBUF_SIZE,
-            pac_file: Optional[bytes] = DEFAULT_PAC_FILE,
+            pac_file: Optional[str] = DEFAULT_PAC_FILE,
             pac_file_url_path: Optional[bytes] = DEFAULT_PAC_FILE_URL_PATH,
             plugins: Optional[Dict[bytes, List[type]]] = None,
             disable_headers: Optional[List[bytes]] = None,
@@ -897,13 +905,12 @@ class HttpProtocolBasePlugin(ABC):
 class ProxyConnectionFailed(HttpProtocolException):
     """Exception raised when HttpProxyPlugin is unable to establish connection to upstream server."""
 
-    RESPONSE_PKT = CRLF.join([
-        b'HTTP/1.1 502 Bad Gateway',
-        PROXY_AGENT_HEADER,
-        b'Content-Length: 11',
-        b'Connection: close',
-        CRLF
-    ]) + b'Bad Gateway'
+    RESPONSE_PKT = HttpParser.build_response(
+        502, reason=b'Bad Gateway',
+        headers={PROXY_AGENT_HEADER_KEY: PROXY_AGENT_HEADER_VALUE,
+                 b'Connection': b'close'},
+        body=b'Bad Gateway'
+    )
 
     def __init__(self, host: str, port: int, reason: str):
         self.host: str = host
@@ -922,14 +929,12 @@ class ProxyAuthenticationFailed(HttpProtocolException):
     """Exception raised when Http Proxy auth is enabled and
     incoming request doesn't present necessary credentials."""
 
-    RESPONSE_PKT = CRLF.join([
-        b'HTTP/1.1 407 Proxy Authentication Required',
-        PROXY_AGENT_HEADER,
-        b'Content-Length: 29',
-        b'Connection: close',
-        b'Proxy-Authenticate: Basic',
-        CRLF
-    ]) + b'Proxy Authentication Required'
+    RESPONSE_PKT = HttpParser.build_response(
+        407, reason=b'Proxy Authentication Required',
+        headers={PROXY_AGENT_HEADER_KEY: PROXY_AGENT_HEADER_VALUE,
+                 b'Connection': b'close',
+                 b'Proxy-Authenticate': b'Basic'},
+        body=b'Proxy Authentication Required')
 
     def response(self, _request: HttpParser) -> bytes:
         return self.RESPONSE_PKT
@@ -985,10 +990,9 @@ class HttpProxyBasePlugin(ABC):
 class HttpProxyPlugin(HttpProtocolBasePlugin):
     """HttpProtocolHandler plugin which implements HttpProxy specifications."""
 
-    PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT = CRLF.join([
-        b'HTTP/1.1 200 Connection established',
-        CRLF
-    ])
+    PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT = HttpParser.build_response(
+        200, reason=b'Connection established'
+    )
 
     # Used to synchronize with other HttpProxyPlugin instances while
     # generating certificates
@@ -1231,53 +1235,132 @@ class HttpProxyPlugin(HttpProtocolBasePlugin):
             raise HttpProtocolException()
 
 
-class HttpWebServerPlugin(HttpProtocolBasePlugin):
-    """HttpProtocolHandler plugin which handles incoming requests to local webserver."""
+class HttpWebServerRoutePlugin(ABC):
+    """Web Server Route Plugin."""
+
+    def __init__(
+            self,
+            config: HttpProtocolConfig,
+            client: TcpClientConnection):
+        self.config = config
+        self.client = client
+
+    def routes(self) -> List[Tuple[int, bytes]]:
+        """Return List(protocol, path) that this plugin handles."""
+        raise NotImplementedError()
 
-    DEFAULT_404_RESPONSE = CRLF.join([
-        b'HTTP/1.1 404 NOT FOUND',
-        b'Server: proxy.py v%s' % version,
-        b'Connection: Close',
-        CRLF
-    ])
+    def handle_request(self, request: HttpParser) -> None:
+        """Handle the request and serve response."""
+        raise NotImplementedError()
 
-    PAC_FILE_RESPONSE_PREFIX = CRLF.join([
-        b'HTTP/1.1 200 OK',
-        b'Content-Type: application/x-ns-proxy-autoconfig',
-        b'Connection: close',
-        CRLF
-    ])
+
+def route(path: bytes, protocols: Optional[List[int]] = None) -> \
+        Callable[[Callable[[HttpParser], bytes]], Type[HttpWebServerRoutePlugin]]:
+    def decorator(func: Callable[[HttpParser], bytes]) -> Type[HttpWebServerRoutePlugin]:
+        class HttpWebServerRouteHandler(HttpWebServerRoutePlugin):
+            @staticmethod
+            def name() -> str:
+                return func.__name__
+
+            def routes(self) -> List[Tuple[int, bytes]]:
+                p = protocols
+                if p is None:
+                    p = [httpProtocolTypes.HTTP, httpProtocolTypes.HTTPS]
+                return [(protocol, path) for protocol in p]
+
+            def handle_request(self, request: HttpParser) -> None:
+                self.client.queue(func(request))
+                self.client.flush()
+        return HttpWebServerRouteHandler
+    return decorator
+
+
+class HttpWebServerPacFilePlugin(HttpWebServerRoutePlugin):
 
     def __init__(
             self,
             config: HttpProtocolConfig,
-            client: TcpClientConnection,
-            request: HttpParser):
-        super().__init__(config, client, request)
+            client: TcpClientConnection):
+        super().__init__(config, client)
+        self.pac_file_response: Optional[bytes] = None
+        self.cache_pac_file_response()
+
+    def routes(self) -> List[Tuple[int, bytes]]:
+        if self.config.pac_file_url_path:
+            return [(httpProtocolTypes.HTTP, bytes_(self.config.pac_file_url_path))]
+        return []
+
+    def handle_request(self, request: HttpParser) -> None:
+        if self.config.pac_file and self.pac_file_response:
+            self.client.queue(self.pac_file_response)
+            self.client.flush()
+
+    def cache_pac_file_response(self) -> None:
         if self.config.pac_file:
             try:
                 with open(self.config.pac_file, 'rb') as f:
                     logger.debug('Will serve pac file from disk')
-                    self.pac_file_content = f.read()
+                    content = f.read()
             except IOError:
                 logger.debug('Will serve pac file content from buffer')
-                self.pac_file_content = self.config.pac_file
+                content = bytes_(self.config.pac_file)
+            self.pac_file_response = HttpParser.build_response(
+                200, reason=b'OK', headers={
+                    b'Content-Type': b'application/x-ns-proxy-autoconfig',
+                    b'Connection': b'close'
+                }, body=content
+            )
+
+
+class HttpWebServerPlugin(HttpProtocolBasePlugin):
+    """HttpProtocolHandler plugin which handles incoming requests to local webserver."""
+
+    DEFAULT_404_RESPONSE = HttpParser.build_response(
+        404, reason=b'NOT FOUND',
+        headers={b'Server': PROXY_AGENT_HEADER_VALUE,
+                 b'Connection': b'close'}
+    )
+
+    def __init__(
+            self,
+            config: HttpProtocolConfig,
+            client: TcpClientConnection,
+            request: HttpParser):
+        super().__init__(config, client, request)
+
+        self.routes: Dict[int, Dict[bytes, HttpWebServerRoutePlugin]] = {
+            httpProtocolTypes.HTTP: {},
+            httpProtocolTypes.HTTPS: {},
+        }
+
+        if b'HttpWebServerRoutePlugin' in self.config.plugins:
+            for klass in self.config.plugins[b'HttpWebServerRoutePlugin']:
+                instance = klass(self.config, self.client)
+                for (protocol, path) in instance.routes():
+                    self.routes[protocol][path] = instance
 
     def on_request_complete(self) -> Union[socket.socket, bool]:
         if self.request.has_upstream_server():
             return False
 
-        if self.config.pac_file and self.request.url and \
-                self.request.url.path == self.config.pac_file_url_path:
-            self.client.queue(self.PAC_FILE_RESPONSE_PREFIX)
-            self.client.queue(self.pac_file_content)
-            self.client.flush()
+        url = self.request.build_url()
+
+        # Routing
+        if self.request.method != b'CONNECT':
+            registered_routes = self.routes[httpProtocolTypes.HTTP]
         else:
-            # Catch all unhandled web server requests, return 404
-            self.client.queue(self.DEFAULT_404_RESPONSE)
-            # But is client ready for flush?
-            self.client.flush()
+            registered_routes = self.routes[httpProtocolTypes.HTTPS]
+        for route_ in registered_routes:
+            if url == route_:
+                registered_routes[route_].handle_request(self.request)
+                # But is client ready for flush?
+                self.client.flush()
+                return True
 
+        # Catch all unhandled web server requests, return 404
+        self.client.queue(self.DEFAULT_404_RESPONSE)
+        # But is client ready for flush?
+        self.client.flush()
         return True
 
     def access_log(self) -> None:
@@ -1477,8 +1560,8 @@ class HttpProtocolHandler(threading.Thread):
                 try:
                     self.client.conn.shutdown(socket.SHUT_RDWR)
                     self.client.close()
-                except OSError as e:
-                    logger.exception('OSError: %s', str(e), exc_info=e)
+                except OSError:
+                    pass
 
             # Invoke plugin.on_client_connection_close
             for plugin in self.plugins.values():
@@ -1513,19 +1596,26 @@ def load_plugins(plugins: bytes) -> Dict[bytes, List[type]]:
     a list of respective Python classes."""
     p: Dict[bytes, List[type]] = {
         b'HttpProtocolBasePlugin': [],
-        b'HttpProxyBasePlugin': []
+        b'HttpProxyBasePlugin': [],
+        b'HttpWebServerRoutePlugin': [],
     }
     for plugin in plugins.split(COMMA):
         plugin = plugin.strip()
         if plugin == b'':
             continue
         module_name, klass_name = plugin.rsplit(DOT, 1)
-        module_name = __name__ if module_name == b'proxy' else text_(module_name)
-        module = importlib.import_module(__name__)
-        klass = getattr(module, text_(klass_name))
-        base_klass = inspect.getmro(klass)[::-1][2:][0]
+        if module_name == 'proxy':
+            klass = getattr(sys.modules[__name__], text_(klass_name))
+        else:
+            klass = getattr(importlib.import_module(text_(module_name)), text_(klass_name))
+        base_klass = inspect.getmro(klass)[1]
         p[bytes_(base_klass.__name__)].append(klass)
-        logger.info('Loaded plugin %s', klass)
+        logger.info(
+            'Loaded %s %s.%s',
+            'plugin' if klass.__name__ != 'HttpWebServerRouteHandler' else 'route',
+            text_(module_name),
+            # HttpWebServerRouteHandler route decorator adds a special staticmethod to return decorated function name
+            klass.__name__ if klass.__name__ != 'HttpWebServerRouteHandler' else klass.name())
     return p
 
 
@@ -1667,7 +1757,7 @@ def init_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         '--pac-file-url-path',
         type=str,
-        default=DEFAULT_PAC_FILE_URL_PATH,
+        default=text_(DEFAULT_PAC_FILE_URL_PATH),
         help='Default: %s. Web server path to serve the PAC file.' %
              text_(DEFAULT_PAC_FILE_URL_PATH))
     parser.add_argument(
@@ -1734,8 +1824,8 @@ def main(input_args: List[str]) -> None:
             auth_code=auth_code,
             server_recvbuf_size=args.server_recvbuf_size,
             client_recvbuf_size=args.client_recvbuf_size,
-            pac_file=args.pac_file,
-            pac_file_url_path=args.pac_file_url_path,
+            pac_file=bytes_(args.pac_file),
+            pac_file_url_path=bytes_(args.pac_file_url_path),
             disable_headers=[
                 header.lower() for header in bytes_(
                     args.disable_headers).split(COMMA) if header.strip() != b''],
@@ -1749,14 +1839,15 @@ def main(input_args: List[str]) -> None:
             port=args.port,
             backlog=args.backlog,
             num_workers=args.num_workers if args.num_workers > 0 else multiprocessing.cpu_count())
-        if config.pac_file is not None:
-            args.enable_web_server = True
 
         default_plugins = ''
         if not args.disable_http_proxy:
             default_plugins += 'proxy.HttpProxyPlugin,'
-        if args.enable_web_server:
+        if args.enable_web_server or config.pac_file is not None:
             default_plugins += 'proxy.HttpWebServerPlugin,'
+        if config.pac_file is not None:
+            default_plugins += 'proxy.HttpWebServerPacFilePlugin,'
+
         config.plugins = load_plugins(
             bytes_(
                 '%s%s' %
diff --git a/requirements-release.txt b/requirements-release.txt
new file mode 100644
index 0000000..f3c2cd3
--- /dev/null
+++ b/requirements-release.txt
@@ -0,0 +1,1 @@
+twine==2.0.0


Test patch:
diff --git a/requirements-testing.txt b/requirements-testing.txt
index 874643d..e567c42 100644
--- a/requirements-testing.txt
+++ b/requirements-testing.txt
@@ -1,7 +1,6 @@
 python-coveralls==2.9.3
 coverage==4.5.4
 flake8==3.7.8
-twine==2.0.0
 pytest==5.1.3
 autopep8==1.4.4
 mypy==0.720
diff --git a/tests.py b/tests.py
index c1d7cb8..4c08d51 100644
--- a/tests.py
+++ b/tests.py
@@ -241,6 +241,7 @@ class MockHttpProxy:
     def start(self) -> None:
         self.client.conn.sendall(proxy.CRLF.join(
             [b'HTTP/1.1 200 OK', proxy.CRLF]))
+        self.client.conn.shutdown(socket.SHUT_RDWR)
         self.client.conn.close()
 
 
@@ -250,6 +251,8 @@ def mock_tcp_proxy_side_effect(client: proxy.TcpClientConnection, **kwargs) -> M
 
 @unittest.skipIf(os.getenv('TESTING_ON_TRAVIS', 0),
                  'Opening sockets not allowed on Travis')
+@unittest.skipIf(os.getenv('GITHUB_WORKFLOW', 0),
+                 'This test fails on GitHub Windows environment')
 class TestMultiCoreRequestDispatcherIntegration(unittest.TestCase):
     tcp_port = None
     tcp_server = None
@@ -276,11 +279,10 @@ class TestMultiCoreRequestDispatcherIntegration(unittest.TestCase):
                     try:
                         sock.connect(
                             (str(proxy.DEFAULT_IPV4_HOSTNAME), self.tcp_port))
-                        sock.send(proxy.CRLF.join([
-                            b'GET http://httpbin.org/get HTTP/1.1',
-                            b'Host: httpbin.org',
-                            proxy.CRLF
-                        ]))
+                        sock.send(proxy.HttpParser.build_request(
+                            b'GET', b'http://httpbin.org/get', b'HTTP/1.1',
+                            headers={b'Host': b'httpbin.org'})
+                        )
                         data = sock.recv(proxy.DEFAULT_BUFFER_SIZE)
                         self.assertEqual(data, proxy.CRLF.join(
                             [b'HTTP/1.1 200 OK', proxy.CRLF]))
@@ -1034,7 +1036,7 @@ class TestHttpProtocolHandler(unittest.TestCase):
     @mock.patch('select.select')
     def test_pac_file_served_from_disk(self, mock_select: mock.Mock) -> None:
         mock_select.return_value = [self._conn], [], []
-        config = proxy.HttpProtocolConfig(pac_file=b'proxy.pac')
+        config = proxy.HttpProtocolConfig(pac_file='proxy.pac')
         self.init_and_make_pac_file_request(config)
         self.proxy.run_once()
         self.assertEqual(
@@ -1043,14 +1045,18 @@ class TestHttpProtocolHandler(unittest.TestCase):
         with open('proxy.pac', 'rb') as pac_file:
             self.assertEqual(
                 self._conn.received,
-                proxy.HttpWebServerPlugin.PAC_FILE_RESPONSE_PREFIX +
-                pac_file.read())
+                proxy.HttpParser.build_response(
+                    200, reason=b'OK', headers={
+                        b'Content-Type': b'application/x-ns-proxy-autoconfig',
+                        b'Connection': b'close'
+                    }, body=pac_file.read()
+                ))
 
     @mock.patch('select.select')
     def test_pac_file_served_from_buffer(self, mock_select: mock.Mock) -> None:
         pac_file_content = b'function FindProxyForURL(url, host) { return "PROXY localhost:8899; DIRECT"; }'
         mock_select.return_value = [self._conn], [], []
-        config = proxy.HttpProtocolConfig(pac_file=pac_file_content)
+        config = proxy.HttpProtocolConfig(pac_file=proxy.text_(pac_file_content))
         self.init_and_make_pac_file_request(config)
         self.proxy.run_once()
         self.assertEqual(
@@ -1058,8 +1064,12 @@ class TestHttpProtocolHandler(unittest.TestCase):
             proxy.httpParserStates.COMPLETE)
         self.assertEqual(
             self._conn.received,
-            proxy.HttpWebServerPlugin.PAC_FILE_RESPONSE_PREFIX +
-            pac_file_content)
+            proxy.HttpParser.build_response(
+                200, reason=b'OK', headers={
+                    b'Content-Type': b'application/x-ns-proxy-autoconfig',
+                    b'Connection': b'close'
+                }, body=pac_file_content
+            ))
 
     @mock.patch('select.select')
     def test_default_web_server_returns_404(self, mock_select: mock.Mock) -> None:
@@ -1100,7 +1110,7 @@ class TestHttpProtocolHandler(unittest.TestCase):
 
     def init_and_make_pac_file_request(self, config: proxy.HttpProtocolConfig) -> None:
         config.plugins = proxy.load_plugins(
-            b'proxy.HttpProxyPlugin,proxy.HttpWebServerPlugin')
+            b'proxy.HttpProxyPlugin,proxy.HttpWebServerPlugin,proxy.HttpWebServerPacFilePlugin')
         self.proxy = proxy.HttpProtocolHandler(
             proxy.TcpClientConnection(
                 self._conn, self._addr), config=config)