minor cleanup

This commit is contained in:
Maximilian Hils 2016-09-22 01:58:01 -07:00
parent ec588f8822
commit e5b79a6d72
3 changed files with 11 additions and 17 deletions

View File

@ -163,8 +163,11 @@ class ServerConnectionMixin(object):
self.server_conn.close() self.server_conn.close()
self.channel.tell("serverdisconnect", self.server_conn) self.channel.tell("serverdisconnect", self.server_conn)
self.server_conn = models.ServerConnection(address, self.server_conn = models.ServerConnection(
(self.server_conn.source_address.host, 0), self.config.options.spoof_source_address) address,
(self.server_conn.source_address.host, 0),
self.config.options.spoof_source_address
)
def connect(self): def connect(self):
""" """

View File

@ -58,7 +58,3 @@ class InvalidCertificateException(TlsException):
class Timeout(TcpException): class Timeout(TcpException):
pass pass
class ProtocolException(NetlibException):
pass

View File

@ -484,14 +484,12 @@ class _Connection(object):
if not isinstance(self.connection, SSL.Connection): if not isinstance(self.connection, SSL.Connection):
if not getattr(self.wfile, "closed", False): if not getattr(self.wfile, "closed", False):
try: try:
if self.wfile: self.wfile.flush()
self.wfile.flush() self.wfile.close()
self.wfile.close()
except exceptions.TcpDisconnect: except exceptions.TcpDisconnect:
pass pass
if self.rfile: self.rfile.close()
self.rfile.close()
else: else:
try: try:
self.connection.shutdown() self.connection.shutdown()
@ -731,11 +729,7 @@ class TCPClient(_Connection):
def connect(self): def connect(self):
try: try:
# Allow the socket to be manipulated by using the server_conn stub. connection = socket.socket(self.address.family, socket.SOCK_STREAM)
if not self.connection:
connection = socket.socket(self.address.family, socket.SOCK_STREAM)
else:
connection = self.connection
if self.spoof_source_address: if self.spoof_source_address:
try: try:
@ -744,7 +738,8 @@ class TCPClient(_Connection):
connection.setsockopt(socket.SOL_IP, 19, 1) connection.setsockopt(socket.SOL_IP, 19, 1)
except socket.error as e: except socket.error as e:
raise exceptions.TcpException( raise exceptions.TcpException(
"Failed to spoof the source address: " + e.strerror) "Failed to spoof the source address: " + e.strerror
)
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())