1
2"""Extends standard socket with "dump traffic" feature"""
3
4import sys
5from socket import _realsocket
6
7LOG_FILE = sys.stderr
8
9class recv_buffer:
10    socket = None
11    buf = ''
12
13    @classmethod
14    def flush(cls):
15        if cls.socket and cls.buf:
16            cls.socket.log_recv(cls.buf)
17        cls.socket = None
18        cls.buf = ''
19
20    @classmethod
21    def append(cls, socket, data):
22        if socket is cls.socket:
23            cls.buf += data
24        else:
25            cls.flush()
26            cls.buf = data
27            cls.socket = socket
28
29flush = recv_buffer.flush
30
31class logging_socket(_realsocket):
32
33    def log_recv(self, data):
34        params = self.getsockname() + self.getpeername() + (len(data), data)
35        LOG_FILE.write('%s:%s <-- %s:%s [%s]\n%s\n' % params)
36
37    def __log_recv(self, data):
38        recv_buffer.append(self, data)
39
40    def __log_send(self, data):
41        flush()
42        params = self.getsockname() + self.getpeername() + (len(data), data)
43        LOG_FILE.write('%s:%s --> %s:%s [%s]\n%s\n' % params)
44
45    def close(self):
46        flush()
47        _realsocket.close(self)
48
49    def recv(self, bufsize, flags=0):
50        result = _realsocket.recv(self, bufsize)
51        self.__log_recv(result)
52        return result
53
54    def send(self, data, flags=0):
55        result = _realsocket.send(self, data, flags)
56        self.__log_send(data[:result])
57        return result
58
59    def sendall(self, data, flags=0):
60        result = _realsocket.sendall(self, data, flags)
61        self.__log_send(data)
62        return result
63    # XXX for tcp there's also recv_into
64    # QQQ for udp there're also recvfrom, recvfrom_into, sendto
65
66def _install():
67    import socket
68    socket._realsocket = logging_socket
69
70