Skip to content

Instantly share code, notes, and snippets.

@sampritipanda
Created July 13, 2018 12:33
Show Gist options
  • Save sampritipanda/f7450b9b80510d08ab8ac04af3113f5d to your computer and use it in GitHub Desktop.
Save sampritipanda/f7450b9b80510d08ab8ac04af3113f5d to your computer and use it in GitHub Desktop.
#!/usr/bin/env python
import logging
import socket
import urllib
import ipaddress
import tornado.ioloop
import tornado.iostream
import tornado.web
import tornado.httpclient
import tornado.httputil
logger = logging.getLogger('tornado_proxy')
def fetch_request(url, callback, **kwargs):
req = tornado.httpclient.HTTPRequest(url, **kwargs)
client = tornado.httpclient.AsyncHTTPClient()
client.fetch(req, callback, raise_error=False)
def resolve_host(host):
# Remove port from host if it exists
parts = host.split(":")
hostname = parts[0]
try:
host_ip = socket.gethostbyname(hostname)
logger.debug('Resolved %s to %s', hostname, host_ip)
ip_addr = ipaddress.IPv4Address(host_ip)
if not ip_addr.is_global:
return
netloc = host_ip
# If there was a port in the original hostname
if len(parts) > 1:
netloc += ':' + parts[1]
return netloc
except (socket.gaierror, UnicodeError, ipaddress.AddressValueError):
logger.debug('Failed to resolve %s', host)
return
class ProxyHandler(tornado.web.RequestHandler):
SUPPORTED_METHODS = tornado.web.RequestHandler.SUPPORTED_METHODS + ('CONNECT',)
def compute_etag(self):
return None
@tornado.web.asynchronous
def get(self):
logger.debug('Handle %s request to %s', self.request.method,
self.request.uri)
def handle_response(response):
if (response.error and not
isinstance(response.error, tornado.httpclient.HTTPError)):
self.set_status(500)
self.write('Internal server error:\n' + str(response.error))
else:
self.set_status(response.code, response.reason)
self._headers = response.headers
if response.body:
self.write(response.body)
self.finish()
body = self.request.body
if not body:
body = None
if 'Proxy-Connection' in self.request.headers:
del self.request.headers['Proxy-Connection']
if 'Connection' in self.request.headers:
del self.request.headers['Connection']
parsed_url = urllib.parse.urlparse(self.request.uri)
# The host header can be spoofed, so we get the host from the url itself.
actual_host = parsed_url.netloc
if actual_host != self.request.host:
self.set_status(400, 'Invalid Host Header')
self.finish()
return
resolved_netloc = resolve_host(self.request.host)
if not resolved_netloc:
self.set_status(404, 'Not Found')
self.finish()
return
self.request.headers.add('Connection', 'Close')
if 'Host' not in self.request.headers:
self.request.headers.add('Host', self.request.host)
final_url = urllib.parse.urlunparse((parsed_url.scheme, resolved_netloc,
parsed_url.path, parsed_url.params,
parsed_url.query, parsed_url.fragment))
fetch_request(
final_url, handle_response,
method=self.request.method, body=body,
headers=self.request.headers, follow_redirects=False,
allow_nonstandard_methods=True, decompress_response=False)
@tornado.web.asynchronous
def head(self):
return self.get()
@tornado.web.asynchronous
def post(self):
return self.get()
@tornado.web.asynchronous
def delete(self):
return self.get()
@tornado.web.asynchronous
def patch(self):
return self.get()
@tornado.web.asynchronous
def put(self):
return self.get()
@tornado.web.asynchronous
def options(self):
return self.get()
@tornado.web.asynchronous
def connect(self):
logger.debug('Start CONNECT to %s', self.request.uri)
host, port = self.request.uri.split(':')
client = self.request.connection.stream
def read_from_client(data):
upstream.write(data)
def read_from_upstream(data):
client.write(data)
def client_close(data=None):
if upstream.closed():
return
if data:
upstream.write(data)
upstream.close()
def upstream_close(data=None):
if client.closed():
return
if data:
client.write(data)
client.close()
def start_tunnel():
logger.debug('CONNECT tunnel established to %s', self.request.uri)
client.read_until_close(client_close, read_from_client)
upstream.read_until_close(upstream_close, read_from_upstream)
client.write(b'HTTP/1.0 200 Connection established\r\n\r\n')
host_ip = resolve_host(host)
if not host_ip:
self.set_status(404, 'Not Found')
self.finish()
return
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
upstream = tornado.iostream.IOStream(s)
upstream.connect((host_ip, int(port)), start_tunnel)
def run_proxy(port):
app = tornado.web.Application([
(r'.*', ProxyHandler),
])
app.listen(port)
logger.debug('Serving on http://localhost:' + str(port))
tornado.ioloop.IOLoop.current().start()
if __name__ == '__main__':
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - pid:%(process)d - %(message)s')
run_proxy(8899)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment