Rooba
/
proxy.py
mirror of https://github.com/abhinavsingh/proxy.py.git
1
1
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add --timeout flag with default value of 10 second. (#129) 

* Add --timeout flag with default value of 5.  This value was previously hardcoded to 30

* --timeout=10 by default

* Dispatch 408 timeout when connection is dropped due to inactivity

* Add httpStatusCodes named tuple

* Update plugin client connection reference after TLS connection upgrade
This commit is contained in:
Abhinav Singh 2019-10-12 21:02:17 -07:00 committed by GitHub
parent 6455a34973
commit dc560be6ea
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 116 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored
View File

@ -14,5 +14,6 @@ build
proxy.py.egg-info
proxy.py.iml
*.pyc
*.pem
ca-*.pem
https-*.pem
benchmark.py

8
README.md
View File

@ -905,9 +905,10 @@ usage: proxy.py [-h] [--backlog BACKLOG] [--basic-auth BASIC_AUTH]
[--pac-file-url-path PAC_FILE_URL_PATH] [--pid-file PID_FILE]
[--plugins PLUGINS] [--port PORT]
[--server-recvbuf-size SERVER_RECVBUF_SIZE]
[--static-server-dir STATIC_SERVER_DIR] [--version]
[--static-server-dir STATIC_SERVER_DIR] [--timeout TIMEOUT]
[--version]
proxy.py v1.1.0
proxy.py v1.2.0
optional arguments:
-h, --help show this help message and exit
@ -994,6 +995,9 @@ optional arguments:
server root directory. This option is only applicable
when static server is also enabled. See --enable-
static-server.
--timeout TIMEOUT Default: 10. Number of seconds after which an inactive
connection must be dropped. Inactivity is defined by
no data sent or received by the client.
--version, -v Prints proxy.py version.
Proxy.py not working? Report at:

17
plugin_examples.py
View File

@ -81,14 +81,16 @@ class ProposedRestApiPlugin(proxy.HttpProxyBasePlugin):
assert request.path
if request.path in self.REST_API_SPEC:
self.client.queue(proxy.build_http_response(
200, reason=b'OK',
proxy.httpStatusCodes.OK,
reason=b'OK',
headers={b'Content-Type': b'application/json'},
body=proxy.bytes_(json.dumps(
self.REST_API_SPEC[request.path]))
))
else:
self.client.queue(proxy.build_http_response(
404, reason=b'NOT FOUND', body=b'Not Found'
proxy.httpStatusCodes.NOT_FOUND,
reason=b'NOT FOUND', body=b'Not Found'
))
return None
@ -134,7 +136,7 @@ class FilterByUpstreamHostPlugin(proxy.HttpProxyBasePlugin):
def before_upstream_connection(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]:
if request.host in self.FILTERED_DOMAINS:
raise proxy.HttpRequestRejected(
status_code=418, reason=b'I\'m a tea pot')
status_code=proxy.httpStatusCodes.I_AM_A_TEAPOT, reason=b'I\'m a tea pot')
return request
def handle_client_request(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]:
@ -194,7 +196,8 @@ class ManInTheMiddlePlugin(proxy.HttpProxyBasePlugin):
def handle_upstream_chunk(self, chunk: bytes) -> bytes:
return proxy.build_http_response(
200, reason=b'OK', body=b'Hello from man in the middle')
proxy.httpStatusCodes.OK,
reason=b'OK', body=b'Hello from man in the middle')
def on_upstream_connection_close(self) -> None:
pass
@ -212,9 +215,11 @@ class WebServerPlugin(proxy.HttpWebServerBasePlugin):
def handle_request(self, request: proxy.HttpParser) -> None:
if request.path == b'/http-route-example':
self.client.queue(proxy.build_http_response(200, body=b'HTTP route response'))
self.client.queue(proxy.build_http_response(
proxy.httpStatusCodes.OK, body=b'HTTP route response'))
elif request.path == b'/https-route-example':
self.client.queue(proxy.build_http_response(200, body=b'HTTPS route response'))
self.client.queue(proxy.build_http_response(
proxy.httpStatusCodes.OK, body=b'HTTPS route response'))
def on_websocket_open(self) -> None:
proxy.logger.info('Websocket open')

129
proxy.py
View File

