diff --git a/TWS/twsproxy/__init__.py b/TWS/twsproxy/__init__.py index 3fbe1c340b17d567806f40bea799aa3245de21f6..0cc145742ba65d5bf53e6845f5a1b232047ec378 100644 --- a/TWS/twsproxy/__init__.py +++ b/TWS/twsproxy/__init__.py @@ -41,7 +41,7 @@ class TWSProxy(threading.Thread): header[bytessofar:bytessofar+len(partial)] = partial bytessofar = bytessofar + len(partial) logger.debug('inittws, checking headers') - port = TWSProxy.verifyauth(header[0:bytessofar]) + port = self.verifyauth(header[0:bytessofar]) if port is not None: logger.debug('inittws, found auth token and got port {}'.format(port)) keepreading = False @@ -49,7 +49,7 @@ class TWSProxy(threading.Thread): logger.debug('inittws, no authtok found in the first {} bytes'.format(bytessofar)) else: logger.debug('inittws, select returned with no more info, verifying headers for the last time') - port = TWSProxy.verifyauth(header[0:bytessofar]) + port = self.verifyauth(header[0:bytessofar]) keepreading = False if initcount > failthresh: logger.debug('inittws, checked headers enough times, got {} bytes with no success'.format(bytessofar)) @@ -77,6 +77,12 @@ class TWSProxy(threading.Thread): return (None,None) + def insert_header(self,buff,insertafter,extraheader): + try: + index = buff.decode().index(insertafter.decode()) + except: + return buff + return buff[0:index] + extraheader + buff[index:] def run(self): import logging @@ -85,6 +91,11 @@ class TWSProxy(threading.Thread): initshutdown = threading.Event() initshutdown.clear() (header, bytessofar) = self.inittws(initshutdown) + if self.authtok is not None: + replyheader = b'Set-Cookie: twsproxyauth2='+self.authtok+b' ; HttpOnly\r\n' + else: + replyheader = None + insertafter = b'HTTP/1.1 200 OK\r\n' logger.debug('connecting {} to {}'.format(self.csock,self.ssock)) if initshutdown.isSet(): logger.debug('NOT connecting {} inittws did not connect us'.format(self.csock)) @@ -97,7 +108,7 @@ class TWSProxy(threading.Thread): TWSProxy.reliablesend(self.ssock,header[0:bytessofar],bytessofar) logger.debug('creating threads') # TWSProxy.twosocks(self.csock,self.ssock,initshutdown) - t1 = threading.Thread(target=TWSProxy.sockcopy, args=(self.ssock, self.csock, initshutdown),name='s2c') + t1 = threading.Thread(target=TWSProxy.sockcopy, args=(self.ssock, self.csock, initshutdown,insertafter,replyheader),name='s2c') t2 = threading.Thread(target=TWSProxy.sockcopy, args=(self.csock, self.ssock, initshutdown),name='c2s') t1.start() t2.start() @@ -111,12 +122,11 @@ class TWSProxy(threading.Thread): if self.csock is not None: self.csock.close() - @staticmethod - def verifyauth(header): + def verifyauth(self,header): import re import requests logger = logging.getLogger() - token = b'twsproxyauth=(?P<authtok>\w+)[\W|$]' + token = b'twsproxyauth[0-9]*=(?P<authtok>\w+)[\W|$]' m = re.search(token,header) if m: authtok = m.groupdict()['authtok'] @@ -134,6 +144,7 @@ class TWSProxy(threading.Thread): except: logger.error('authtok found port found {}'.format(port)) raise Exception('unable to get a port number for the authtok {}'.format(r.text)) + self.authtok = authtok return port return None # if m: @@ -203,7 +214,7 @@ class TWSProxy(threading.Thread): @staticmethod - def sockcopy(src,dest,initshutdown): + def sockcopy(src,dest,initshutdown,insertafter=None,insert=None): shuttype = socket.SHUT_RD import threading logger = logging.getLogger() @@ -220,6 +231,14 @@ class TWSProxy(threading.Thread): msglength = -1 try: buff = src.recv(TWSProxy.MAXBUFF) + if insert is not None: + try: + index = buff.decode().index(insertafter.decode()) + newbuff = buff[0:index+len(insertafter)] + insert + buff[index+len(insertafter):] + buff = newbuff + insert = None + except: + pass if buff is None: continue except ConnectionResetError as e: