Merge pull request #48 from mitmproxy/mitmproxy_issue_341

Adjust netlib to mitmproxy changes
This commit is contained in:
Aldo Cortesi 2014-09-07 12:47:49 +12:00
commit 754b627937
3 changed files with 39 additions and 31 deletions

View File

@ -216,10 +216,16 @@ class Address(object):
def use_ipv6(self, b): def use_ipv6(self, b):
self.family = socket.AF_INET6 if b else socket.AF_INET self.family = socket.AF_INET6 if b else socket.AF_INET
def __repr__(self):
return repr(self.address)
def __eq__(self, other): def __eq__(self, other):
other = Address.wrap(other) other = Address.wrap(other)
return (self.address, self.family) == (other.address, other.family) return (self.address, self.family) == (other.address, other.family)
def __ne__(self, other):
return not self.__eq__(other)
class _Connection(object): class _Connection(object):
def get_current_cipher(self): def get_current_cipher(self):
@ -313,6 +319,8 @@ class TCPClient(_Connection):
if self.source_address: if self.source_address:
connection.bind(self.source_address()) connection.bind(self.source_address())
connection.connect(self.address()) connection.connect(self.address())
if not self.source_address:
self.source_address = Address(connection.getsockname())
self.rfile = Reader(connection.makefile('rb', self.rbufsize)) self.rfile = Reader(connection.makefile('rb', self.rbufsize))
self.wfile = Writer(connection.makefile('wb', self.wbufsize)) self.wfile = Writer(connection.makefile('wb', self.wbufsize))
except (socket.error, IOError), err: except (socket.error, IOError), err:

View File

@ -9,15 +9,15 @@ class ClientConn:
class Flow: class Flow:
def __init__(self, client_conn): def __init__(self, address, request):
self.client_conn = client_conn self.client_conn = ClientConn(address)
self.request = request
class Request: class Request:
def __init__(self, client_conn, scheme, method, path, headers, content): def __init__(self, scheme, method, path, headers, content):
self.scheme, self.method, self.path = scheme, method, path self.scheme, self.method, self.path = scheme, method, path
self.headers, self.content = headers, content self.headers, self.content = headers, content
self.flow = Flow(client_conn)
def date_time_string(): def date_time_string():
@ -39,37 +39,37 @@ class WSGIAdaptor:
def __init__(self, app, domain, port, sversion): def __init__(self, app, domain, port, sversion):
self.app, self.domain, self.port, self.sversion = app, domain, port, sversion self.app, self.domain, self.port, self.sversion = app, domain, port, sversion
def make_environ(self, request, errsoc, **extra): def make_environ(self, flow, errsoc, **extra):
if '?' in request.path: if '?' in flow.request.path:
path_info, query = request.path.split('?', 1) path_info, query = flow.request.path.split('?', 1)
else: else:
path_info = request.path path_info = flow.request.path
query = '' query = ''
environ = { environ = {
'wsgi.version': (1, 0), 'wsgi.version': (1, 0),
'wsgi.url_scheme': request.scheme, 'wsgi.url_scheme': flow.request.scheme,
'wsgi.input': cStringIO.StringIO(request.content), 'wsgi.input': cStringIO.StringIO(flow.request.content),
'wsgi.errors': errsoc, 'wsgi.errors': errsoc,
'wsgi.multithread': True, 'wsgi.multithread': True,
'wsgi.multiprocess': False, 'wsgi.multiprocess': False,
'wsgi.run_once': False, 'wsgi.run_once': False,
'SERVER_SOFTWARE': self.sversion, 'SERVER_SOFTWARE': self.sversion,
'REQUEST_METHOD': request.method, 'REQUEST_METHOD': flow.request.method,
'SCRIPT_NAME': '', 'SCRIPT_NAME': '',
'PATH_INFO': urllib.unquote(path_info), 'PATH_INFO': urllib.unquote(path_info),
'QUERY_STRING': query, 'QUERY_STRING': query,
'CONTENT_TYPE': request.headers.get('Content-Type', [''])[0], 'CONTENT_TYPE': flow.request.headers.get('Content-Type', [''])[0],
'CONTENT_LENGTH': request.headers.get('Content-Length', [''])[0], 'CONTENT_LENGTH': flow.request.headers.get('Content-Length', [''])[0],
'SERVER_NAME': self.domain, 'SERVER_NAME': self.domain,
'SERVER_PORT': str(self.port), 'SERVER_PORT': str(self.port),
# FIXME: We need to pick up the protocol read from the request. # FIXME: We need to pick up the protocol read from the request.
'SERVER_PROTOCOL': "HTTP/1.1", 'SERVER_PROTOCOL': "HTTP/1.1",
} }
environ.update(extra) environ.update(extra)
if request.flow.client_conn.address: if flow.client_conn.address:
environ["REMOTE_ADDR"], environ["REMOTE_PORT"] = request.flow.client_conn.address() environ["REMOTE_ADDR"], environ["REMOTE_PORT"] = flow.client_conn.address()
for key, value in request.headers.items(): for key, value in flow.request.headers.items():
key = 'HTTP_' + key.upper().replace('-', '_') key = 'HTTP_' + key.upper().replace('-', '_')
if key not in ('HTTP_CONTENT_TYPE', 'HTTP_CONTENT_LENGTH'): if key not in ('HTTP_CONTENT_TYPE', 'HTTP_CONTENT_LENGTH'):
environ[key] = value environ[key] = value

