diff --git a/scripts/proxy/cliside.py b/scripts/proxy/cliside.py index acc9180..ed5e0c1 100644 --- a/scripts/proxy/cliside.py +++ b/scripts/proxy/cliside.py @@ -4,11 +4,9 @@ import json import random import time import threading -import tornado.ioloop -import tornado.web -import tornado.gen -import tornado.websocket -import tornado.tcpserver +from tornado import ioloop +from tornado import gen +from tornado.websocket import websocket_connect from tornado.tcpclient import TCPClient class ClientSide: @@ -16,11 +14,22 @@ class ClientSide: def __init__(self, local_ip, remote_ip): self._local_ip = local_ip self._remote_ip = remote_ip + self.local_conn = None + self.remote_conn = None - @gen.coroutine - def co_connect(): - [host, port] = self._local_ip.split(':') - stream = yield TCPClient().connect(host, port) + @gen.coroutine + def co_connect(self): + while True: + if not self.local_conn: + [local_host, local_port] = self._local_ip.split(':') + 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) + remote_conn = yield websocket_connect(url) + + if self.local_conn: + pass def run(self): - ioloop.IOLoop.current().run_sync(self,connec) + ioloop.IOLoop.current().run_sync(self.co_connect) diff --git a/scripts/proxy/srvside.py b/scripts/proxy/srvside.py index 74d4964..f7ff080 100644 --- a/scripts/proxy/srvside.py +++ b/scripts/proxy/srvside.py @@ -14,6 +14,7 @@ app = None class ServerSide: def __init__(self, local_port, remote_port): + global app self._local_port = int(local_port) self._remote_port = int(remote_port) app = self @@ -131,6 +132,7 @@ class RemoteServerConnection(object): self.localSocket = None async def handle_stream(self): + global app while True: try: data = await self.stream.read_until(b"\n") @@ -151,6 +153,7 @@ class RemoteServerConnection(object): class RemoteServer(tornado.tcpserver.TCPServer): async def handle_stream(self, stream, address): + global app conn = RemoteServerConnection(steam, address) if not app.addRemoteConn(conn): stream.close() @@ -164,14 +167,17 @@ class RemoteServer(tornado.tcpserver.TCPServer): class LocalSocketHandler(tornado.websocket.WebSocketHandler): def open(self): + global app self._recvBuf = '' app.addLocalConn(self) def on_message(self, message): + global app self._recvBuf += message self.parsePacket() def on_close(self): + global app app.removeLocalConn(self) def parsePacket(self): @@ -186,6 +192,7 @@ class LocalSocketHandler(tornado.websocket.WebSocketHandler): self.dispatchMsg(msg) def dispatchMsg(self, msg): + global app if msg['cmd'] == 'connectOk': app.onLocalConnectOk(msg) elif msg['cmd'] == 'forwardData': @@ -196,4 +203,3 @@ class LocalSocketHandler(tornado.websocket.WebSocketHandler): def sendMsg(self, msg): data = json.dumps(msg) + '\n' self.write_message(data) -