@ -50,7 +50,7 @@ if os.name != 'nt':
PROXY_PY_DIR = os.path.dirname(os.path.realpath(__file__))
PROXY_PY_START_TIME = time.time()
VERSION = (1, 1, 2)
VERSION = (1, 2, 0)
__version__ = '.'.join(map(str, VERSION[0:3]))
__description__ = '⚡⚡⚡ Fast, Lightweight, Programmable Proxy Server in a single Python file.'
__author__ = 'Abhinav Singh'
@ -62,35 +62,36 @@ __license__ = 'BSD'
# Defaults
DEFAULT_BACKLOG = 100
DEFAULT_BASIC_AUTH = None
DEFAULT_CA_KEY_FILE = None
DEFAULT_BUFFER_SIZE = 1024 * 1024
DEFAULT_CA_CERT_DIR = None
DEFAULT_CA_CERT_FILE = None
DEFAULT_CA_KEY_FILE = None
DEFAULT_CA_SIGNING_KEY_FILE = None
DEFAULT_CERT_FILE = None
DEFAULT_BUFFER_SIZE = 1024 * 1024
DEFAULT_CLIENT_RECVBUF_SIZE = DEFAULT_BUFFER_SIZE
DEFAULT_SERVER_RECVBUF_SIZE = DEFAULT_BUFFER_SIZE
DEFAULT_DEVTOOLS_WS_PATH = b'/devtools'
DEFAULT_DISABLE_HEADERS: List[bytes] = []
DEFAULT_DISABLE_HTTP_PROXY = False
DEFAULT_ENABLE_DEVTOOLS = False
DEFAULT_ENABLE_STATIC_SERVER = False
DEFAULT_ENABLE_WEB_SERVER = False
DEFAULT_IPV4_HOSTNAME = ipaddress.IPv4Address('127.0.0.1')
DEFAULT_IPV6_HOSTNAME = ipaddress.IPv6Address('::1')
DEFAULT_KEY_FILE = None
DEFAULT_PORT = 8899
DEFAULT_DISABLE_HTTP_PROXY = False
DEFAULT_ENABLE_DEVTOOLS = False
DEFAULT_DEVTOOLS_WS_PATH = b'/devtools'
DEFAULT_ENABLE_STATIC_SERVER = False
DEFAULT_ENABLE_WEB_SERVER = False
DEFAULT_LOG_FILE = None
DEFAULT_LOG_FORMAT = '%(asctime)s - pid:%(process)d [%(levelname)-.1s] %(funcName)s:%(lineno)d - %(message)s'
DEFAULT_LOG_LEVEL = 'INFO'
DEFAULT_NUM_WORKERS = 0
DEFAULT_OPEN_FILE_LIMIT = 1024
DEFAULT_PAC_FILE = None
DEFAULT_PAC_FILE_URL_PATH = b'/'
DEFAULT_PID_FILE = None
DEFAULT_NUM_WORKERS = 0
DEFAULT_PLUGINS = '' # Comma separated list of plugins
DEFAULT_PLUGINS = ''
DEFAULT_PORT = 8899
DEFAULT_SERVER_RECVBUF_SIZE = DEFAULT_BUFFER_SIZE
DEFAULT_STATIC_SERVER_DIR = os.path.join(PROXY_PY_DIR, 'public')
DEFAULT_TIMEOUT = 10
DEFAULT_VERSION = False
DEFAULT_LOG_FORMAT = '%(asctime)s - pid:%(process)d [%(levelname)-.1s] %(funcName)s:%(lineno)d - %(message)s'
DEFAULT_LOG_FILE = None
UNDER_TEST = False # Set to True if under test
logger = logging.getLogger(__name__)
@ -140,6 +141,35 @@ ChunkParserStates = NamedTuple('ChunkParserStates', [
])
chunkParserStates = ChunkParserStates(1, 2, 3)
HttpStatusCodes = NamedTuple('HttpStatusCodes', [
# 1xx
('CONTINUE', int),
('SWITCHING_PROTOCOLS', int),
# 2xx
('OK', int),
# 4xx
('BAD_REQUEST', int),
('UNAUTHORIZED', int),
('FORBIDDEN', int),
('NOT_FOUND', int),
('PROXY_AUTH_REQUIRED', int),
('REQUEST_TIMEOUT', int),
('I_AM_A_TEAPOT', int),
# 5xx
('INTERNAL_SERVER_ERROR', int),
('NOT_IMPLEMENTED', int),
('BAD_GATEWAY', int),
('GATEWAY_TIMEOUT', int),
('NETWORK_READ_TIMEOUT_ERROR', int),
('NETWORK_CONNECT_TIMEOUT_ERROR', int),
])
httpStatusCodes = HttpStatusCodes(
100, 101,
200,
400, 401, 403, 404, 407, 408, 418,
500, 501, 502, 504, 598, 599
)
HttpMethods = NamedTuple('HttpMethods', [
('GET', bytes),
('HEAD', bytes),
@ -893,7 +923,8 @@ class ProxyConnectionFailed(ProtocolException):
"""Exception raised when HttpProxyPlugin is unable to establish connection to upstream server."""
RESPONSE_PKT = build_http_response(
502, reason=b'Bad Gateway',
httpStatusCodes.BAD_GATEWAY,
reason=b'Bad Gateway',
headers={
PROXY_AGENT_HEADER_KEY: PROXY_AGENT_HEADER_VALUE,
b'Connection': b'close'
@ -915,7 +946,8 @@ class ProxyAuthenticationFailed(ProtocolException):
incoming request doesn't present necessary credentials."""
RESPONSE_PKT = build_http_response(
407, reason=b'Proxy Authentication Required',
httpStatusCodes.PROXY_AUTH_REQUIRED,
reason=b'Proxy Authentication Required',
headers={
PROXY_AGENT_HEADER_KEY: PROXY_AGENT_HEADER_VALUE,
b'Proxy-Authenticate': b'Basic',
@ -965,7 +997,9 @@ class ProtocolConfig:
static_server_dir: str = DEFAULT_STATIC_SERVER_DIR,
enable_static_server: bool = DEFAULT_ENABLE_STATIC_SERVER,
devtools_event_queue: Optional[DevtoolsEventQueueType] = None,
devtools_ws_path: bytes = DEFAULT_DEVTOOLS_WS_PATH) -> None:
devtools_ws_path: bytes = DEFAULT_DEVTOOLS_WS_PATH,
timeout: int = DEFAULT_TIMEOUT) -> None:
self.timeout = timeout
self.auth_code = auth_code
self.server_recvbuf_size = server_recvbuf_size
self.client_recvbuf_size = client_recvbuf_size
@ -1130,7 +1164,7 @@ class HttpProxyBasePlugin(ABC):
Raise HttpRequestRejected or ProtocolException directly to
teardown the connection with client.
"""
return request
return request # pragma: no cover
@abstractmethod
def handle_upstream_chunk(self, chunk: bytes) -> bytes:
@ -1150,7 +1184,8 @@ class HttpProxyPlugin(ProtocolHandlerPlugin):
"""ProtocolHandler plugin which implements HttpProxy specifications."""
PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT = build_http_response(
200, reason=b'Connection established'
httpStatusCodes.OK,
reason=b'Connection established'
)
# Used to synchronize with other HttpProxyPlugin instances while
@ -1207,7 +1242,6 @@ class HttpProxyPlugin(ProtocolHandlerPlugin):
) and self.server and not self.server.closed and self.server.connection in r:
logger.debug('Server is ready for reads, reading')
raw = self.server.recv(self.config.server_recvbuf_size)
# self.last_activity = ProtocolHandler.now()
if not raw:
logger.debug('Server closed connection, tearing down...')
return True
@ -1393,6 +1427,9 @@ class HttpProxyPlugin(ProtocolHandlerPlugin):
self.client.connection.setblocking(False)
logger.info(
'TLS interception using %s', generated_cert)
# Update all plugin connection reference
for plugin in self.plugins.values():
plugin.client._conn = self.client.connection
return self.client.connection
elif self.server:
# - proxy-connection header is a mistake, it doesn't seem to be
@ -1859,13 +1896,15 @@ class HttpWebServerPlugin(ProtocolHandlerPlugin):
"""ProtocolHandler plugin which handles incoming requests to local web server."""
DEFAULT_404_RESPONSE = build_http_response(
404, reason=b'NOT FOUND',
httpStatusCodes.NOT_FOUND,
reason=b'NOT FOUND',
headers={b'Server': PROXY_AGENT_HEADER_VALUE,
b'Connection': b'close'}
)
DEFAULT_501_RESPONSE = build_http_response(
501, reason=b'NOT IMPLEMENTED',
httpStatusCodes.NOT_IMPLEMENTED,
reason=b'NOT IMPLEMENTED',
headers={b'Server': PROXY_AGENT_HEADER_VALUE,
b'Connection': b'close'}
)
@ -1900,10 +1939,12 @@ class HttpWebServerPlugin(ProtocolHandlerPlugin):
if content_type is None:
content_type = 'text/plain'
self.client.queue(build_http_response(
200, reason=b'OK', headers={
httpStatusCodes.OK,
reason=b'OK',
headers={
b'Content-Type': bytes_(content_type),
}, body=content
))
},
body=content))
return False
except IOError:
self.client.queue(self.DEFAULT_404_RESPONSE)
@ -2087,12 +2128,12 @@ class ProtocolHandler(threading.Thread):
return (self.now() - self.last_activity).seconds
def is_connection_inactive(self) -> bool:
# TODO: Add input argument option for timeout
return self.connection_inactive_for() > 30
return self.connection_inactive_for() > self.config.timeout
def handle_writables(self, writables: List[Union[int, _HasFileno]]) -> bool:
if self.client.buffer_size() > 0 and self.client.connection in writables:
logger.debug('Client is ready for writes, flushing buffer')
self.last_activity = self.now()
# Invoke plugin.on_response_chunk
chunk = self.client.buffer
@ -2112,8 +2153,9 @@ class ProtocolHandler(threading.Thread):
def handle_readables(self, readables: List[Union[int, _HasFileno]]) -> bool:
if self.client.connection in readables:
logger.debug('Client is ready for reads, reading')
client_data = self.client.recv(self.config.client_recvbuf_size)
self.last_activity = self.now()
client_data = self.client.recv(self.config.client_recvbuf_size)
if not client_data:
logger.debug('Client closed connection, tearing down...')
self.client.closed = True
@ -2209,12 +2251,19 @@ class ProtocolHandler(threading.Thread):
return True
# Teardown if client buffer is empty and connection is inactive
if self.client.buffer_size() == 0:
if self.is_connection_inactive():
logger.debug(
'Client buffer is empty and maximum inactivity has reached '
'between client and server connection, tearing down...')
return True
if not self.client.has_buffer() and \
self.is_connection_inactive():
self.client.queue(build_http_response(
httpStatusCodes.REQUEST_TIMEOUT, reason=b'Request Timeout',
headers={
b'Server': PROXY_AGENT_HEADER_VALUE,
b'Connection': b'close',
}
))
logger.debug(
'Client buffer is empty and maximum inactivity has reached '
'between client and server connection, tearing down...')
return True
return False
@ -2241,7 +2290,6 @@ class ProtocolHandler(threading.Thread):
if teardown:
return True
return False
def flush(self) -> None:
@ -2703,6 +2751,14 @@ def init_parser() -> argparse.ArgumentParser:
'This option is only applicable when static server is also enabled. '
'See --enable-static-server.'
)
parser.add_argument(
'--timeout',
type=int,
default=DEFAULT_TIMEOUT,
help='Default: ' + str(DEFAULT_TIMEOUT) + '. Number of seconds after which '
'an inactive connection must be dropped. Inactivity is defined by no '
'data sent or received by the client.'
)
parser.add_argument(
'--version',
'-v',
@ -2783,7 +2839,8 @@ def main(input_args: List[str]) -> None:
static_server_dir=args.static_server_dir,
enable_static_server=args.enable_static_server,
devtools_event_queue=devtools_event_queue,
devtools_ws_path=args.devtools_ws_path)
devtools_ws_path=args.devtools_ws_path,
timeout=args.timeout)
config.plugins = load_plugins(
bytes_(

7
tests.py
View File

@ -1662,7 +1662,7 @@ class TestHttpProxyTlsInterception(unittest.TestCase):
data=None), selectors.EVENT_READ)], ]
self.proxy.run_once()
# Assert our mocked plugin invocations
# Assert our mocked plugins invocations
self.plugin.return_value.get_descriptors.assert_called()
self.plugin.return_value.write_to_descriptors.assert_called_with([])
self.plugin.return_value.on_client_data.assert_called_with(connect_request)
@ -1698,8 +1698,7 @@ class TestHttpProxyTlsInterception(unittest.TestCase):
# Assert connection references for all other plugins is updated
self.assertEqual(self.plugin.return_value.client._conn, self.mock_ssl_wrap.return_value)
# Currently proxy doesn't update it's own plugin connection reference
# self.assertEqual(self.proxy_plugin.return_value.client._conn, self.mock_ssl_wrap.return_value)
self.assertEqual(self.proxy_plugin.return_value.client._conn, self.mock_ssl_wrap.return_value)
class TestHttpRequestRejected(unittest.TestCase):
@ -1761,6 +1760,7 @@ class TestMain(unittest.TestCase):
mock_args.enable_devtools = proxy.DEFAULT_ENABLE_DEVTOOLS
mock_args.devtools_event_queue = None
mock_args.devtools_ws_path = proxy.DEFAULT_DEVTOOLS_WS_PATH
mock_args.timeout = proxy.DEFAULT_TIMEOUT
@mock.patch('time.sleep')
@mock.patch('proxy.load_plugins')
@ -1817,6 +1817,7 @@ class TestMain(unittest.TestCase):
enable_static_server=mock_args.enable_static_server,
devtools_event_queue=None,
devtools_ws_path=proxy.DEFAULT_DEVTOOLS_WS_PATH,
timeout=proxy.DEFAULT_TIMEOUT
)
mock_acceptor_pool.assert_called_with(