95 lines
3.8 KiB
Python
95 lines
3.8 KiB
Python
from __future__ import absolute_import
|
|
import select
|
|
import socket
|
|
from .primitives import ProtocolHandler
|
|
from netlib.utils import cleanBin
|
|
from netlib.tcp import NetLibError
|
|
|
|
|
|
class TCPHandler(ProtocolHandler):
|
|
"""
|
|
TCPHandler acts as a generic TCP forwarder.
|
|
Data will be .log()ed, but not stored any further.
|
|
"""
|
|
|
|
chunk_size = 4096
|
|
|
|
def __init__(self, c, log=True):
|
|
super(TCPHandler, self).__init__(c)
|
|
self.log = log
|
|
|
|
def handle_messages(self):
|
|
self.c.establish_server_connection()
|
|
|
|
server = "%s:%s" % self.c.server_conn.address()[:2]
|
|
buf = memoryview(bytearray(self.chunk_size))
|
|
conns = [self.c.client_conn.rfile, self.c.server_conn.rfile]
|
|
|
|
try:
|
|
while True:
|
|
r, _, _ = select.select(conns, [], [], 10)
|
|
for rfile in r:
|
|
if self.c.client_conn.rfile == rfile:
|
|
src, dst = self.c.client_conn, self.c.server_conn
|
|
direction = "-> tcp ->"
|
|
src_str, dst_str = "client", server
|
|
else:
|
|
dst, src = self.c.client_conn, self.c.server_conn
|
|
direction = "<- tcp <-"
|
|
dst_str, src_str = "client", server
|
|
|
|
closed = False
|
|
if src.ssl_established:
|
|
# Unfortunately, pyOpenSSL lacks a recv_into function.
|
|
# We need to read a single byte before .pending()
|
|
# becomes usable
|
|
contents = src.rfile.read(1)
|
|
contents += src.rfile.read(src.connection.pending())
|
|
if not contents:
|
|
closed = True
|
|
else:
|
|
size = src.connection.recv_into(buf)
|
|
if not size:
|
|
closed = True
|
|
|
|
if closed:
|
|
conns.remove(src.rfile)
|
|
# Shutdown connection to the other peer
|
|
if dst.ssl_established:
|
|
dst.connection.shutdown()
|
|
else:
|
|
dst.connection.shutdown(socket.SHUT_WR)
|
|
|
|
if len(conns) == 0:
|
|
return
|
|
continue
|
|
|
|
if src.ssl_established or dst.ssl_established:
|
|
# if one of the peers is over SSL, we need to send
|
|
# bytes/strings
|
|
if not src.ssl_established:
|
|
# we revc'd into buf but need bytes/string now.
|
|
contents = buf[:size].tobytes()
|
|
if self.log:
|
|
self.c.log(
|
|
"%s %s\r\n%s" % (
|
|
direction, dst_str, cleanBin(contents)
|
|
),
|
|
"info"
|
|
)
|
|
# Do not use dst.connection.send here, which may raise OpenSSL-specific errors.
|
|
dst.send(contents)
|
|
else:
|
|
# socket.socket.send supports raw bytearrays/memoryviews
|
|
if self.log:
|
|
self.c.log(
|
|
"%s %s\r\n%s" % (
|
|
direction, dst_str, cleanBin(buf.tobytes())
|
|
),
|
|
"info"
|
|
)
|
|
dst.connection.send(buf[:size])
|
|
except (socket.error, NetLibError) as e:
|
|
self.c.log("TCP connection closed unexpectedly.", "debug")
|
|
return
|