rwlaunchpad/plugins/rwimagemgr/rift/tasklets/rwimagemgr/lib/quickproxy/proxy.py

   1 import os
   2 import sys
   3 import urllib.parse as urlparse
   4 import pprint
   5 import http.cookies as Cookie
   6 import datetime
   7 import dateutil.parser
   8 from copy import copy
   9
  10 import tornado.httpserver
  11 import tornado.ioloop
  12 import tornado.iostream
  13 import tornado.web
  14 import tornado.httpclient
  15 import tornado.escape
  16
  17 __all__ = ['run_proxy', 'RequestObj', 'ResponseObj']
  18
  19 DEFAULT_CALLBACK = lambda r: r
  20
  21
  22 class Bunch(object):
  23     def __init__(self, **kwds):
  24         self.__dict__.update(kwds)
  25
  26     def __str__(self):
  27         return str(self.__dict__)
  28
  29
  30 class RequestObj(Bunch):
  31     '''
  32     An HTTP request object that contains the following request attributes:
  33
  34     protocol: either 'http' or 'https'
  35     host: the destination hostname of the request
  36     port: the port for the request
  37     path: the path of the request ('/index.html' for example)
  38     query: the query string ('?key=value&other=value')
  39     fragment: the hash fragment ('#fragment')
  40     method: request method ('GET', 'POST', etc)
  41     username: always passed as None, but you can set it to override the user
  42     password: None, but can be set to override the password
  43     body: request body as a string
  44     headers: a dictionary of header / value pairs
  45         (for example {'Content-Type': 'text/plain', 'Content-Length': 200})
  46     follow_redirects: true to follow redirects before returning a response
  47     validate_cert: false to turn off SSL cert validation
  48     context: a dictionary to place data that will be accessible to the response
  49     '''
  50     pass
  51
  52
  53 class ResponseObj(Bunch):
  54     '''
  55     An HTTP response object that contains the following request attributes:
  56
  57     code: response code, such as 200 for 'OK'
  58     headers: the response headers
  59     pass_headers: a list or set of headers to pass along in the response. All
  60         other headeres will be stripped out. By default this includes:
  61         ('Date', 'Cache-Control', 'Server', 'Content-Type', 'Location')
  62     body: response body as a string
  63     context: the context object from the request
  64     '''
  65
  66     def __init__(self, **kwargs):
  67         kwargs.setdefault('code', 200)
  68         kwargs.setdefault('headers', {})
  69         kwargs.setdefault('pass_headers', True)
  70         kwargs.setdefault('body', '')
  71         kwargs.setdefault('context', {})
  72         super(ResponseObj, self).__init__(**kwargs)
  73
  74
  75 class ResponseStreamWriterFuture(tornado.concurrent.Future):
  76     def __init__(self, write_fn, *args, **kwargs):
  77         self.write_fn = write_fn
  78         super().__init__(*args, **kwargs)
  79
  80
  81 def _make_proxy(methods, io_loop, req_callback, resp_callback, err_callback, debug_level=0):
  82
  83     @tornado.web.stream_request_body
  84     class ProxyHandler(tornado.web.RequestHandler):
  85
  86         SUPPORTED_METHODS = methods
  87
  88         def initialize(self):
  89             self.proxy_request_ready = tornado.concurrent.Future()
  90             self.request_future = None
  91
  92         def on_connection_close(self):
  93             if self.request_future is not None:
  94                 self.request_future.set_result(False)
  95
  96         def create_body_producer_future(self, write_fn):
  97             self.request_future = ResponseStreamWriterFuture(write_fn)
  98             self.proxy_request_ready.set_result(True)
  99             return self.request_future
 100
 101         @tornado.gen.coroutine
 102         def data_received(self, chunk):
 103             yield self.proxy_request_ready
 104
 105             yield self.request_future.write_fn(chunk)
 106
 107         def make_requestobj(self, request):
 108             '''
 109             creates a request object for this request
 110             '''
 111
 112             # get url for request
 113             # surprisingly, tornado's HTTPRequest sometimes
 114             # has a uri field with the full uri (http://...)
 115             # and sometimes it just contains the path. :(
 116
 117             url = request.uri
 118             if not url.startswith(u'http'):
 119                 url = u"{proto}://{netloc}{path}".format(
 120                     proto=request.protocol,
 121                     netloc=request.host,
 122                     path=request.uri
 123                 )
 124
 125             parsedurl = urlparse.urlparse(url)
 126
 127             # Passing on the transfer encoding header, causes Tornado to not
 128             # transmit valid chunks
 129             headers = request.headers.copy()
 130             if "Transfer-encoding" in headers:
 131                 del headers["Transfer-Encoding"]
 132
 133             # create request object
 134
 135             requestobj = RequestObj(
 136                 method=request.method,
 137                 protocol=parsedurl.scheme,
 138                 username=None,
 139                 password=None,
 140                 host=parsedurl.hostname,
 141                 port=parsedurl.port or 80,
 142                 path=parsedurl.path,
 143                 query=parsedurl.query,
 144                 fragment=parsedurl.fragment,
 145                 #body=request.body,
 146                 headers=headers,
 147                 follow_redirects=False,
 148                 validate_cert=True,
 149                 context={}
 150             )
 151
 152             return requestobj, parsedurl
 153
 154
 155         def make_request(self, obj, parsedurl):
 156             '''
 157             converts a request object into an HTTPRequest
 158             '''
 159
 160             obj.headers.setdefault('Host', obj.host)
 161
 162             if obj.username or parsedurl.username or \
 163                 obj.password or parsedurl.password:
 164
 165                 auth = u"{username}:{password}@".format(
 166                     username=obj.username or parsedurl.username,
 167                     password=obj.password or parsedurl.password
 168                 )
 169
 170             else:
 171                 auth = ''
 172
 173             url = u"{proto}://{auth}{host}{port}{path}{query}{frag}"
 174             url = url.format(
 175                 proto=obj.protocol,
 176                 auth=auth,
 177                 host=obj.host,
 178                 port=(u':' + str(obj.port)) if (obj.port and obj.port != 80) else u'',
 179                 path=u'/'+obj.path.lstrip(u'/') if obj.path else u'',
 180                 query=u'?'+obj.query.lstrip(u'?') if obj.query else u'',
 181                 frag=obj.fragment
 182             )
 183
 184             body_producer = None
 185             if "Transfer-encoding" in self.request.headers and \
 186                     self.request.headers["Transfer-Encoding"] == "chunked":
 187                 body_producer = self.create_body_producer_future
 188
 189             req = tornado.httpclient.HTTPRequest(
 190                 url=url,
 191                 method=obj.method,
 192                 body_producer=body_producer,
 193                 decompress_response=False,
 194                 headers=obj.headers,
 195                 follow_redirects=obj.follow_redirects,
 196                 allow_nonstandard_methods=True,
 197                 request_timeout=1*60*60 #1 hour
 198             )
 199
 200             return req
 201
 202         def prepare(self):
 203
 204             request = self.request
 205             if debug_level >= 4:
 206                 print("<<<<<<<< REQUEST <<<<<<<<")
 207                 pprint.pprint(request.__dict__)
 208
 209             MB = 1024 * 1024
 210             GB = 1024 * MB
 211
 212             MAX_STREAMED_SIZE = 50 * GB
 213             request.connection.set_max_body_size(MAX_STREAMED_SIZE)
 214
 215             requestobj, parsedurl = self.make_requestobj(request)
 216
 217             if debug_level >= 3:
 218                 print("<<<<<<<< REQUESTOBJ <<<<<<<<")
 219                 pprint.pprint(requestobj.__dict__)
 220
 221             if debug_level >= 1:
 222                 debugstr = "serving request from %s:%d%s " % (requestobj.host,
 223                                                               requestobj.port or 80,
 224                                                               requestobj.path)
 225
 226             modrequestobj = req_callback(requestobj)
 227
 228             if isinstance(modrequestobj, ResponseObj):
 229                 self.handle_response(modrequestobj)
 230                 return
 231
 232             if debug_level >= 1:
 233                 print(debugstr + "to %s:%d%s" % (modrequestobj.host,
 234                                                  modrequestobj.port or 80,
 235                                                  modrequestobj.path))
 236
 237             outreq = self.make_request(modrequestobj, parsedurl)
 238
 239             if debug_level >= 2:
 240                 print(">>>>>>>> REQUEST >>>>>>>>")
 241                 print("%s %s" % (outreq.method, outreq.url))
 242                 for k, v in outreq.headers.items():
 243                     print( "%s: %s" % (k, v))
 244
 245             # send the request
 246
 247             def _resp_callback(response):
 248                 self.handle_response(response, context=modrequestobj.context)
 249
 250             client = tornado.httpclient.AsyncHTTPClient(io_loop=io_loop)
 251             try:
 252                 client.fetch(outreq, _resp_callback,
 253                              validate_cert=modrequestobj.validate_cert)
 254             except tornado.httpclient.HTTPError as e:
 255                 if hasattr(e, 'response') and e.response:
 256                     self.handle_response(e.response,
 257                                          context=modrequestobj.context,
 258                                          error=True)
 259                 else:
 260                     self.set_status(500)
 261                     self.write('Internal server error:\n' + str(e))
 262                     self.finish()
 263
 264
 265         def handle_response(self, response, context={}, error=False):
 266
 267             if not isinstance(response, ResponseObj):
 268                 if debug_level >= 4:
 269                     print("<<<<<<<< RESPONSE <<<<<<<")
 270                     pprint.pprint(response.__dict__)
 271
 272                 responseobj = ResponseObj(
 273                     code=response.code,
 274                     headers=response.headers,
 275                     pass_headers=True,
 276                     body=response.body,
 277                     context=context,
 278                 )
 279             else:
 280                 responseobj = response
 281
 282             if debug_level >= 3:
 283                 print("<<<<<<<< RESPONSEOBJ <<<<<<<")
 284                 responseprint = copy(responseobj)
 285                 responseprint.body = "-- body content not displayed --"
 286                 pprint.pprint(responseprint.__dict__)
 287
 288             if not error:
 289                 mod = resp_callback(responseobj)
 290             else:
 291                 mod = err_callback(responseobj)
 292
 293             # set the response status code
 294
 295             if mod.code == 599:
 296                 self.set_status(500)
 297                 self.write('Internal server error. Server unreachable.')
 298                 self.finish()
 299                 return
 300
 301             self.set_status(mod.code)
 302
 303             # set the response headers
 304
 305             if type(mod.pass_headers) == bool:
 306                 header_keys = mod.headers.keys() if mod.pass_headers else []
 307             else:
 308                 header_keys = mod.pass_headers
 309             for key in header_keys:
 310                 if key.lower() == "set-cookie":
 311                     cookies = Cookie.BaseCookie()
 312                     cookies.load(tornado.escape.native_str(mod.headers.get(key)))
 313                     for cookie_key in cookies:
 314                         cookie = cookies[cookie_key]
 315                         params = dict(cookie)
 316                         expires = params.pop('expires', None)
 317                         if expires:
 318                             expires = dateutil.parser.parse(expires)
 319                         self.set_cookie(
 320                             cookie.key,
 321                             cookie.value,
 322                             expires = expires,
 323                             **params
 324                         )
 325                 else:
 326                     val = mod.headers.get(key)
 327                     self.set_header(key, val)
 328
 329             if debug_level >= 2:
 330                 print(">>>>>>>> RESPONSE (%s) >>>>>>>" % mod.code)
 331                 for k, v in self._headers.items():
 332                     print("%s: %s" % (k, v))
 333                 if hasattr(self, '_new_cookie'):
 334                     print(self._new_cookie.output())
 335
 336             # set the response body
 337
 338             if mod.body:
 339                 self.write(mod.body)
 340
 341             self.finish()
 342
 343         @tornado.web.asynchronous
 344         def get(self):
 345             pass
 346
 347         @tornado.web.asynchronous
 348         def options(self):
 349             pass
 350
 351         @tornado.web.asynchronous
 352         def head(self):
 353             pass
 354
 355         @tornado.web.asynchronous
 356         def put(self):
 357             self.request_future.set_result(True)
 358
 359         @tornado.web.asynchronous
 360         def patch(self):
 361             self.request_future.set_result(True)
 362
 363         @tornado.web.asynchronous
 364         def post(self):
 365             self.request_future.set_result(True)
 366
 367         @tornado.web.asynchronous
 368         def delete(self):
 369             pass
 370
 371
 372     return ProxyHandler
 373
 374
 375 def run_proxy(port,
 376               methods=['GET', 'POST', 'PUT', 'DELETE', 'HEAD'],
 377               req_callback=DEFAULT_CALLBACK,
 378               resp_callback=DEFAULT_CALLBACK,
 379               err_callback=DEFAULT_CALLBACK,
 380               test_ssl=False,
 381               debug_level=0,
 382               io_loop=None,
 383               ):
 384
 385     """
 386     Run proxy on the specified port.
 387
 388     methods: the HTTP methods this proxy will support
 389     req_callback: a callback that is passed a RequestObj that it should
 390         modify and then return
 391     resp_callback: a callback that is given a ResponseObj that it should
 392         modify and then return
 393     err_callback: in the case of an error, this callback will be called.
 394         there's no difference between how this and the resp_callback are
 395         used.
 396     test_ssl: if true, will wrap the socket in an self signed ssl cert
 397     start_ioloop: if True (default), the tornado IOLoop will be started
 398         immediately.
 399     debug_level: 0 no debug, 1 basic, 2 verbose
 400     """
 401
 402     io_loop = tornado.ioloop.IOLoop.instance() if io_loop is None else io_loop
 403
 404     app = tornado.web.Application([
 405         (r'.*', _make_proxy(methods=methods,
 406                             io_loop=io_loop,
 407                             req_callback=req_callback,
 408                             resp_callback=resp_callback,
 409                             err_callback=err_callback,
 410                             debug_level=debug_level)),
 411     ])
 412
 413     if test_ssl:
 414         this_dir, this_filename = os.path.split(__file__)
 415         kwargs = {
 416             "ssl_options": {
 417                 "certfile": os.path.join(this_dir, "data", "test.crt"),
 418                 "keyfile": os.path.join(this_dir, "data", "test.key"),
 419             },
 420             "io_loop": io_loop,
 421         }
 422     else:
 423         kwargs = {"io_loop": io_loop}
 424
 425     http_server = tornado.httpserver.HTTPServer(app, **kwargs)
 426     http_server.listen(port)
 427     return http_server
 428
 429
 430 if __name__ == '__main__':
 431     port = 8888
 432     if len(sys.argv) > 1:
 433         port = int(sys.argv[1])
 434
 435     print("Starting HTTP proxy on port %d" % port)
 436     run_proxy(port)