set an additional cookie which might fix a problem with safari

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: