diff --git a/scripts/proxy/cliside.py b/scripts/proxy/cliside.py index 9b0ed09..3a30e3b 100644 --- a/scripts/proxy/cliside.py +++ b/scripts/proxy/cliside.py @@ -3,6 +3,7 @@ import os import json import random import time +import base64 import threading from tornado import ioloop from tornado import gen @@ -17,30 +18,30 @@ class ClientSide: self.local_conn = None self.local_read_future = None self.remote_conn = None + self.conn_hash = {} @gen.coroutine - def co_connect2(self): - while True: - if not self.local_conn: - [local_host, local_port] = self._local_ip.split(':') - self.local_conn = yield TCPClient().connect(local_host, local_port) - if not self.remote_conn: - [remote_host, remote_port] = self._remote_ip.split(':') - url = 'ws://%s:%s/websocket' % (remote_host, remote_port) - self.remote_conn = yield websocket_connect(url) - - if self.local_conn and not self.local_read_future: - print('ok1') - self.local_read_future = self.local_conn.read_until(b"\n") - print('ok') - - if self.local_read_future: - self.local_read_future.fetch_next_chunk() - - @gen.coroutine - def co_localConnect(self): + def co_localConnect(self, idx): + print(idx, flush=True) [local_host, local_port] = self._local_ip.split(':') - self.local_conn = yield TCPClient().connect(local_host, local_port) + local_conn = yield TCPClient().connect(local_host, local_port) + self.conn_hash[idx] = local_conn + self.sendRemoteMsg({ + 'cmd' : 'connectOk', + 'remoteConnIdx' : idx, + }) + while True: + try: + data = yield local_conn.read_until_close() + if data: + self.sendRemoteMsg({ + 'cmd' : 'forwardData', + 'remoteConnIdx' : idx, + 'data' : str(base64.b64encode(data)) + }) + except: + print('qqqqq', flush=True) + break @gen.coroutine def co_remoteConnect(self): @@ -48,26 +49,39 @@ class ClientSide: url = 'ws://%s:%s/websocket' % (remote_host, remote_port) self.remote_conn = yield websocket_connect(url) data = '' - while True: - data += yield self.remote_conn.read_message() - if not data: - continue - lines = data.split('\n') - if data[-1] == '\n': - data = lines[-1] - lines = lines[:-1] - for line in lines: - print(line) - msg = json.loads(line) - self.dispatchRemoteMsg(msg) + try: + while True: + data += yield self.remote_conn.read_message() + print(data, flush=True) + if not data: + continue + print(data, flush=True) + lines = data.split('\n') + if data[-1] == '\n': + data = lines[-1] + lines = lines[:-1] + for line in lines: + msg = json.loads(line) + self.dispatchRemoteMsg(msg) + except: + print('wwwwwwwwww', flush=True) + def dispatchRemoteMsg(self, msg): if msg['cmd'] == 'connect': - ioloop.IOLoop.current().spawn_callback(self.co_localConnect) + ioloop.IOLoop.current().spawn_callback(self.co_localConnect, msg['remoteConnIdx']) elif msg['cmd'] == 'socketClose': pass elif msg['cmd'] == 'forwardData': - pass + if msg['remoteConnIdx'] in self.conn_hash: + conn = self.conn_hash[msg['remoteConnIdx']] + data = base64.b64decode(msg['data'][2:-1]) + print(data, flush=True) + conn.write(data) + + def sendRemoteMsg(self, msg): + data = json.dumps(msg) + '\n' + self.remote_conn.write_message(data) @gen.coroutine def co_connect(self): diff --git a/scripts/proxy/srvside.py b/scripts/proxy/srvside.py index 49dda5e..ffce7f8 100644 --- a/scripts/proxy/srvside.py +++ b/scripts/proxy/srvside.py @@ -3,6 +3,7 @@ import os import json import random import time +import base64 import threading import tornado.ioloop import tornado.web @@ -55,6 +56,7 @@ class ServerSide: 'remoteConnIdx' : conn.idx, 'cmd' : 'connect' }) + print('addRemoteConn') return True def addLocalConn(self, conn): @@ -83,7 +85,7 @@ class ServerSide: def removeRemoteConn(self, conn): self.remoteConnHashMutex.acquire() if conn.idx in self.remoteConnHash: - del self.remoteConnHash[conn] + del self.remoteConnHash[conn.idx] self.remoteConnHashMutex.release() def randLocalConn(self): @@ -110,11 +112,14 @@ class ServerSide: self.remoteConnHashMutex.release() def onLocalForwardData(self, msg): - self.remotePendingMutex.acquire() - if msg['remoteConnIdx'] in self.remotePending: - local_conn = self.remotePending[msg['remoteConnIdx']] - local_conn.stream.write(base64.b64decode(msg['data'])) - self.remotePendingMutex.release() +# print(msg) + self.remoteConnHashMutex.acquire() + if msg['remoteConnIdx'] in self.remoteConnHash: + local_conn = self.remoteConnHash[msg['remoteConnIdx']] + data = base64.b64decode(msg['data'][2:-1]) + print(data) + local_conn.stream.write(data) + self.remoteConnHashMutex.release() def isConnectOk(self, connIdx): isOk = False @@ -138,7 +143,8 @@ class RemoteServerConnection(object): data = await self.stream.read_until(b"\n") self.localSocket.sendMsg({ 'cmd' : 'forwardData', - 'data' : base64.b64encode(data) + 'remoteConnIdx' : self.idx, + 'data' : str(base64.b64encode(data)) }) except tornado.iostream.StreamClosedError: break @@ -153,6 +159,7 @@ class RemoteServerConnection(object): class RemoteServer(tornado.tcpserver.TCPServer): async def handle_stream(self, stream, address): + print('zzzz') global app conn = RemoteServerConnection(stream, address) if not app.addRemoteConn(conn): @@ -179,6 +186,7 @@ class LocalSocketHandler(tornado.websocket.WebSocketHandler): def on_close(self): global app app.removeLocalConn(self) + print('on_close') def parsePacket(self): if len(self._recvBuf) <= 0: @@ -196,10 +204,12 @@ class LocalSocketHandler(tornado.websocket.WebSocketHandler): if msg['cmd'] == 'connectOk': app.onLocalConnectOk(msg) elif msg['cmd'] == 'forwardData': - app.onLocalForward(msg) + app.onLocalForwardData(msg) elif msg['cmd'] == 'socketClose': - app.removeRemoteConn(msg['remoteConnIdx']) + pass +# app.removeRemoteConn(msg['remoteConnIdx']) def sendMsg(self, msg): data = json.dumps(msg) + '\n' + print(data) self.write_message(data)