Merged in test (pull request #19)

Test

coco/httpd.py

@@ -3,12 +3,10 @@
 #
 import os
 import socket
-from flask_socketio import SocketIO, Namespace, emit, join_room, leave_room
-from flask import Flask, send_from_directory, render_template, request, jsonify
 import uuid
+from flask_socketio import SocketIO, Namespace, join_room, leave_room
+from flask import Flask, request, current_app, redirect
 
-# Todo: Remove for future
-from jms.models import User
 from .models import Request, Client, WSProxy
 from .proxy import ProxyServer
 from .utils import get_logger
@@ -19,31 +17,31 @@ BASE_DIR = os.path.dirname(os.path.dirname(__file__))
 logger = get_logger(__file__)
 
 
-class BaseWebSocketHandler:
+class BaseNamespace(Namespace):
     clients = None
     current_user = None
 
-    def app(self, app):
-        self.app = app
-        return self
+    @property
+    def app(self):
+        app = current_app.config['coco']
+        return app
 
-    def prepare(self, request):
-        # self.app = self.settings["app"]
-        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
-        req = Request((remote_ip, 0))
-        req.user = self.current_user
-        req.meta = {
-            "width": self.clients[request.sid]["cols"],
-            "height": self.clients[request.sid]["rows"]
-        }
-        self.clients[request.sid]["request"] = req
+    def on_connect(self):
+        self.current_user = self.get_current_user()
+        if self.current_user is None:
+            return redirect(current_app.config['LOGIN_URL'])
+        logger.debug("{} connect websocket".format(self.current_user))
 
-    def check_origin(self, origin):
-        return True
+    def get_current_user(self):
+        session_id = request.cookies.get('sessionid', '')
+        csrf_token = request.cookies.get('csrftoken', '')
+        token = request.headers.get("Authorization")
+        user = None
+        if session_id and csrf_token:
+            user = self.app.service.check_user_cookie(session_id, csrf_token)
+        if token:
+            user = self.app.service.check_user_with_token(token)
+        return user
 
     def close(self):
         try:
@@ -52,99 +50,164 @@ class BaseWebSocketHandler:
             pass
 
 
-class SSHws(Namespace, BaseWebSocketHandler):
+class ProxyNamespace(BaseNamespace):
     def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
         self.clients = dict()
         self.rooms = dict()
-        super().__init__(*args, **kwargs)
 
-    def on_connect(self):
+    def new_client(self):
         room = str(uuid.uuid4())
-        self.clients[request.sid] = {
+        client = {
             "cols": int(request.cookies.get('cols', 80)),
             "rows": int(request.cookies.get('rows', 24)),
             "room": room,
-            # "chan": dict(),
             "proxy": dict(),
             "client": dict(),
             "forwarder": dict(),
-            "request": None,
+            "request": self.make_coco_request()
         }
-        self.rooms[room] = {
+        return client
+
+    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_request = request.cookies.get('cols')
+        rows_request = request.cookies.get('rows')
+        if width_request and width_request.isdigit():
+            width = int(width_request)
+        else:
+            width = 80
+
+        if rows_request and rows_request.isdigit():
+            rows = int(rows_request)
+        else:
+            rows = 24
+        req = Request((remote_ip, 0))
+        req.user = self.current_user
+        req.meta = {
+            "width": width,
+            "height": rows,
+        }
+        return req
+
+    def on_connect(self):
+        logger.debug("On connect event trigger")
+        super().on_connect()
+        client = self.new_client()
+        self.clients[request.sid] = client
+        self.rooms[client['room']] = {
             "admin": request.sid,
             "member": [],
             "rw": []
         }
-        join_room(room)
-        self.current_user = self.app.service.check_user_cookie(
-            session_id=request.cookies.get('sessionid', ''),
-            csrf_token=request.cookies.get('csrftoken', '')
-        )
-        self.prepare(request)
+        join_room(client['room'])
 
     def on_data(self, message):
-        if message['room'] and self.clients[request.sid]["proxy"][message['room']]:
-            self.clients[request.sid]["proxy"][message['room']].send({"data": message['data']})
+        """
+        收到浏览器请求
+        :param message: {"data": "xxx", "room": "xxx"}
+        :return:
+        """
+        room = message.get('room')
+        if not room:
+            return
+        room_proxy = self.clients[request.sid]['proxy'].get(room)
+        if room_proxy:
+            room_proxy.send({"data": message['data']})
 
     def on_host(self, message):
         # 此处获取主机的信息
+        logger.debug("On host event trigger")
         connection = str(uuid.uuid4())
         asset_id = message.get('uuid', None)
         user_id = message.get('userid', None)
-        self.emit('room', {'room': connection, 'secret': message['secret']})
-
-        if asset_id and user_id:
-            asset = self.app.service.get_asset(asset_id)
-            system_user = self.app.service.get_system_user(user_id)
-
-            if 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.clients[request.sid]["forwarder"][
-                    connection] = ProxyServer(
-                    self.app, self.clients[request.sid]["client"][connection]
-                )
-                self.app.clients.append(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()
-        else:
-            self.on_disconnect()
+        secret = message.get('secret', None)
+
+        self.emit('room', {'room': connection, 'secret': secret})
+
+        if not asset_id or not user_id:
+            # self.on_connect()
+            return
+
+        asset = self.app.service.get_asset(asset_id)
+        system_user = self.app.service.get_system_user(user_id)
+
+        if not asset or not system_user:
+            self.on_connect()
+            return
+
+        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.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
+        )
+
+    def on_token(self, message):
+        # 此处获取token含有的主机的信息
+        logger.debug("On token trigger")
+        token = message.get('token', None)
+        secret = message.get('secret', None)
+        host = self.app.service.get_token_asset(token).json()
+        logger.debug(host)
+        # {
+        #     "user": {UUID},
+        #     "asset": {UUID},
+        #     "system_user": {UUID}
+        # }
+        self.on_host({'secret': secret, 'uuid': host['asset'], 'userid': host['system_user']})
 
     def on_resize(self, message):
-        if self.clients[request.sid]["request"]:
-            self.clients[request.sid]["request"].meta['width'] = message.get('cols', 80)
-            self.clients[request.sid]["request"].meta['height'] = message.get('rows', 24)
+        cols = message.get('cols')
+        rows = message.get('rows')
+        logger.debug("On resize event trigger: {}*{}".format(cols, rows))
+        if cols and rows and self.clients[request.sid]["request"]:
+            self.clients[request.sid]["request"].meta['width'] = cols
+            self.clients[request.sid]["request"].meta['height'] = rows
             self.clients[request.sid]["request"].change_size_event.set()
 
-    def on_room(self, sessionid):
-        if sessionid not in self.clients.keys():
-            self.emit('error', "no such session", room=self.clients[request.sid]["room"])
+    def on_room(self, session_id):
+        logger.debug("On room event trigger")
+        if session_id not in self.clients.keys():
+            self.emit(
+                'error', "no such session",
+                room=self.clients[request.sid]["room"]
+            )
         else:
-            self.emit('room', self.clients[sessionid]["room"], room=self.clients[request.sid]["room"])
+            self.emit(
+                'room', self.clients[session_id]["room"],
+                room=self.clients[request.sid]["room"]
+            )
 
     def on_join(self, room):
+        logger.debug("On join room event trigger")
         self.on_leave(self.clients[request.id]["room"])
         self.clients[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):
+        logger.debug("On disconnect event trigger")
         self.on_leave(self.clients[request.sid]["room"])
         try:
             for connection in self.clients[request.sid]["client"]:
@@ -154,10 +217,11 @@ class SSHws(Namespace, BaseWebSocketHandler):
             pass
 
     def on_logout(self, connection):
-        logger.debug("{} logout".format(connection))
+        logger.debug("On logout event trigger")
         if connection:
             if connection in self.clients[request.sid]["proxy"].keys():
                 self.clients[request.sid]["proxy"][connection].close()
+                del self.clients[request.sid]['proxy'][connection]
 
     def logout(self, connection):
         if connection and (request.sid in self.clients.keys()):
@@ -171,28 +235,31 @@ class SSHws(Namespace, BaseWebSocketHandler):
 
 class HttpServer:
     # prepare may be rewrite it
-    settings = {
-        'cookie_secret': '',
-        'app': None,
-        'login_url': '/login'
+    config = {
+        'SECRET_KEY': '',
+        'coco': None,
+        'LOGIN_URL': '/login'
     }
+    async_mode = "threading"
 
-    def __init__(self, app):
-        self.app = app
-        # self.settings['cookie_secret'] = self.app.config['SECRET_KEY']
-        # self.settings['app'] = self.app
+    def __init__(self, coco):
+        config = coco.config
+        config.update(self.config)
+        config['coco'] = coco
+        self.flask_app = Flask(__name__, template_folder='dist')
+        self.flask_app.config.update(config)
+        self.socket_io = SocketIO()
+        self.register_routes()
 
-        self.flask = Flask(__name__, template_folder='dist')
-        self.flask.config['SECRET_KEY'] = self.app.config['SECRET_KEY']
-        self.socketio = SocketIO()
+    def register_routes(self):
+        self.socket_io.on_namespace(ProxyNamespace('/ssh'))
 
     def run(self):
-        host = self.app.config["BIND_HOST"]
-        port = self.app.config["HTTPD_PORT"]
+        host = self.flask_app.config["BIND_HOST"]
+        port = self.flask_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, async_mode="threading")
-        self.socketio.run(self.flask, port=port, host=host)
+        self.socket_io.init_app(self.flask_app, async_mode=self.async_mode)
+        self.socket_io.run(self.flask_app, port=port, host=host, debug=False)
 
     def shutdown(self):
         pass

coco/interactive.py

@@ -82,6 +82,7 @@ class InteractiveServer:
         self.client.send(wr(prompt, before=1, after=0))
         while True:
             data = self.client.recv(10)
+            logger.debug(data)
             if len(data) == 0:
                 self.app.remove_client(self.client)
                 break
@@ -96,6 +97,15 @@ class InteractiveServer:
                 self.client.send(data)
                 continue
 
+            if data.startswith(b'\x03'):
+                # Ctrl-C
+                self.client.send(b'^C\r\nOpt> ')
+                input_data = []
+                continue
+            elif data.startswith(b'\x04'):
+                # Ctrl-D
+                return 'q'
+
             # Todo: Move x1b to char
             if data.startswith(b'\x1b') or data in char.UNSUPPORTED_CHAR:
                 self.client.send(b'')
@@ -119,11 +129,12 @@ class InteractiveServer:
                 input_data.append(data)
 
     def dispatch(self, opt):
+        print(repr(opt))
         if opt is None:
             return self._sentinel
         elif opt.startswith("/"):
             self.search_and_display(opt.lstrip("/"))
-        elif opt in ['p', 'P']:
+        elif opt in ['p', 'P', '']:
             self.display_assets()
         elif opt in ['g', 'G']:
             self.display_asset_groups()
@@ -181,7 +192,7 @@ class InteractiveServer:
         amount_max_length = max(len(str(max([group.assets_amount for group in self.asset_groups]))), 10)
         header = '{1:>%d} {0.name:%d} {0.assets_amount:<%s} ' % (id_max_length, name_max_length, amount_max_length)
         comment_length = self.request.meta["width"] - len(header.format(fake_group, id_max_length))
-        line = header + '{0.comment:%s}' % (comment_length//2)  # comment中可能有中文
+        line = header + '{0.comment:%s}' % (comment_length // 2)  # comment中可能有中文
         header += "{0.comment:%s}" % comment_length
         self.client.send(title(header.format(fake_group, "ID")))
         for index, group in enumerate(self.asset_groups, 1):
@@ -194,12 +205,13 @@ class InteractiveServer:
             self.display_asset_groups()
             return
 
-        self.search_result = self.asset_groups[_id-1].assets_granted
+        self.search_result = self.asset_groups[_id - 1].assets_granted
         self.display_search_result()
 
     def display_search_result(self):
         self.search_result = sort_assets(self.search_result, self.app.config["ASSET_LIST_SORT_BY"])
-        fake_asset = Asset(hostname=_("Hostname"), ip=_("IP"), _system_users_name_list=_("LoginAs"), comment=_("Comment"))
+        fake_asset = Asset(hostname=_("Hostname"), ip=_("IP"), _system_users_name_list=_("LoginAs"),
+                           comment=_("Comment"))
         id_max_length = max(len(str(len(self.search_result))), 3)
         hostname_max_length = max(max([len(asset.hostname) for asset in self.search_result + [fake_asset]]), 15)
         sysuser_max_length = max([len(asset.system_users_name_list) for asset in self.search_result + [fake_asset]])

requirements/requirements.txt

@@ -30,5 +30,5 @@ tornado==4.5.2
 urllib3==1.22
 wcwidth==0.1.7
 werkzeug==0.12.2
-jumpserver-python-sdk==0.0.27
+jumpserver-python-sdk==0.0.28
 jms-es-sdk