Commit 96f121a6 authored by liuzheng712's avatar liuzheng712

fix: fix the mutil connection bug, but did not tested yet

parent e6d4a73b
......@@ -6,15 +6,14 @@ import os
import paramiko
import logging
import socket
from flask_socketio import SocketIO, Namespace, emit
from flask_socketio import SocketIO, Namespace, emit, join_room, leave_room
from flask import Flask, send_from_directory, render_template, request, jsonify
from urllib.parse import urlparse
import uuid
# Todo: Remove for future
from jms.models import User
from .models import Request, Client, WSProxy
from .forward import ProxyServer
import http.client
__version__ = '0.4.0'
BASE_DIR = os.path.dirname(os.path.dirname(__file__))
......@@ -23,6 +22,10 @@ logger = logging.getLogger(__file__)
class BaseWebSocketHandler:
def __init__(self, *args, **kwargs):
self.clients = dict()
super().__init__(*args, **kwargs)
def app(self, app):
self.app = app
return self
......@@ -34,13 +37,15 @@ class BaseWebSocketHandler:
remote_ip = request.headers.getlist("X-Forwarded-For")[0]
else:
remote_ip = request.remote_addr
self.request = Request((remote_ip, 0))
self.request.user = self.get_current_user()
self.request.meta = {"width": self.cols, "height": self.rows}
self.clients[request.sid]["request"] = Request((remote_ip, 0))
self.clients[request.sid]["request"].user = self.get_current_user()
self.clients[request.sid]["request"].meta = {"width": self.clients[request.sid]["cols"],
"height": self.clients[request.sid]["rows"]}
# self.request.__dict__.update(request.__dict__)
self.client = Client(parent, self.request)
self.proxy = WSProxy(self, child)
self.app.clients.append(self.client)
self.clients[request.sid]["client"] = Client(parent, self.clients[request.sid]["request"])
self.clients[request.sid]["proxy"] = WSProxy(self, child, self.clients[request.sid]["room"])
self.app.clients.append(self.clients[request.sid]["client"])
self.clients[request.sid]["forwarder"] = ProxyServer(self.app, self.clients[request.sid]["client"])
def get_current_user(self):
return User(id='61c39c1f5b5742688180b6dda235aadd', username="admin", name="admin")
......@@ -50,34 +55,29 @@ class BaseWebSocketHandler:
def close(self):
try:
self.ssh.close()
self.clients[request.sid]["client"].close()
except:
pass
pass
class SSHws(Namespace, BaseWebSocketHandler):
def ssh_with_password(self):
self.ssh = paramiko.SSHClient()
self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
self.ssh.connect("127.0.0.1", 22, "liuzheng", "liuzheng")
self.chan = self.ssh.invoke_shell(term='xterm', width=self.cols, height=self.rows)
self.socketio.start_background_task(self.send_data)
def send_data(self):
while True:
data = self.chan.recv(2048).decode('utf-8', 'replace')
self.emit('data', data)
def on_connect(self):
self.cols = int(request.cookies.get('cols', 80))
self.rows = int(request.cookies.get('rows', 24))
self.clients[request.sid] = {
"cols": int(request.cookies.get('cols', 80)),
"rows": int(request.cookies.get('rows', 24)),
"room": str(uuid.uuid4()),
"chan": None,
"proxy": None,
"client": None,
}
join_room(self.clients[request.sid]["room"])
self.prepare(request)
self.forwarder = ProxyServer(self.app, self.client)
def on_data(self, message):
self.proxy.send({"data": message})
if self.clients[request.sid]["proxy"]:
self.clients[request.sid]["proxy"].send({"data": message})
def on_host(self, message):
# 此处获取主机的信息
......@@ -85,11 +85,10 @@ class SSHws(Namespace, BaseWebSocketHandler):
uuid = message.get('uuid', None)
userid = message.get('userid', None)
if uuid and userid:
self.asset = self.app.service.get_asset(uuid)
asset = self.app.service.get_asset(uuid)
system_user = self.app.service.get_system_user(userid)
print(system_user)
if system_user:
self.socketio.start_background_task(self.forwarder.proxy, self.asset, system_user)
self.socketio.start_background_task(self.clients[request.sid]["forwarder"].proxy, asset, system_user)
# self.forwarder.proxy(self.asset, system_user)
else:
self.close()
......@@ -97,12 +96,12 @@ class SSHws(Namespace, BaseWebSocketHandler):
self.close()
def on_resize(self, message):
self.request.meta['width'] = message.get('cols', 80)
self.request.meta['height'] = message.get('rows', 24)
self.request.change_size_event.set()
self.clients[request.sid]["request"].meta['width'] = message.get('cols', 80)
self.clients[request.sid]["request"].meta['height'] = message.get('rows', 24)
self.clients[request.sid]["request"].change_size_event.set()
def on_disconnect(self):
self.proxy.close()
self.clients[request.sid]["proxy"].close()
# self.ssh.close()
pass
......
......@@ -191,7 +191,7 @@ class WSProxy:
```
"""
def __init__(self, ws, child):
def __init__(self, ws, child, room):
"""
:param ws: websocket instance or handler, have write_message method
:param child: sock child pair
......@@ -199,7 +199,7 @@ class WSProxy:
self.ws = ws
self.child = child
self.stop_event = threading.Event()
self.room = room
self.auto_forward()
def send(self, msg):
......@@ -220,7 +220,7 @@ class WSProxy:
data = self.child.recv(BUF_SIZE)
if len(data) == 0:
self.close()
self.ws.emit("data", data.decode("utf-8"))
self.ws.emit("data", data.decode("utf-8"), room=self.room)
def auto_forward(self):
thread = threading.Thread(target=self.forward, args=())
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment