From a182213c57a52f6443176feccb938304ed21fa3c Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Tue, 23 Jul 2019 17:57:47 -0700 Subject: [PATCH] Refactor 1. Update class name to be more verbose 2. Update docstring for important classes 3. Global default constants --- README.md | 2 +- proxy.py | 125 ++++++++++++++++++++++++++----------------- requirements-dev.txt | 4 +- tests.py | 10 ++-- 4 files changed, 84 insertions(+), 57 deletions(-) diff --git a/README.md b/README.md index 8bc154fee5..dee4ebf50e 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ proxy.py ======== -Lightweight HTTP, HTTPS and WebSockets Proxy Server in Python. +Lightweight HTTP, HTTPS, HTTP2 and WebSockets Proxy Server in Python. ![alt text](https://travis-ci.org/abhinavsingh/proxy.py.svg?branch=develop "Build Status") diff --git a/proxy.py b/proxy.py index 1ccdc23d05..c458963183 100755 --- a/proxy.py +++ b/proxy.py @@ -49,6 +49,20 @@ import urlparse import Queue as queue +# Defaults +DEFAULT_BACKLOG = 100 +DEFAULT_BASIC_AUTH = None +DEFAULT_BUFFER_SIZE = 8192 +DEFAULT_CLIENT_RECVBUF_SIZE = DEFAULT_BUFFER_SIZE +DEFAULT_SERVER_RECVBUF_SIZE = DEFAULT_BUFFER_SIZE +DEFAULT_HOSTNAME = '127.0.0.1' +DEFAULT_PORT = 8899 +DEFAULT_IPV4 = False +DEFAULT_LOG_LEVEL = 'INFO' +DEFAULT_OPEN_FILE_LIMIT = 1024 +DEFAULT_PAC_FILE = None +DEFAULT_NUM_WORKERS = 0 + def text_(s, encoding='utf-8', errors='strict'): # pragma: no cover """Utility to ensure text-like usability. @@ -337,7 +351,7 @@ def split(data): return line, data -class Connection(object): +class TCPConnection(object): """TCP server/client connection abstraction.""" def __init__(self, what): @@ -350,7 +364,7 @@ def send(self, data): # TODO: Gracefully handle BrokenPipeError exceptions return self.conn.send(data) - def recv(self, bufsiz=8192): + def recv(self, bufsiz=DEFAULT_BUFFER_SIZE): try: data = self.conn.recv(bufsiz) if len(data) == 0: @@ -385,11 +399,11 @@ def flush(self): logger.debug('flushed %d bytes to %s' % (sent, self.what)) -class Server(Connection): +class TCPServerConnection(TCPConnection): """Establish connection to destination server.""" def __init__(self, host, port): - super(Server, self).__init__(b'server') + super(TCPServerConnection, self).__init__(b'server') self.addr = (host, int(port)) def __del__(self): @@ -400,11 +414,11 @@ def connect(self): self.conn = socket.create_connection((self.addr[0], self.addr[1])) -class Client(Connection): +class TCPClientConnection(TCPConnection): """Accepted client connection.""" def __init__(self, conn, addr): - super(Client, self).__init__(b'client') + super(TCPClientConnection, self).__init__(b'client') self.conn = conn self.addr = addr @@ -428,14 +442,15 @@ class ProxyAuthenticationFailed(ProxyError): pass -class Proxy(threading.Thread): +class HTTPProxy(threading.Thread): """HTTP proxy implementation. Accepts `Client` connection object and act as a proxy between client and server. """ - def __init__(self, client, auth_code=None, server_recvbuf_size=8192, client_recvbuf_size=8192, pac_file=None): - super(Proxy, self).__init__() + def __init__(self, client, auth_code=DEFAULT_BASIC_AUTH, server_recvbuf_size=DEFAULT_SERVER_RECVBUF_SIZE, + client_recvbuf_size=DEFAULT_CLIENT_RECVBUF_SIZE, pac_file=DEFAULT_PAC_FILE): + super(HTTPProxy, self).__init__() self.start_time = self._now() self.last_activity = self.start_time @@ -494,7 +509,7 @@ def _process_request(self, data): self._serve_pac_file() return True - self.server = Server(host, port) + self.server = TCPServerConnection(host, port) try: logger.debug('connecting to server %s:%s' % (host, port)) self.server.connect() @@ -569,7 +584,7 @@ def _process_rlist(self, r): return self._process_request(data) except (ProxyAuthenticationFailed, ProxyConnectionFailed) as e: logger.exception(e) - self.client.queue(Proxy._get_response_pkt_by_exception(e)) + self.client.queue(HTTPProxy._get_response_pkt_by_exception(e)) self.client.flush() return True @@ -644,13 +659,14 @@ def run(self): logger.debug('Closing proxy for connection %r at address %r' % (self.client.conn, self.client.addr)) -class TCP(object): - """TCP server implementation. +class TCPServer(object): + """TCPServer server implementation. - Subclass MUST implement `handle` method which accepts an instance of accepted `Client` connection. + Inheritor MUST implement `handle` method. It accepts an instance of `TCPClientConnection`. + Optionally, can also implement `setup` and `shutdown` methods for custom bootstrapping and teardown. """ - def __init__(self, hostname='127.0.0.1', port=8899, backlog=100, ipv4=False): + def __init__(self, hostname=DEFAULT_HOSTNAME, port=DEFAULT_PORT, backlog=DEFAULT_BACKLOG, ipv4=DEFAULT_IPV4): self.hostname = hostname self.port = port self.backlog = backlog @@ -676,7 +692,7 @@ def run(self): logger.info('Started server on port %d' % self.port) while True: conn, addr = self.socket.accept() - client = Client(conn, addr) + client = TCPClientConnection(conn, addr) self.handle(client) except Exception as e: logger.exception('Exception while running the server %r' % e) @@ -686,15 +702,19 @@ def run(self): self.socket.close() -class HTTP(TCP): +class HTTPServer(TCPServer): """HTTP server implementation. - Spawns new process which either serve local server content or proxy the accepted client connection. + Pre-spawns worker process to utilize all cores available on the system. Accepted `TCPClientConnection` is + dispatched over a queue to workers. One of the worker picks up the work and starts a new thread to handle the + client request. """ - def __init__(self, hostname='127.0.0.1', port=8899, backlog=100, num_workers=0, - auth_code=None, server_recvbuf_size=8192, client_recvbuf_size=8192, pac_file=None, ipv4=False): - super(HTTP, self).__init__(hostname, port, backlog, ipv4) + def __init__(self, hostname=DEFAULT_HOSTNAME, port=DEFAULT_PORT, backlog=DEFAULT_BACKLOG, + num_workers=DEFAULT_NUM_WORKERS, + auth_code=DEFAULT_BASIC_AUTH, server_recvbuf_size=DEFAULT_SERVER_RECVBUF_SIZE, + client_recvbuf_size=DEFAULT_CLIENT_RECVBUF_SIZE, pac_file=DEFAULT_PAC_FILE, ipv4=DEFAULT_IPV4): + super(HTTPServer, self).__init__(hostname, port, backlog, ipv4) self.auth_code = auth_code self.client_recvbuf_size = client_recvbuf_size self.server_recvbuf_size = server_recvbuf_size @@ -716,7 +736,7 @@ def setup(self): self.workers.append(worker) def handle(self, client): - self.worker_queue.put((Worker.operations.DEFAULT, {'client': client})) + self.worker_queue.put((Worker.operations.DEFAULT, client)) def shutdown(self): logger.info('Shutting down %d workers' % self.num_workers) @@ -727,12 +747,19 @@ def shutdown(self): class Worker(multiprocessing.Process): + """Generic worker class implementation. + + Worker instance accepts (operation, payload) over work queue and + starts a new thread to complete the work. + """ + operations = namedtuple('WorkerOperations', ( - 'DEFAULT', # Default worker action i.e. handle http request. + 'DEFAULT', # Default worker action 'SHUTDOWN', ))(1, 2) - def __init__(self, work_queue, auth_code=None, server_recvbuf_size=8192, client_recvbuf_size=8192, pac_file=None): + def __init__(self, work_queue, auth_code=DEFAULT_BASIC_AUTH, server_recvbuf_size=DEFAULT_SERVER_RECVBUF_SIZE, + client_recvbuf_size=DEFAULT_CLIENT_RECVBUF_SIZE, pac_file=DEFAULT_PAC_FILE): super(Worker, self).__init__() self.work_queue = work_queue self.auth_code = auth_code @@ -745,11 +772,11 @@ def run(self): try: op, payload = self.work_queue.get(True, 1) if op == Worker.operations.DEFAULT: - proxy = Proxy(payload['client'], - auth_code=self.auth_code, - server_recvbuf_size=self.server_recvbuf_size, - client_recvbuf_size=self.client_recvbuf_size, - pac_file=self.pac_file) + proxy = HTTPProxy(payload, + auth_code=self.auth_code, + server_recvbuf_size=self.server_recvbuf_size, + client_recvbuf_size=self.client_recvbuf_size, + pac_file=self.pac_file) proxy.daemon = True proxy.start() elif op == Worker.operations.SHUTDOWN: @@ -778,39 +805,39 @@ def parse_args(args): epilog='Proxy.py not working? Report at: %s/issues/new' % __homepage__ ) # Argument names are ordered alphabetically. - parser.add_argument('--backlog', type=int, default=100, + parser.add_argument('--backlog', type=int, default=DEFAULT_BACKLOG, help='Default: 100. Maximum number of pending connections to proxy server') - parser.add_argument('--basic-auth', type=str, default=None, + parser.add_argument('--basic-auth', type=str, default=DEFAULT_BASIC_AUTH, help='Default: No authentication. Specify colon separated user:password ' 'to enable basic authentication.') - parser.add_argument('--client-recvbuf-size', type=int, default=8192, + parser.add_argument('--client-recvbuf-size', type=int, default=DEFAULT_CLIENT_RECVBUF_SIZE, help='Default: 8 KB. Maximum amount of data received from the ' 'client in a single recv() operation. Bump this ' 'value for faster uploads at the expense of ' 'increased RAM.') - parser.add_argument('--hostname', type=str, default='127.0.0.1', + parser.add_argument('--hostname', type=str, default=DEFAULT_HOSTNAME, help='Default: 127.0.0.1. Server IP address.') - parser.add_argument('--ipv4', action='store_true', default=False, + parser.add_argument('--ipv4', action='store_true', default=DEFAULT_IPV4, help='Whether to listen on IPv4 address. ' 'By default server only listens on IPv6.') - parser.add_argument('--log-level', type=str, default='INFO', + parser.add_argument('--log-level', type=str, default=DEFAULT_LOG_LEVEL, help='Valid options: DEBUG, INFO (default), WARNING, ERROR, CRITICAL. ' 'Both upper and lowercase values are allowed.' 'You may also simply use the leading character e.g. --log-level d') - parser.add_argument('--open-file-limit', type=int, default=1024, + parser.add_argument('--open-file-limit', type=int, default=DEFAULT_OPEN_FILE_LIMIT, help='Default: 1024. Maximum number of files (TCP connections) ' 'that proxy.py can open concurrently.') - parser.add_argument('--port', type=int, default=8899, + parser.add_argument('--port', type=int, default=DEFAULT_PORT, help='Default: 8899. Server port.') - parser.add_argument('--pac-file', type=str, default=None, + parser.add_argument('--pac-file', type=str, default=DEFAULT_PAC_FILE, help='A file (Proxy Auto Configuration) or string to serve when ' 'the server receives a direct file request.') - parser.add_argument('--server-recvbuf-size', type=int, default=8192, + parser.add_argument('--server-recvbuf-size', type=int, default=DEFAULT_SERVER_RECVBUF_SIZE, help='Default: 8 KB. Maximum amount of data received from the ' 'server in a single recv() operation. Bump this ' 'value for faster downloads at the expense of ' 'increased RAM.') - parser.add_argument('--num-workers', type=int, default=0, + parser.add_argument('--num-workers', type=int, default=DEFAULT_NUM_WORKERS, help='Defaults to number of CPU cores.') return parser.parse_args(args) @@ -834,15 +861,15 @@ def main(): if args.basic_auth: auth_code = b'Basic %s' % base64.b64encode(bytes_(args.basic_auth)) - server = HTTP(hostname=args.hostname, - port=args.port, - backlog=args.backlog, - auth_code=auth_code, - server_recvbuf_size=args.server_recvbuf_size, - client_recvbuf_size=args.client_recvbuf_size, - pac_file=args.pac_file, - ipv4=args.ipv4, - num_workers=args.num_workers) + server = HTTPServer(hostname=args.hostname, + port=args.port, + backlog=args.backlog, + auth_code=auth_code, + server_recvbuf_size=args.server_recvbuf_size, + client_recvbuf_size=args.client_recvbuf_size, + pac_file=args.pac_file, + ipv4=args.ipv4, + num_workers=args.num_workers) server.run() except KeyboardInterrupt: pass diff --git a/requirements-dev.txt b/requirements-dev.txt index f28d14fd6e..174eaf8ff4 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,3 +1,3 @@ coverage==4.5.2 -pyflakes==2.0.0 -twine==1.12.1 \ No newline at end of file +flake8==3.7.8 +twine==1.12.1 diff --git a/tests.py b/tests.py index 6e516a301b..89a3b77ecf 100644 --- a/tests.py +++ b/tests.py @@ -15,7 +15,7 @@ import unittest from threading import Thread from contextlib import closing -from proxy import Proxy, ChunkParser, HttpParser, Client +from proxy import HTTPProxy, ChunkParser, HttpParser, TCPClientConnection from proxy import ProxyAuthenticationFailed, ProxyConnectionFailed from proxy import CRLF, version, PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT @@ -423,7 +423,7 @@ def tearDownClass(cls): def setUp(self): self._conn = MockConnection() self._addr = ('127.0.0.1', 54382) - self.proxy = Proxy(Client(self._conn, self._addr)) + self.proxy = HTTPProxy(TCPClientConnection(self._conn, self._addr)) def test_http_get(self): # Send request line @@ -506,7 +506,7 @@ def test_proxy_connection_failed(self): ])) def test_proxy_authentication_failed(self): - self.proxy = Proxy(Client(self._conn, self._addr), b'Basic %s' % base64.b64encode(b'user:pass')) + self.proxy = HTTPProxy(TCPClientConnection(self._conn, self._addr), b'Basic %s' % base64.b64encode(b'user:pass')) with self.assertRaises(ProxyAuthenticationFailed): self.proxy._process_request(CRLF.join([ @@ -516,7 +516,7 @@ def test_proxy_authentication_failed(self): ])) def test_authenticated_proxy_http_get(self): - self.proxy = Proxy(Client(self._conn, self._addr), b'Basic %s' % base64.b64encode(b'user:pass')) + self.proxy = HTTPProxy(TCPClientConnection(self._conn, self._addr), b'Basic %s' % base64.b64encode(b'user:pass')) self.proxy.client.conn.queue((b'GET http://localhost:%d HTTP/1.1' % self.http_server_port) + CRLF) self.proxy._process_request(self.proxy.client.recv()) @@ -549,7 +549,7 @@ def test_authenticated_proxy_http_get(self): self.assertEqual(int(self.proxy.response.code), 200) def test_authenticated_proxy_http_tunnel(self): - self.proxy = Proxy(Client(self._conn, self._addr), b'Basic %s' % base64.b64encode(b'user:pass')) + self.proxy = HTTPProxy(TCPClientConnection(self._conn, self._addr), b'Basic %s' % base64.b64encode(b'user:pass')) self.proxy.client.conn.queue(CRLF.join([ b'CONNECT localhost:%d HTTP/1.1' % self.http_server_port,