View File

@ -2,11 +2,11 @@ import cStringIO, sys
from netlib import wsgi, odict from netlib import wsgi, odict
def treq(): def tflow():
cc = wsgi.ClientConn(("127.0.0.1", 8888))
h = odict.ODictCaseless() h = odict.ODictCaseless()
h["test"] = ["value"] h["test"] = ["value"]
return wsgi.Request(cc, "http", "GET", "/", h, "") req = wsgi.Request("http", "GET", "/", h, "")
return wsgi.Flow(("127.0.0.1", 8888), req)
class TestApp: class TestApp:
@ -24,22 +24,22 @@ class TestApp:
class TestWSGI: class TestWSGI:
def test_make_environ(self): def test_make_environ(self):
w = wsgi.WSGIAdaptor(None, "foo", 80, "version") w = wsgi.WSGIAdaptor(None, "foo", 80, "version")
tr = treq() tf = tflow()
assert w.make_environ(tr, None) assert w.make_environ(tf, None)
tr.path = "/foo?bar=voing" tf.request.path = "/foo?bar=voing"
r = w.make_environ(tr, None) r = w.make_environ(tf, None)
assert r["QUERY_STRING"] == "bar=voing" assert r["QUERY_STRING"] == "bar=voing"
def test_serve(self): def test_serve(self):
ta = TestApp() ta = TestApp()
w = wsgi.WSGIAdaptor(ta, "foo", 80, "version") w = wsgi.WSGIAdaptor(ta, "foo", 80, "version")
r = treq() f = tflow()
r.host = "foo" f.request.host = "foo"
r.port = 80 f.request.port = 80
wfile = cStringIO.StringIO() wfile = cStringIO.StringIO()
err = w.serve(r, wfile) err = w.serve(f, wfile)
assert ta.called assert ta.called
assert not err assert not err
@ -49,11 +49,11 @@ class TestWSGI:
def _serve(self, app): def _serve(self, app):
w = wsgi.WSGIAdaptor(app, "foo", 80, "version") w = wsgi.WSGIAdaptor(app, "foo", 80, "version")
r = treq() f = tflow()
r.host = "foo" f.request.host = "foo"
r.port = 80 f.request.port = 80
wfile = cStringIO.StringIO() wfile = cStringIO.StringIO()
err = w.serve(r, wfile) err = w.serve(f, wfile)
return wfile.getvalue() return wfile.getvalue()
def test_serve_empty_body(self): def test_serve_empty_body(self):