Package restkit :: Package contrib :: Module wsgi_proxy
[hide private]
[frames] | no frames]

Source Code for Module restkit.contrib.wsgi_proxy

  1  # -*- coding: utf-8 - 
  2  # 
  3  # This file is part of restkit released under the MIT license. 
  4  # See the NOTICE for more information. 
  5   
  6  import urlparse 
  7  from restkit import request 
  8  from restkit.pool.simple import SimplePool 
  9  from restkit.util.sock import MAX_BODY 
 10   
 11  ALLOWED_METHODS = ['GET', 'HEAD', 'POST', 'PUT', 'DELETE'] 
 12   
 13  BLOCK_SIZE = 4096 * 16 
 14   
 15  WEBOB_ERROR = ("Content-Length is set to -1. This usually mean that WebOb has " 
 16          "already parsed the content body. You should set the Content-Length " 
 17          "header to the correct value before forwarding your request to the " 
 18          "proxy: ``req.content_length = str(len(req.body));`` " 
 19          "req.get_response(proxy)") 
 20   
21 -class ResponseIter(object):
22
23 - def __init__(self, response):
24 response.CHUNK_SIZE = BLOCK_SIZE 25 self.body = response.body_stream()
26
27 - def next(self):
28 data = self.body.read(BLOCK_SIZE) 29 if not data: 30 raise StopIteration 31 return data
32
33 - def __iter__(self):
34 return self
35
36 -class Proxy(object):
37 """A proxy wich redirect the request to SERVER_NAME:SERVER_PORT 38 and send HTTP_HOST header""" 39
40 - def __init__(self, pool=None, allowed_methods=ALLOWED_METHODS, 41 strip_script_name=True, **kwargs):
42 self.pool = pool or SimplePool() 43 self.allowed_methods = allowed_methods 44 self.strip_script_name = strip_script_name
45
46 - def extract_uri(self, environ):
47 port = None 48 scheme = environ['wsgi.url_scheme'] 49 if 'SERVER_NAME' in environ: 50 host = environ['SERVER_NAME'] 51 else: 52 host = environ['HTTP_HOST'] 53 if ':' in host: 54 host, port = host.split(':') 55 56 if not port: 57 if 'SERVER_PORT' in environ: 58 port = environ['SERVER_PORT'] 59 else: 60 port = scheme == 'https' and '443' or '80' 61 62 uri = '%s://%s:%s' % (scheme, host, port) 63 return uri
64
65 - def __call__(self, environ, start_response):
66 method = environ['REQUEST_METHOD'] 67 if method not in self.allowed_methods: 68 start_response('403 Forbidden', ()) 69 return [''] 70 71 if self.strip_script_name: 72 path_info = '' 73 else: 74 path_info = environ['SCRIPT_NAME'] 75 path_info += environ['PATH_INFO'] 76 77 query_string = environ['QUERY_STRING'] 78 if query_string: 79 path_info += '?' + query_string 80 81 host_uri = self.extract_uri(environ) 82 uri = host_uri + path_info 83 84 new_headers = {} 85 for k, v in environ.items(): 86 if k.startswith('HTTP_'): 87 k = k[5:].replace('_', '-').title() 88 new_headers[k] = v 89 90 for k, v in (('CONTENT_TYPE', None), ('CONTENT_LENGTH', '0')): 91 v = environ.get(k, None) 92 if v is not None: 93 new_headers[k.replace('_', '-').title()] = v 94 95 if new_headers.get('Content-Length', '0') == '-1': 96 raise ValueError(WEBOB_ERROR) 97 98 response = request(uri, method, 99 body=environ['wsgi.input'], headers=new_headers, 100 pool_instance=self.pool) 101 102 if 'location' in response: 103 headers = [] 104 for k, v in response.headerslist: 105 if k.lower() == 'location': 106 # rewrite location with a relative path. dont want to deal 107 # with complex url rebuild stuff 108 if v.startswith(host_uri): 109 v = v[len(host_uri):] 110 if self.strip_script_name: 111 v = environ['SCRIPT_NAME'] + v 112 headers.append((k, v)) 113 else: 114 headers = response.headerslist 115 116 start_response(response.status, headers) 117 118 if 'content-length' in response and \ 119 int(response['content-length']) <= MAX_BODY: 120 return [response.body] 121 122 return ResponseIter(response)
123
124 -class TransparentProxy(Proxy):
125 """A proxy based on HTTP_HOST environ variable""" 126
127 - def extract_uri(self, environ):
128 port = None 129 scheme = environ['wsgi.url_scheme'] 130 host = environ['HTTP_HOST'] 131 if ':' in host: 132 host, port = host.split(':') 133 134 if not port: 135 port = scheme == 'https' and '443' or '80' 136 137 uri = '%s://%s:%s' % (scheme, host, port) 138 return uri
139 140
141 -class HostProxy(Proxy):
142 """A proxy to redirect all request to a specific uri""" 143
144 - def __init__(self, uri, **kwargs):
145 super(HostProxy, self).__init__(**kwargs) 146 self.uri = uri.rstrip('/') 147 self.scheme, self.net_loc = urlparse.urlparse(self.uri)[0:2]
148
149 - def extract_uri(self, environ):
150 environ['HTTP_HOST'] = self.net_loc 151 return self.uri
152
153 -def get_config(local_config):
154 """parse paste config""" 155 config = {} 156 allowed_methods = local_config.get('allowed_methods', None) 157 if allowed_methods: 158 config['allowed_methods'] = [m.upper() for m in allowed_methods.split()] 159 strip_script_name = local_config.get('strip_script_name', 'true') 160 if strip_script_name.lower() in ('false', '0'): 161 config['strip_script_name'] = False 162 config['max_connections'] = int(local_config.get('max_connections', '5')) 163 return config
164
165 -def make_proxy(global_config, **local_config):
166 """TransparentProxy entry_point""" 167 config = get_config(local_config) 168 print 'Running TransparentProxy with %s' % config 169 return TransparentProxy(**config)
170
171 -def make_host_proxy(global_config, uri=None, **local_config):
172 """HostProxy entry_point""" 173 uri = uri.rstrip('/') 174 config = get_config(local_config) 175 print 'Running HostProxy on %s with %s' % (uri, config) 176 return HostProxy(uri, **config)
177