Thransparent TCP proxy and logger

pyspy.py

import argparse
import os
import select
import socket
import sys
import threading

parser = argparse.ArgumentParser(
    description='Transparent TCP proxy and logger')
parser.add_argument('-a', '--host', required=True,
                    help='address/host to connect')
parser.add_argument('-p', '--port', required=True, type=int,
                    help='remote port to connect')
parser.add_argument('-l', '--listen', required=True, type=int,
                    help='local port to listen')
parser.add_argument('-L', '--log', help='log file')
parser.add_argument('-c', action='store_true', default=False,
                    help='suppress console output')

def printable(ch):
    return (ch < 32 or ch > 126) and '.' or chr(ch)

printable_map = [printable(x) for x in range(256)]

class Spy:
    log_cond = threading.Condition()
    queue = []
    remote = None
    running = True

    def __init__(self, host, port, local_port, logfile, no_console):
        self.remote = (host, port)
        self.local = ('0.0.0.0', local_port)
        self.console = not no_console
        self.logfile = logfile

    def logger(self):
        while self.running:
            with self.log_cond:
                while len(self.queue) == 0:
                    self.log_cond.wait()
                if self.logfile:
                    try:
                        with open(self.logfile, 'a+') as logfile:
                            logfile.writelines(x + '\n' for x in self.queue)
                    except:
                        pass
                if self.console:
                    for line in self.queue:
                        print(line)
                self.queue = []

    def log(self, thread, msg):
        if self.console or self.logfile:
            with self.log_cond:
                self.queue.append('%04d: %s' % (thread, msg))
                self.log_cond.notify()

    def log_dump(self, thread, msg):
        if self.console or self.logfile:
            with self.log_cond:
                width = 16
                header = ''.join('%02X-' % x for x in range(width))[0:-1]
                self.queue.append("%04d: ----: %s" % (thread, header))
                self.queue.append("%04d:       %s" % (thread, '-' * width * 3))
                i = 0
                while True:
                    line = msg[i:i+width]
                    if len(line) == 0:
                        break
                    def conv(ch):
                        if isinstance(ch, int):
                            return ch
                        return ord(ch)
                    dump = ''.join('%02X ' % conv(x) for x in line)
                    char = ''.join(printable_map[conv(x)] for x in line)
                    self.queue.append("%04X: %04X: %-*s| %-*s" %
                                      (thread, i, width*3, dump, width, char))
                    i = i + width
                self.log_cond.notify()

    def spy_thread(self, local, addr, thread_id):
        self.log(thread_id, 'Thread started')
        try:
            self.log(thread_id, 'Connecting to %s...' % (self.remote,))
            remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            remote.connect(self.remote)
        except Exception as e:
            self.log(thread_id, 'Unable connect to %s -> %s'
                     % (self.remote, e))
            local.close()
            return
        self.log(thread_id, 'Remote host: %s' % (addr,))
        try:
            running = True
            while running:
                rd, _, er = select.select([local, remote], [],
                                          [local, remote], 3600)
                for sock in er:
                    if sock == local:
                        self.log(thread_id, 'Connection error from %s' %
                                 (addr,))
                        running = False
                    if sock == remote:
                        self.log(thread_id, 'Connection error from %s' %
                                 (self.remote,))
                        running = False
                def proxy(from_sock, to_sock, from_addr, to_addr):
                    val = from_sock.recv(1024)
                    if val:
                        self.log(thread_id, 'Recevied from %s (%d)' %
                                 (from_addr, len(val)))
                        self.log_dump(thread_id, val)
                        to_sock.send(val)
                        self.log(thread_id, 'Sent to %s (%d)' %
                                 (to_addr, len(val)))
                    else:
                        self.log(thread_id, 'Connection reset by %s' %
                                 (from_addr,))
                        return False
                    return True
                for sock in rd:
                    if sock == local:
                        running = proxy(local, remote, addr, self.remote)
                    if sock == remote:
                        running = proxy(remote, local, self.remote, addr)
        except Exception as e:
            self.log(thread_id, 'Connection terminated: ' + str(e))
        remote.close()
        local.close()
        self.log(thread_id, 'Connection close')
    def run(self):
        try:
            srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            srv.bind(self.local)
        except Exception as e:
            print('Error: ' + str(e))
            sys.exit(1)
        counter = 1
        threading.Thread(target=Spy.logger, args=[self]).start()
        self.log(0, 'Listen at port %s, remote host %s' %
                 (self.local, self.remote))
        try:
            while self.running:
                srv.listen(1)
                local, addr = srv.accept()
                self.log(0, 'Connection accepted from %s, thread %d launched' %
                         (addr, counter))
                threading.Thread(target=Spy.spy_thread,
                                 args=[self, local, addr, counter]).start()
                counter = counter + 1
        except KeyboardInterrupt:
            print('Keyboard interrupt')
            self.running = False
            try:
                sys.exit(1)
            except SystemExit:
                os._exit(1)
        except Exception as e:
            print('Error: ' + str(e))
            self.running = False
            sys.exit(1)

def main():
    args = parser.parse_args()
    obj = Spy(args.host, args.port, args.listen, args.log, args.c)
    obj.run()

if __name__ == '__main__':
    main()