From 739f23da9e2ab0beb865c0284eec8577f00c93b2 Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Sat, 26 Jul 2014 09:48:44 -0700
Subject: [PATCH] https proxy support, CONNECT verb handling (uwsgi only)

---
 README.rst                      |  2 +-
 pywb/framework/proxy.py         | 72 ++++++++++++++++++++++++++++++++-
 pywb/framework/wsgi_wrappers.py | 33 +++++++++++++++
 setup.py                        |  2 +-
 4 files changed, 106 insertions(+), 3 deletions(-)

diff --git a/README.rst b/README.rst
index 9c4b380d..6aa256ac 100644
--- a/README.rst
+++ b/README.rst
@@ -1,4 +1,4 @@
-PyWb 0.5.1
+PyWb 0.5.2
 ==========
 
 .. image:: https://travis-ci.org/ikreymer/pywb.png?branch=develop
diff --git a/pywb/framework/proxy.py b/pywb/framework/proxy.py
index 62bc06b0..6754ecd7 100644
--- a/pywb/framework/proxy.py
+++ b/pywb/framework/proxy.py
@@ -62,12 +62,16 @@ class ProxyRouter(object):
         self.unaltered = proxy_options.get('unaltered_replay', False)
 
     def __call__(self, env):
+        if env['REQUEST_METHOD'] == 'CONNECT':
+            if not self.handle_connect(env):
+                return None
+
         url = env['REL_REQUEST_URI']
 
         if url.endswith('/proxy.pac'):
             return self.make_pac_response(env)
 
-        if not url.startswith('http://'):
+        if not url.startswith(('http://', 'https://')):
             return None
 
         proxy_auth = env.get('HTTP_PROXY_AUTHORIZATION')
@@ -122,6 +126,72 @@ class ProxyRouter(object):
 
         return route.handler(wbrequest)
 
+    def handle_connect(self, env):
+        import uwsgi
+        import socket
+        import ssl
+        from io import BytesIO
+
+        fd = uwsgi.connection_fd()
+        conn = socket.fromfd(fd, socket.AF_INET, socket.SOCK_STREAM)
+        sock = socket.socket(_sock=conn)
+
+        if (self.use_default_coll or
+            len(self.routes) == 1 or
+            env.get('HTTP_PROXY_AUTHORIZATION') is not None):
+
+            sock.send('HTTP/1.0 200 Connection Established\r\n')
+            sock.send('Server: pywb proxy\r\n')
+            sock.send('\r\n')
+        else:
+            env['pywb.proxy_statusline'] = '407 Proxy Auth Required'
+            sock.send('HTTP/1.0 407 Proxy Auth Required\r\n')
+            sock.send('Server: pywb proxy\r\n')
+            sock.send('\r\n')
+            return False
+
+        ssl_sock = ssl.wrap_socket(sock, server_side=True,
+                                   certfile='/tmp/testcert.pem',
+                                   ssl_version=ssl.PROTOCOL_SSLv23)
+
+        env['pywb.proxy_ssl_sock'] = ssl_sock
+
+        buff = ssl_sock.recv(4096)
+
+        buffreader = BytesIO(buff)
+
+        statusline = buffreader.readline()
+        statusparts = statusline.split(' ')
+
+        if len(statusparts) < 3:
+            return
+
+        env['REQUEST_METHOD'] = statusparts[0]
+        env['REL_REQUEST_URI'] = ('https://' +
+                                  env['REL_REQUEST_URI'].replace(':443', '') +
+                                  statusparts[1])
+
+        env['SERVER_PROTOCOL'] = statusparts[2].strip()
+
+        queryparts = env['REL_REQUEST_URI'].split('?', 1)
+        env['PATH_INFO'] = queryparts[0]
+        env['QUERY_STRING'] = queryparts[1] if len(queryparts) > 1 else ''
+
+        while True:
+            line = buffreader.readline()
+            if not line:
+                break
+
+            parts = line.split(':')
+            if len(parts) < 2:
+                continue
+
+            name = 'HTTP_' + parts[0].replace('-', '_').upper()
+            env[name] = parts[1]
+
+        return True
+
+
     # Proxy Auto-Config (PAC) script for the proxy
     def make_pac_response(self, env):
         import os
diff --git a/pywb/framework/wsgi_wrappers.py b/pywb/framework/wsgi_wrappers.py
index 3729a660..b40b5678 100644
--- a/pywb/framework/wsgi_wrappers.py
+++ b/pywb/framework/wsgi_wrappers.py
@@ -50,6 +50,39 @@ class WSGIApp(object):
 
     # Top-level wsgi application
     def __call__(self, env, start_response):
+        if env['REQUEST_METHOD'] == 'CONNECT':
+            return self.handle_connect(env, start_response)
+        else:
+            return self.handle_methods(env, start_response)
+
+    def handle_connect(self, env, start_response):
+        def ssl_start_response(statusline, headers):
+            ssl_sock = env.get('pywb.proxy_ssl_sock')
+            if not ssl_sock:
+                return
+
+            env['pywb.proxy_statusline'] = statusline
+
+            ssl_sock.write('HTTP/1.1 ' + statusline + '\r\n')
+            for name, value in headers:
+                ssl_sock.write(name + ': ' + value + '\r\n')
+
+        resp_iter = self.handle_methods(env, ssl_start_response)
+
+        ssl_sock = env.get('pywb.proxy_ssl_sock')
+        if ssl_sock:
+            ssl_sock.write('\r\n')
+
+            for obj in resp_iter:
+                ssl_sock.write(obj)
+
+            ssl_sock.close()
+
+        start_response(env['pywb.proxy_statusline'], [])
+
+        return []
+
+    def handle_methods(self, env, start_response):
         if env.get('SCRIPT_NAME') or not env.get('REQUEST_URI'):
             env['REL_REQUEST_URI'] = rel_request_uri(env)
         else:
diff --git a/setup.py b/setup.py
index 3e89abed..a6e9c885 100755
--- a/setup.py
+++ b/setup.py
@@ -34,7 +34,7 @@ class PyTest(TestCommand):
 
 setup(
     name='pywb',
-    version='0.5.1',
+    version='0.5.2',
     url='https://github.com/ikreymer/pywb',
     author='Ilya Kreymer',
     author_email='ikreymer@gmail.com',