1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
|
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:
# We can't half-close a connection, so we just close everything here.
# Sockets will be cleaned up on a higher level.
return
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
|
