Skip to content

C
coco
Commit 1a754877 authored by ibuler's avatar ibuler
Browse files
Options

[Update] 优化coco

parent cb161095
Show whitespace changes
Inline Side-by-side
Showing with 121 additions and 106 deletions
coco/httpd.py
View file @ 1a754877
...@@ -4,8 +4,10 @@ ...@@ -4,8 +4,10 @@
import os import os
import socket import socket
import uuid import uuid
import traceback
from flask_socketio import SocketIO, Namespace, join_room, leave_room from flask_socketio import SocketIO, Namespace, join_room, leave_room
from flask import Flask, request, current_app, redirect from flask import Flask, request, current_app, redirect
from copy import deepcopy
from .models import Request, Client, WSProxy from .models import Request, Client, WSProxy
from .proxy import ProxyServer from .proxy import ProxyServer
...@@ -52,70 +54,93 @@ class BaseNamespace(Namespace): ...@@ -52,70 +54,93 @@ class BaseNamespace(Namespace):
class ProxyNamespace(BaseNamespace): class ProxyNamespace(BaseNamespace):
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
"""
:param args:
:param kwargs:
self.connections = {
"request_sid": {
"room_id": {
"id": room_id,
"proxy": None,
"client": None,
"forwarder": None,
"request": None,
"cols": 80,
"rows": 24
},
...
},
...
}
"""
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
self.connections = dict() self.connections = dict()
self.rooms = dict()
def new_connection(self): def new_connection(self):
self.connections[request.sid] = dict()
def new_room(self):
room_id = str(uuid.uuid4()) room_id = str(uuid.uuid4())
connection = { room = {
"cols": int(request.cookies.get('cols', 80)), "id": room_id,
"rows": int(request.cookies.get('rows', 24)), "proxy": None,
"room": room_id, "client": None,
"proxy": dict(), "forwarder": None,
"client": dict(), "request": self.make_coco_request(),
"forwarder": dict(), "cols": 80,
"request": self.make_coco_request() "rows": 24
} }
return connection self.connections[request.sid][room_id] = room
return room
def make_coco_request(self): @staticmethod
x_forwarded_for = request.headers.get("X-Forwarded-For", '').split(',') def get_win_size():
if x_forwarded_for and x_forwarded_for[0]: cols_request = request.cookies.get('cols')
remote_ip = x_forwarded_for[0]
else:
remote_ip = request.remote_addr
width_request = request.cookies.get('cols')
rows_request = request.cookies.get('rows') rows_request = request.cookies.get('rows')
if width_request and width_request.isdigit(): if cols_request and cols_request.isdigit():
width = int(width_request) cols = int(cols_request)
else: else:
width = 80 cols = 80
if rows_request and rows_request.isdigit(): if rows_request and rows_request.isdigit():
rows = int(rows_request) rows = int(rows_request)
else: else:
rows = 24 rows = 24
return cols, rows
def make_coco_request(self):
x_forwarded_for = request.headers.get("X-Forwarded-For", '').split(',')
if x_forwarded_for and x_forwarded_for[0]:
remote_ip = x_forwarded_for[0]
else:
remote_ip = request.remote_addr
width, height = self.get_win_size()
req = Request((remote_ip, 0)) req = Request((remote_ip, 0))
req.user = self.current_user req.user = self.current_user
req.meta = { req.meta = {
"width": width, "width": width,
"height": rows, "height": height,
} }
return req return req
def on_connect(self): def on_connect(self):
logger.debug("On connect event trigger") logger.debug("On connect event trigger")
super().on_connect() super().on_connect()
connection = self.new_connection() self.new_connection()
self.connections[request.sid] = connection
self.rooms[connection['room']] = {
"admin": request.sid,
"member": [],
"rw": []
}
join_room(connection['room'])
def on_host(self, message): def on_host(self, message):
# 此处获取主机的信息 # 此处获取主机的信息
logger.debug("On host event trigger") logger.debug("On host event trigger")
room_id = str(uuid.uuid4())
asset_id = message.get('uuid', None) asset_id = message.get('uuid', None)
user_id = message.get('userid', None) user_id = message.get('userid', None)
secret = message.get('secret', None) secret = message.get('secret', None)
room = self.new_room()
self.emit('room', {'room': room_id, 'secret': secret}) self.emit('room', {'room': room["id"], 'secret': secret})
join_room(room["id"])
if not asset_id or not user_id: if not asset_id or not user_id:
# self.on_connect() # self.on_connect()
return return
...@@ -128,18 +153,14 @@ class ProxyNamespace(BaseNamespace): ...@@ -128,18 +153,14 @@ class ProxyNamespace(BaseNamespace):
return return
child, parent = socket.socketpair() child, parent = socket.socketpair()
self.connections[request.sid]["client"][room_id] = Client( client = Client(parent, room["request"])
parent, self.connections[request.sid]["request"] forwarder = ProxyServer(self.app, client)
) room["client"] = client
self.connections[request.sid]["proxy"][room_id] = WSProxy( room["forwarder"] = forwarder
self, child, self.connections[request.sid]["room"], room_id room["proxy"] = WSProxy(self, child, room["id"])
) room["cols"], room["rows"] = self.get_win_size()
self.connections[request.sid]["forwarder"][room_id] = ProxyServer(
self.app, self.connections[request.sid]["client"][room_id]
)
self.socketio.start_background_task( self.socketio.start_background_task(
self.connections[request.sid]["forwarder"][room_id].proxy, forwarder.proxy, asset, system_user
asset, system_user
) )
def on_data(self, message): def on_data(self, message):
...@@ -148,93 +169,84 @@ class ProxyNamespace(BaseNamespace): ...@@ -148,93 +169,84 @@ class ProxyNamespace(BaseNamespace):
:param message: {"data": "xxx", "room": "xxx"} :param message: {"data": "xxx", "room": "xxx"}
:return: :return:
""" """
room = message.get('room') room_id = message.get('room')
room = self.connections.get(request.sid, {}).get(room_id)
if not room: if not room:
return return
room_proxy = self.connections[request.sid]['proxy'].get(room) room["proxy"].send({"data": message['data']})
if room_proxy:
room_proxy.send({"data": message['data']})
def on_token(self, message): def on_token(self, message):
# 此处获取token含有的主机的信息 # 此处获取token含有的主机的信息
logger.debug("On token trigger") logger.debug("On token trigger")
token = message.get('token', None) token = message.get('token', None)
secret = message.get('secret', None) secret = message.get('secret', None)
connection = str(uuid.uuid4()) room = self.new_room()
self.emit('room', {'room': connection, 'secret': secret}) self.emit('room', {'room': room["id"], 'secret': secret})
if not (token or secret): if not token or not secret:
logger.debug("token or secret is None") logger.debug("Token or secret is None")
self.emit('data', {'data': "\nOperation not permitted!", 'room': connection}) self.emit('data', {'data': "\nOperation not permitted!",
'room': room["id"]})
self.emit('disconnect') self.emit('disconnect')
return None return None
info = self.app.service.get_token_asset(token) info = self.app.service.get_token_asset(token)
logger.debug(info) logger.debug(info)
if not info: if not info:
logger.debug("host is None") logger.debug("Token info is None")
self.emit('data', {'data': "\nOperation not permitted!", 'room': connection}) self.emit('data', {'data': "\nOperation not permitted!",
'room': room["id"]})
self.emit('disconnect') self.emit('disconnect')
return None return None
user_id = info.get('user', None) user_id = info.get('user', None)
logger.debug("self.current_user")
self.current_user = self.app.service.get_user_profile(user_id) self.current_user = self.app.service.get_user_profile(user_id)
self.connections[request.sid]["request"].user = self.current_user room["request"].user = self.current_user
logger.debug(self.current_user) logger.debug(self.current_user)
self.on_host({'secret': secret, 'uuid': info['asset'], 'userid': info['system_user']}) self.on_host({
'secret': secret,
'uuid': info['asset'],
'userid': info['system_user'],
})
def on_resize(self, message): def on_resize(self, message):
cols = message.get('cols') cols, rows = message.get('cols', None), message.get('rows', None)
rows = message.get('rows')
logger.debug("On resize event trigger: {}*{}".format(cols, rows)) logger.debug("On resize event trigger: {}*{}".format(cols, rows))
if cols and rows and self.connections[request.sid]["request"]: rooms = self.connections.get(request.sid)
self.connections[request.sid]["request"].meta['width'] = cols if not rooms:
self.connections[request.sid]["request"].meta['height'] = rows return
self.connections[request.sid]["request"].change_size_event.set() room = list(rooms.values())[0]
if rooms and (room["cols"], room["rows"]) != (cols, rows):
# def on_room(self, session_id): for room in rooms.values():
# logger.debug("On room event trigger") room["request"].meta.update({
# if session_id not in self.connections.keys(): 'width': cols, 'height': rows
# self.emit( })
# 'error', "no such session", room["request"].change_size_event.set()
# room=self.connections[request.sid]["room"] room.update({"cols": cols, "rows": rows})
# )
# else:
# self.emit(
# 'room', self.connections[session_id]["room"],
# room=self.connections[request.sid]["room"]
# )
#
# def on_join(self, room):
# logger.debug("On join room event trigger")
# self.on_leave(self.connections[request.id]["room"])
# self.connections[request.sid]["room"] = room
# self.rooms[room]["member"].append(request.sid)
# join_room(room=room)
#
def on_leave(self, room):
logger.debug("On leave room event trigger")
if self.rooms[room]["admin"] == request.sid:
self.emit("data", "\nAdmin leave", room=room)
del self.rooms[room]
leave_room(room=room)
def on_disconnect(self): def on_disconnect(self):
logger.debug("On disconnect event trigger") logger.debug("On disconnect event trigger")
# self.on_leave(self.clients[request.sid]["room"]) room_id_list = list(self.connections.get(request.sid, {}).keys())
print(self.connections[request.sid]["client"]) for room_id in room_id_list:
try: try:
for connection in self.connections[request.sid]["client"]: self.on_logout(room_id)
self.on_logout(connection) except Exception as e:
logger.warn(e)
del self.connections[request.sid] del self.connections[request.sid]
except IndexError:
pass
def on_logout(self, connection): def on_logout(self, room_id):
logger.debug("On logout event trigger") room = self.connections.get(request.sid, {}).get(room_id)
if connection in self.connections[request.sid]["proxy"].keys(): if room:
self.connections[request.sid]["proxy"][connection].close() room["proxy"].close()
del self.connections[request.sid]['proxy'][connection] self.close_room(room_id)
del self.connections[request.sid][room_id]
del room
@staticmethod
def on_error_default(e):
traceback.print_exc()
logger.warn(e)
error_handler = on_error_default
class HttpServer: class HttpServer:
...@@ -258,10 +270,14 @@ class HttpServer: ...@@ -258,10 +270,14 @@ class HttpServer:
self.flask_app.config.update(config) self.flask_app.config.update(config)
self.socket_io = SocketIO() self.socket_io = SocketIO()
self.register_routes() self.register_routes()
self.add_error_handler()
def register_routes(self): def register_routes(self):
self.socket_io.on_namespace(ProxyNamespace('/ssh')) self.socket_io.on_namespace(ProxyNamespace('/ssh'))
def add_error_handler(self):
self.socket_io.on_error_default(ProxyNamespace.on_error_default)
def run(self): def run(self):
host = self.flask_app.config["BIND_HOST"] host = self.flask_app.config["BIND_HOST"]
port = self.flask_app.config["HTTPD_PORT"] port = self.flask_app.config["HTTPD_PORT"]
......
coco/models.py
View file @ 1a754877
...@@ -218,7 +218,7 @@ class WSProxy: ...@@ -218,7 +218,7 @@ class WSProxy:
``` ```
""" """
def __init__(self, ws, child, room, connection): def __init__(self, ws, child, room_id):
""" """
:param ws: websocket instance or handler, have write_message method :param ws: websocket instance or handler, have write_message method
:param child: sock child pair :param child: sock child pair
...@@ -226,9 +226,8 @@ class WSProxy: ...@@ -226,9 +226,8 @@ class WSProxy:
self.ws = ws self.ws = ws
self.child = child self.child = child
self.stop_event = threading.Event() self.stop_event = threading.Event()
self.room = room self.room_id = room_id
self.auto_forward() self.auto_forward()
self.connection = connection
def send(self, msg): def send(self, msg):
""" """
...@@ -252,7 +251,8 @@ class WSProxy: ...@@ -252,7 +251,8 @@ class WSProxy:
if len(data) == 0: if len(data) == 0:
self.close() self.close()
data = data.decode(errors="ignore") data = data.decode(errors="ignore")
self.ws.emit("data", {'data': data, 'room': self.connection}, room=self.room) self.ws.emit("data", {'data': data, 'room': self.room_id},
room=self.room_id)
if len(data) == BUF_SIZE: if len(data) == BUF_SIZE:
time.sleep(0.1) time.sleep(0.1)
...@@ -265,7 +265,6 @@ class WSProxy: ...@@ -265,7 +265,6 @@ class WSProxy:
self.stop_event.set() self.stop_event.set()
self.child.shutdown(1) self.child.shutdown(1)
self.child.close() self.child.close()
self.ws.logout(self.connection)
logger.debug("Proxy {} closed".format(self)) logger.debug("Proxy {} closed".format(self))
......
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