#!/usr/bin/env python3

#
# @RDE_HEADER@
# $Id$
# $HeadURL$
# $Author$
# Created by Christian Fröbel, SLSE, 2016-07-07
#

import socket
from socketserver import ThreadingMixIn, TCPServer, StreamRequestHandler
import threading
from threading import Thread, Lock
from argparse import ArgumentParser
import time

# normale Farben
BLACK      = '\x1b[30m'
RED        = '\x1b[31m'
GREEN      = '\x1b[32m'
ORANGE     = '\x1b[33m'
BLUE       = '\x1b[34m'
PURPLE     = '\x1b[35m'
CYAN       = '\x1b[36m'
LIGHT_GREY = '\x1b[37m'

# intensive Farben
DARK_GREY    = '\x1b[30;1m'
LIGHT_RED    = '\x1b[31;1m'
LIGHT_GREEN  = '\x1b[32;1m'
YELLOW       = '\x1b[33;1m'
LIGHT_BLUE   = '\x1b[34;1m'
LIGHT_PURPLE = '\x1b[35;1m'
LIGHT_CYAN   = '\x1b[36;1m'
WHITE        = '\x1b[37;1m'

RESET   = '\x1b[0m'

IN_MARKER  = '{}<=={}'.format(LIGHT_PURPLE, RESET)
OUT_MARKER = '{}==>{}'.format(LIGHT_GREEN, RESET)

class _ThreadedTCPServer(ThreadingMixIn, TCPServer):
    allow_reuse_address = True


def _parseCmdLine():
    parser = ArgumentParser()
    parser.add_argument('--listen', '-l', type=int, required=True, help='Listening Port für den Client')
    parser.add_argument('--host', default='localhost', help='Server Host, default: localhost')
    parser.add_argument('--port', '-p', type=int, required=True, help='Server Port')
    parser.add_argument('--delay', '-d', type=int, default=0, help='Verzögerung von Requests und Responses in ms')
    return parser.parse_args()


class _Acceptor(object):
    def __init__(self):
        self._doRun = True
        self._handlers = []
        self._mutex = Lock()

    def run(self, connectionHandlerFactory):
        acceptor = _ThreadedTCPServer(('', options.listen), connectionHandlerFactory)
        acceptor.allow_reuse_address = True # entspricht socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) für den listening-Socket des TCPServers
        acceptor.timeout = 1

        while self._doRun:
            acceptor.handle_request() # mit timeout (siehe __init__)
            self._joinFinishedHandlers()

        self._joinAllHandlers()
    
    def initShutdown(self):
        self._doRun = False

    #
    # @param handler: Eine _ConnectionHandler-Instanz
    #
    def registerHandler(self, handler):
        with self._mutex:
            self._handlers.append(handler)

    def _joinFinishedHandlers(self):
        with self._mutex:
            for i in range(len(self._handlers)-1, -1, -1): # rückwärts durchlaufen, weil wir ggf. auch gleich das Element aus der Liste löschen wollen
                handler = self._handlers[i]
                handler.join(0) # ohne Timeout
                if not handler.isAlive():
                    del self._handlers[i]

    def _joinAllHandlers(self):
        # Ohne Mutex, weil zu diesem Zeitpunkt keine weiteren Handler mehr gestartet werden können.
        for handler in self._handlers:
            handler.initShutdown()
        for handler in self._handlers:
            handler.join()


#
# Die Factory brauche ich, damit ich dem _MyHandler-Konstruktor
# meine Parameter mitgeben kann. 
#
class _ConnectionHandlerFactory(object):

    def __init__(self, acceptor, options):
        self._acceptor = acceptor
        self._options = options

    def __call__(self, *args):
        return _MyHandler(self._acceptor, self._options, *args)


#
# Wird für jede neue Verbindung instantiiert.
#
class _MyHandler(StreamRequestHandler):
    def __init__(self, acceptor, options, *args):
        self._options = options
        self._delay = self._options.delay
        self._acceptor = acceptor
        self._thread = None
        StreamRequestHandler.__init__(self, *args)
    
    def handle(self):
        print('new connection: {}:{}'.format(self.client_address[0], self.client_address[1]))

        self._thread = threading.currentThread()
        self._acceptor.registerHandler(self)
        
        client = Client(self._options.host, self._options.port, self.wfile, self._options.delay)
        client.start()
        try:
            self._readAndForward(client)
        finally:
            client.initShutdown()
            client.join()
            print('connection closed: {}:{}'.format(self.client_address[0], self.client_address[1]))

    # Schließen des Sockets (der Leseseite), damit der Service aus dem blockierenden read() rauskommt. 
    def initShutdown(self):
        _shutdownAndCloseSocket(self.request)

    def join(self, timeout = None):
        assert self._thread is not None, 'Darf erst gerufen werden, nachdem der Thread gestartet wurde'
        self._thread.join(timeout)

    def isAlive(self):
        return self._thread.isAlive()

    #
    # Das Forwarding vom Client zum Server
    #
    def _readAndForward(self, client):
        while True:
            line = self.rfile.readline()
            if len(line) == 0:
                break # EOF
            print('{} {}'.format(OUT_MARKER, line))
            if self._delay > 0:
                time.sleep(self._delay/1000)
            client.send(line)


class Client(object):
    #
    # @param outFile: File-like, um an den Client des Proxys senden zu können.
    #
    def __init__(self, host, port, outFile, delay):
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._sock.connect((host, port))
        self._outFile = outFile
        self._delay = delay
        self._thread = Thread(target=self._readAndForward)
    
    def start(self):
        self._thread.start()
    
    def send(self, line):
        self._sock.sendall(line)

    def initShutdown(self):
        _shutdownAndCloseSocket(self._sock)
    
    def join(self):
        self._thread.join()
    
    #
    # Das Forwarding vom Server zum Client
    # Diese Methode wird in einem separaten Thread ausgeführt.
    # Sie liest die Zeilen vom Server und leitet sie den Client weiter.
    #
    def _readAndForward(self):
        inFile = self._sock.makefile('rb')
        while True:
            line = inFile.readline()
            if len(line) == 0:
                break # EOF
            print('{} {}'.format(IN_MARKER, line))
            if self._delay > 0:
                time.sleep(self._delay/1000)
            self._outFile.write(line)


def _shutdownAndCloseSocket(sock):
    try:
        sock.shutdown(socket.SHUT_RD) # wichtig, close() alleine reicht nicht
    except socket.error:
        print('socket error')
        # kann passieren, wenn der Socket schon geschlossen wurde
        pass
    sock.close()


if __name__ == '__main__':
    options = _parseCmdLine()
    acceptor = _Acceptor()
    connectionHandlerFactory = _ConnectionHandlerFactory(acceptor, options)
    t = Thread(target=acceptor.run, args=[connectionHandlerFactory])
    t.start()
    try:
        while True:
            t.join(0.5)
            if not t.is_alive():
                break
    except KeyboardInterrupt:
        acceptor.initShutdown()