__init__.py

import threading
import socket, array
import select
import logging
TES =  'http://localhost:8080/'
failthresh = 10
class TWSProxy(threading.Thread):

    TIMEOUT = 900
    AUTHTIMEOUT = 60
    MAXBUFF = 8192
    MAXHEADERS = 8192

    def __init__(self,socket):
        super().__init__()
        self.csock = socket
        self.ssock = None
        self.authtok = None
        self._timeout = 10



    def inittws(self,closed):
        """ Read from the client socket until one of
            1) We can verify the request based on headers
            2) We hit a timeout (i.e. no more headers will be sent and the client is waiting for a response)
            If we can verify 
            open a new socket to the server and return this
        """

        logger = logging.getLogger()
        bytessofar = 0
        header=bytearray(TWSProxy.MAXHEADERS)
        keepreading = True
        initcount=0
        while keepreading:
            initcount=initcount+1
            r,w,e = select.select([self.csock],[],[],5)
            if len(r) > 0:
                partial = self.csock.recv(TWSProxy.MAXBUFF)
                header[bytessofar:bytessofar+len(partial)] = partial
                bytessofar = bytessofar + len(partial)
                port = TWSProxy.verifyauth(header[0:bytessofar])
                if port is not None:
                    keepreading = False
            else:
                port = TWSProxy.verifyauth(header[0:bytessofar])
                if port is None:
                    logger.error('select returned no new bytes, verification failed, headers are {}'.format(header[0:bytessofar]))
                keepreading = False
            if initcount > failthresh:
                logger.error('Exceeded threhold, terminating proxy after {} bytes'.format(bytessofar))
                keepreading = False

        if port is not None:
            logger.debug('authenticated connection {}'.format(self.csock))
            self.ssock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.ssock.setblocking(True)
            try:
                self.ssock.connect(('127.0.0.1',port))
                if bytessofar > 0:
                    return (header,bytessofar)
                    #TWSProxy.reliablesend(self.ssock,header[0:bytessofar],bytessofar)
            except ConnectionRefusedError as e:
                logger.error('Couldn\'t open connection to the SSH Tunnel')
                self.ssock.close()
                self.csock.close()
                closed.set()
                return (None,None)
        else:
            logger.error('Failed to authenticate to the SSH Tunnel')
            self.csock.close()
            closed.set()
            return (None,None)

        

    def run(self):
        import logging
        logger = logging.getLogger()
        initshutdown = threading.Event()
        initshutdown.clear()
        (header, bytessofar) = self.inittws(initshutdown)

        t1 = None
        t2 = None
        if not initshutdown.isSet():
            TWSProxy.reliablesend(self.ssock,header[0:bytessofar],bytessofar)
#            TWSProxy.twosocks(self.csock,self.ssock,initshutdown)
            t1 = threading.Thread(target=TWSProxy.sockcopy, args=(self.ssock, self.csock, initshutdown),name='s2c')
            t2 = threading.Thread(target=TWSProxy.sockcopy, args=(self.csock, self.ssock, initshutdown),name='c2s')
            t1.start()
            t2.start()
            t1.join()
            t2.join()
        if self.ssock is not None:
            self.ssock.close()
        if self.csock is not None:
            self.csock.close()

    @staticmethod
    def verifyauth(header):
        # We are looking for either 
        # 1. An Authentication: token ... header that we can map to an ssh session
        # 2. the string token=....  (usually appearing as a url query parameter)
        # 3. A cookie called twsproxyauth that we can make to an ssh sesssion
        import re
        import requests
        logger = logging.getLogger()
        token_formats = [ b'Authorization: token (?P<authtok>\w+)[\W|$]',
                         b'token=(?P<authtok>\w+)[&|\W|$]',
                         b'twsproxyauth=(?P<authtok>\w+)[\W|$]']

        for token in token_formats:
            m = re.search(token,header)
            if m:
                try:
                    authtok = m.groupdict()['authtok'].rstrip()
                    s = requests.Session()
                    url = TES+'tunnelstat/'+authtok.decode()
                    try:
                        r = s.get(url)
                        port = r.json()
                    except:
                        raise Exception('unable to get a port number for the authtok {}'.format(r.text))
                    if port is not None:
                        return port
                except Exception as e:
                    import traceback
                    logger.error('Exception')
                    logger.error(e)
                    logger.error(traceback.format_exc())
                    raise e
        return None


    @staticmethod
    def reliablesend(socket,buff,msglength):
        totalsent = 0
        while totalsent < msglength:
            sent = socket.send(buff[totalsent:])
            if sent == 0:
                raise RuntimeError("socket connection broken")
            totalsent = totalsent + sent


    @staticmethod
    def sockcopy(src,dest,initshutdown):
        shuttype = socket.SHUT_RD
        import threading
        logger = logging.getLogger()
        closed = False
        name = threading.current_thread().name
        while not closed:
            r,w,e = select.select([src],[],[],TWSProxy.TIMEOUT)
            if len(r) > 0:
                buff = None
                msglength = -1
                try:
                    buff = src.recv(TWSProxy.MAXBUFF)
                    if buff is None:
                        continue
                except ConnectionResetError as e:
                    closed = True
                    continue
                except Exception as e:
                    import traceback
                    logger.error(traceback.format_exc())
                    closed = True
                msglength = len(buff)
                if msglength > 0:
                    TWSProxy.reliablesend(dest,buff,msglength)
                if msglength == 0:
                    #dest.send(buff)
                    dest.shutdown(shuttype)
                    initshutdown.set()
                    closed = True

        dest.shutdown(shuttype)
        initshutdown.set()

def mainserver(port=None):
    from . import server
    import logging
    import sys
    import os
    to_log=None
    try:
        from logging.handlers import TimedRotatingFileHandler
        logger = logging.getLogger()
        handler = TimedRotatingFileHandler(filename="/var/log/strudel2/tws.log",when='h',interval=24,backupCount=7)
        formatter = logging.Formatter("%(asctime)s %(levelname)s:%(process)s: %(message)s")
        handler.setFormatter(formatter)
        logger.addHandler(handler)
    except Exception as e:
        import traceback
        to_log = "{}\n".format(e)
        to_log = to_log + traceback.format_exc()
        logging.basicConfig(filename=os.path.expanduser("~/.tws.log"),format="%(asctime)s %(levelname)s:%(process)s: %(message)s")
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    if to_log is not None:
        logger.debug(to_log)
    logger.debug("starting TWS proxy")
    if port is None:
        try:
            port = int(sys.argv[1])
        except:
            port = 8090
    server = server.TWSServer(port,5)
    logger.debug("initialised server object")
    server.run()