aboutsummaryrefslogtreecommitdiffstats
path: root/libmproxy/protocol/tcp.py
blob: a56bf07b37d78b11a754cc2606dc7b4e436bfe7a (plain)
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
from __future__ import absolute_import
import select
import socket
from .primitives import ProtocolHandler
from netlib.utils import cleanBin


class TCPHandler(ProtocolHandler):
    """
    TCPHandler acts as a generic TCP forwarder.
    Data will be .log()ed, but not stored any further.
    """

    chunk_size = 4096

    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:
                            # only ssl to dst, i.e. we revc'd into buf but need
                            # bytes/string now.
                            contents = buf[:size].tobytes()
                        self.c.log(
                            "%s %s\r\n%s" % (
                                direction, dst_str, cleanBin(contents)
                            ),
                            "debug"
                        )
                        dst.connection.send(contents)
                    else:
                        # socket.socket.send supports raw bytearrays/memoryviews
                        self.c.log(
                            "%s %s\r\n%s" % (
                                direction,
                                dst_str,
                                cleanBin(buf.tobytes())
                            ),
                            "debug"
                        )
                        dst.connection.send(buf[:size])
        except socket.error as e:
            self.c.log("TCP connection closed unexpectedly.", "debug")
            return