Merged in test (pull request #10)

Test

README.md

@@ -10,7 +10,7 @@ pre version.
 
 ## Install
 
-    $ git clone http://xxxx
+    $ git clone https://github.com/jumpserver/coco.git
 
 ## Setting
 
@@ -31,7 +31,7 @@ Also some config you need kown:
 
 ## Start
 
-    # python ssh_server.py
+    # python run_server.py
 
 When your start ssh server, It will register with jumpserver api,
 

coco/app.py

@@ -150,10 +150,10 @@ class Coco:
                 for s in self.sessions:
                     if not s.stop_evt.is_set():
                         continue
-                    if s.date_finished is None:
+                    if s.date_end is None:
                         self.remove_session(s)
                         continue
-                    delta = datetime.datetime.now() - s.date_finished
+                    delta = datetime.datetime.now() - s.date_end
                     if delta > datetime.timedelta(seconds=interval*5):
                         self.remove_session(s)
                 time.sleep(interval)

coco/httpd.py

@@ -26,7 +26,6 @@ class BaseWebSocketHandler:
 
     def prepare(self, request):
         # self.app = self.settings["app"]
-        child, parent = socket.socketpair()
         if request.headers.getlist("X-Forwarded-For"):
             remote_ip = request.headers.getlist("X-Forwarded-For")[0]
         else:
@@ -36,10 +35,6 @@ class BaseWebSocketHandler:
         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.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 check_origin(self, origin):
         return True
@@ -64,9 +59,10 @@ class SSHws(Namespace, BaseWebSocketHandler):
             "cols": int(request.cookies.get('cols', 80)),
             "rows": int(request.cookies.get('rows', 24)),
             "room": room,
-            "chan": None,
-            "proxy": None,
-            "client": None,
+            # "chan": dict(),
+            "proxy": dict(),
+            "client": dict(),
+            "forwarder": dict(),
             "request": None,
         }
         self.rooms[room] = {
@@ -80,18 +76,31 @@ class SSHws(Namespace, BaseWebSocketHandler):
         self.prepare(request)
 
     def on_data(self, message):
-        if self.clients[request.sid]["proxy"]:
-            self.clients[request.sid]["proxy"].send({"data": message})
+        if message['room'] and self.clients[request.sid]["proxy"][message['room']]:
+            self.clients[request.sid]["proxy"][message['room']].send({"data": message['data']})
 
     def on_host(self, message):
         # 此处获取主机的信息
-        uuid = message.get('uuid', None)
+        connection = str(uuid.uuid4())
+        assetID = message.get('uuid', None)
         userid = message.get('userid', None)
-        if uuid and userid:
-            asset = self.app.service.get_asset(uuid)
+        self.emit('room', {'room': connection, 'secret': message['secret']})
+
+        if assetID and userid:
+            asset = self.app.service.get_asset(assetID)
             system_user = self.app.service.get_system_user(userid)
             if system_user:
-                self.socketio.start_background_task(self.clients[request.sid]["forwarder"].proxy, asset, system_user)
+
+                child, parent = socket.socketpair()
+                self.clients[request.sid]["client"][connection] = Client(parent, self.clients[request.sid]["request"])
+                self.clients[request.sid]["proxy"][connection] = WSProxy(self, child, self.clients[request.sid]["room"],
+                                                                         connection)
+                self.app.clients.append(self.clients[request.sid]["client"][connection])
+                self.clients[request.sid]["forwarder"][connection] = ProxyServer(self.app,
+                                                                                 self.clients[request.sid]["client"][connection])
+
+                self.socketio.start_background_task(self.clients[request.sid]["forwarder"][connection].proxy, asset,
+                                                    system_user)
                 # self.forwarder.proxy(self.asset, system_user)
             else:
                 self.on_disconnect()
@@ -125,13 +134,21 @@ class SSHws(Namespace, BaseWebSocketHandler):
     def on_disconnect(self):
         self.on_leave(self.clients[request.sid]["room"])
         try:
-            # todo: there maybe have bug
-            self.clients[request.sid]["proxy"].close()
+            del self.clients[request.sid]
         except:
             pass
         # self.ssh.close()
         pass
 
+    def on_logout(self, connection):
+        print("logout", connection)
+        if connection:
+            self.clients[request.sid]["proxy"][connection].close()
+            del self.clients[request.sid]["proxy"][connection]
+            del self.clients[request.sid]["forwarder"][connection]
+            self.clients[request.sid]["client"][connection].close()
+            del self.clients[request.sid]["client"][connection]
+
 
 class HttpServer:
     # prepare may be rewrite it
@@ -155,7 +172,7 @@ class HttpServer:
         port = self.app.config["HTTPD_PORT"]
         print('Starting websocket server at {}:{}'.format(host, port))
         self.socketio.on_namespace(SSHws('/ssh').app(self.app))
-        self.socketio.init_app(self.flask)
+        self.socketio.init_app(self.flask, async_mode="threading")
         self.socketio.run(self.flask, port=port, host=host)
 
     def shutdown(self):

coco/interactive.py

@@ -28,7 +28,7 @@ class InteractiveServer:
         self.client = client
         self.request = client.request
         self.assets = None
-        self.search_result = None
+        self._search_result = None
         self.asset_groups = None
         self.get_user_assets_async()
         self.get_user_asset_groups_async()
@@ -37,6 +37,18 @@ class InteractiveServer:
     def app(self):
         return self._app()
 
+    @property
+    def search_result(self):
+        if self._search_result:
+            return self._search_result
+        else:
+            return None
+
+    @search_result.setter
+    def search_result(self, value):
+        value = self.filter_system_users(value)
+        self._search_result = value
+
     def display_banner(self):
         self.client.send(char.CLEAR_CHAR)
         logo_path = os.path.join(self.app.root_path, "logo.txt")
@@ -219,14 +231,13 @@ class InteractiveServer:
     def filter_system_users(assets):
         for asset in assets:
             system_users_granted = asset.system_users_granted
-            high_priority = max([s.priority for s in system_users_granted])
+            high_priority = max([s.priority for s in system_users_granted]) if system_users_granted else 1
             system_users_cleaned = [s for s in system_users_granted if s.priority == high_priority]
             asset.system_users_granted = system_users_cleaned
         return assets
 
     def get_user_assets(self):
-        assets = self.app.service.get_user_assets(self.client.user)
-        self.assets = self.filter_system_users(assets)
+        self.assets = self.app.service.get_user_assets(self.client.user)
         logger.debug("Get user {} assets total: {}".format(self.client.user, len(self.assets)))
 
     def get_user_assets_async(self):

coco/models.py

@@ -186,7 +186,7 @@ class WSProxy:
     ```
     """
 
-    def __init__(self, ws, child, room):
+    def __init__(self, ws, child, room, connection):
         """
         :param ws: websocket instance or handler, have write_message method
         :param child: sock child pair
@@ -196,6 +196,7 @@ class WSProxy:
         self.stop_event = threading.Event()
         self.room = room
         self.auto_forward()
+        self.connection = connection
 
     def send(self, msg):
         """
@@ -215,7 +216,7 @@ class WSProxy:
             data = self.child.recv(BUF_SIZE)
             if len(data) == 0:
                 self.close()
-            self.ws.emit("data", data.decode("utf-8"), room=self.room)
+            self.ws.emit("data", {'data': data.decode("utf-8"), 'room': self.connection}, room=self.room)
 
     def auto_forward(self):
         thread = threading.Thread(target=self.forward, args=())
@@ -225,4 +226,3 @@ class WSProxy:
     def close(self):
         self.stop_event.set()
         self.child.close()
-        self.ws.on_disconnect()

coco/proxy.py

@@ -7,8 +7,8 @@ import threading
 import logging
 import time
 import weakref
-
 import paramiko
+from paramiko.ssh_exception import SSHException
 
 from .session import Session
 from .models import Server
@@ -130,7 +130,10 @@ class ProxyServer:
             width = self.request.meta.get('width', 80)
             height = self.request.meta.get('height', 24)
             logger.debug("Change win size: %s - %s" % (width, height))
+            try:
                 self.server.chan.resize_pty(width=width, height=height)
+            except SSHException:
+                break
 
     def watch_win_size_change_async(self):
         thread = threading.Thread(target=self.watch_win_size_change)

coco/sshd.py

@@ -9,6 +9,8 @@ import threading
 import paramiko
 import sys
 
+import time
+
 from .utils import ssh_key_gen
 from .interface import SSHInterface
 from .interactive import InteractiveServer
@@ -48,13 +50,13 @@ class SSHServer:
             try:
                 sock, addr = self.sock.accept()
                 logger.info("Get ssh request from {}: {}".format(addr[0], addr[1]))
-                thread = threading.Thread(target=self.handle, args=(sock, addr))
+                thread = threading.Thread(target=self.handle_connection, args=(sock, addr))
                 thread.daemon = True
                 thread.start()
             except Exception as e:
                 logger.error("Start SSH server error: {}".format(e))
 
-    def handle(self, sock, addr):
+    def handle_connection(self, sock, addr):
         transport = paramiko.Transport(sock, gss_kex=False)
         try:
             transport.load_server_moduli()
@@ -73,17 +75,23 @@ class SSHServer:
             logger.warning("Handle EOF Error")
             return
 
-        chan = transport.accept(10)
+        while True:
+            chan = transport.accept()
             if chan is None:
-            logger.warning("No ssh channel get")
-            return
-
+                continue
             server.event.wait(5)
             if not server.event.is_set():
                 logger.warning("Client not request a valid request, exiting")
                 return
 
+            t = threading.Thread(target=self.handle_chan, args=(chan, request))
+            t.daemon = True
+            t.start()
+
+    def handle_chan(self, chan, request):
         client = Client(chan, request)
+        print(chan)
+        print(request)
         self.app.add_client(client)
         self.dispatch(client)
 

coco/tasks.py

-# coding: utf-8
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
 import weakref
 import logging