Rooba
/
mitmproxy
mirror of https://github.com/mitmproxy/mitmproxy.git
1
1
Fork 0
Code Issues Packages Projects Releases Wiki Activity

remove code duplication

This commit is contained in:
Maximilian Hils 2016-02-01 20:19:34 +01:00
parent a3af0ce71d
commit 931b5459e9
1 changed files with 18 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
netlib/tcp.py
View File

@ -424,9 +424,10 @@ class _Connection(object):
rbufsize = -1 rbufsize = -1
wbufsize = -1 wbufsize = -1
def __init__(self, connection): def _makefile(self):
if connection: """
self.connection = connection Set up .rfile and .wfile attributes from .connection
"""
# Ideally, we would use the Buffered IO in Python 3 by default. # Ideally, we would use the Buffered IO in Python 3 by default.
# Unfortunately, the implementation of .peek() is broken for n>1 bytes, # Unfortunately, the implementation of .peek() is broken for n>1 bytes,
# as it may just return what's left in the buffer and not all the bytes we want. # as it may just return what's left in the buffer and not all the bytes we want.
@ -438,6 +439,11 @@ class _Connection(object):
else: else:
self.rfile = Reader(socket.SocketIO(self.connection, "rb")) self.rfile = Reader(socket.SocketIO(self.connection, "rb"))
self.wfile = Writer(socket.SocketIO(self.connection, "wb")) self.wfile = Writer(socket.SocketIO(self.connection, "wb"))
def __init__(self, connection):
if connection:
self.connection = connection
self._makefile()
else: else:
self.connection = None self.connection = None
self.rfile = None self.rfile = None
@ -676,20 +682,12 @@ class TCPClient(_Connection):
connection.connect(self.address()) connection.connect(self.address())
if not self.source_address: if not self.source_address:
self.source_address = Address(connection.getsockname()) self.source_address = Address(connection.getsockname())
# See _Connection.__init__ why we do this dance.
if six.PY2:
self.rfile = Reader(connection.makefile('rb', self.rbufsize))
self.wfile = Writer(connection.makefile('wb', self.wbufsize))
else:
self.rfile = Reader(socket.SocketIO(connection, "rb"))
self.wfile = Writer(socket.SocketIO(connection, "wb"))
except (socket.error, IOError) as err: except (socket.error, IOError) as err:
raise TcpException( raise TcpException(
'Error connecting to "%s": %s' % 'Error connecting to "%s": %s' %
(self.address.host, err)) (self.address.host, err))
self.connection = connection self.connection = connection
self._makefile()
def settimeout(self, n): def settimeout(self, n):
self.connection.settimeout(n) self.connection.settimeout(n)