Merged in dev (pull request #57)

Dev

coco/app.py

@@ -219,13 +219,19 @@ class Coco:
     def add_client(self, client):
         with self.lock:
             self.clients.append(client)
-            logger.info("New client {} join, total {} now".format(client, len(self.clients)))
+            logger.info("New client {} join, total {} now".format(
+                    client, len(self.clients)
+                )
+            )
 
     def remove_client(self, client):
         with self.lock:
             try:
                 self.clients.remove(client)
-                logger.info("Client {} leave, total {} now".format(client, len(self.clients)))
+                logger.info("Client {} leave, total {} now".format(
+                       client, len(self.clients)
+                    )
+                )
                 client.close()
             except:
                 pass
@@ -242,4 +248,5 @@ class Coco:
                 self.sessions.remove(session)
                 self.service.finish_session(session.to_json())
             except ValueError:
-                logger.warning("Remove session: {} fail, maybe already removed".format(session))
\ No newline at end of file
+                msg = "Remove session: {} fail, maybe already removed"
+                logger.warning(msg.format(session))

coco/connection.py

@@ -34,7 +34,7 @@ class SSHConnection:
             self.get_system_user_auth(system_user)
 
         if asset.domain:
-            sock = self.get_proxy_sock(asset)
+            sock = self.get_proxy_sock_v2(asset)
 
         try:
             ssh.connect(
@@ -64,29 +64,54 @@ class SSHConnection:
             return None, str(e)
         except (socket.error, TimeoutError) as e:
             return None, str(e)
-        return ssh, None
+        return ssh, sock, None
 
     def get_transport(self, asset, system_user):
-        ssh, msg = self.get_ssh_client(asset, system_user)
+        ssh, sock, msg = self.get_ssh_client(asset, system_user)
         if ssh:
-            return ssh.get_transport(), None
+            return ssh.get_transport(), sock, None
         else:
-            return None, msg
+            return None, None, msg
 
     def get_channel(self, asset, system_user, term="xterm", width=80, height=24):
-        ssh, msg = self.get_ssh_client(asset, system_user)
+        ssh, sock, msg = self.get_ssh_client(asset, system_user)
         if ssh:
             chan = ssh.invoke_shell(term, width=width, height=height)
-            return chan, None
+            return chan, sock, None
         else:
-            return None, msg
+            return None, sock, msg
 
     def get_sftp(self, asset, system_user):
-        ssh, msg = self.get_ssh_client(asset, system_user)
+        ssh, sock, msg = self.get_ssh_client(asset, system_user)
         if ssh:
-            return ssh.open_sftp(), None
+            return ssh.open_sftp(), sock, None
         else:
-            return None, msg
+            return None, sock, msg
+
+    def get_proxy_sock_v2(self, asset):
+        sock = None
+        domain = app_service.get_domain_detail_with_gateway(
+            asset.domain
+        )
+        if not domain.has_ssh_gateway():
+            return None
+        for i in domain.gateways:
+            gateway = domain.random_ssh_gateway()
+            ssh = paramiko.SSHClient()
+            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+            try:
+                ssh.connect(gateway.ip, username=gateway.username,
+                            password=gateway.password,
+                            pkey=gateway.private_key_obj)
+            except(paramiko.AuthenticationException,
+                   paramiko.BadAuthenticationType,
+                   SSHException):
+                continue
+            sock = ssh.get_transport().open_channel(
+                'direct-tcpip', (asset.ip, asset.port), ('127.0.0.1', 0)
+            )
+            break
+        return sock
 
     def get_proxy_sock(self, asset):
         sock = None

coco/httpd.py

@@ -212,8 +212,7 @@ class ProxyNamespace(BaseNamespace):
 
     def on_disconnect(self):
         logger.debug("On disconnect event trigger")
-        room_id_list = list(self.connections.get(request.sid, {}).keys())
-        for room_id in room_id_list:
+        for room_id in self.connections.get(request.sid, {}):
             try:
                 self.on_logout(room_id)
             except Exception as e:
@@ -250,14 +249,14 @@ class HttpServer:
         self.flask_app.config.update(config)
         self.socket_io = SocketIO()
         self.register_routes()
+        self.register_error_handler()
 
     def register_routes(self):
         self.socket_io.on_namespace(ProxyNamespace('/ssh'))
 
     @staticmethod
     def on_error_default(e):
-        traceback.print_exc()
-        logger.warn(e)
+        logger.exception(e)
 
     def register_error_handler(self):
         self.socket_io.on_error_default(self.on_error_default)

coco/interactive.py

@@ -263,7 +263,7 @@ class InteractiveServer:
         self.display_banner()
         while True:
             try:
-                opt = net_input(self.client, prompt='Opt>', before=1)
+                opt = net_input(self.client, prompt='Opt> ', before=1)
                 rv = self.dispatch(opt)
                 if rv is self._sentinel:
                     break

coco/models.py

@@ -94,8 +94,9 @@ class Server:
     """
 
     # Todo: Server name is not very suitable
-    def __init__(self, chan, asset, system_user):
+    def __init__(self, chan, sock, asset, system_user):
         self.chan = chan
+        self.sock = sock
         self.asset = asset
         self.system_user = system_user
         self.send_bytes = 0
@@ -168,6 +169,8 @@ class Server:
         self.stop_evt.set()
         self.chan.close()
         self.chan.transport.close()
+        if self.sock:
+            self.sock.transport.close()
 
     @staticmethod
     def _have_enter_char(s):
@@ -251,7 +254,6 @@ class WSProxy:
             if len(data) == 0:
                 self.close()
             data = data.decode(errors="ignore")
-            print("Send data: {}".format(data))
             self.ws.emit("data", {'data': data, 'room': self.room_id},
                          room=self.room_id)
             if len(data) == BUF_SIZE:
@@ -264,8 +266,11 @@ class WSProxy:
 
     def close(self):
         self.stop_event.set()
+        try:
             self.child.shutdown(1)
             self.child.close()
+        except (OSError, EOFError):
+            pass
         logger.debug("Proxy {} closed".format(self))
 
 

coco/proxy.py

@@ -90,14 +90,14 @@ class ProxyServer:
         width = request.meta.get('width', 80)
         height = request.meta.get('height', 24)
         ssh = SSHConnection()
-        chan, msg = ssh.get_channel(
+        chan, sock, msg = ssh.get_channel(
             asset, system_user, term=term, width=width, height=height
         )
         if not chan:
             self.client.send(warning(wr(msg, before=1, after=0)))
             server = None
         else:
-            server = Server(chan, asset, system_user)
+            server = Server(chan, sock, asset, system_user)
         self.connecting = False
         self.client.send(b'\r\n')
         return server

coco/sftp.py

@@ -2,6 +2,7 @@ import os
 import tempfile
 import paramiko
 import time
+from .ctx import app_service
 from datetime import datetime
 
 from .connection import SSHConnection
@@ -16,6 +17,17 @@ class SFTPServer(paramiko.SFTPServerInterface):
         self._sftp = {}
         self.hosts = self.get_perm_hosts()
 
+    def session_ended(self):
+        super().session_ended()
+        for _, v in self._sftp.items():
+            sftp = v['sftp']
+            sock = v.get('sock')
+            sftp.close()
+            if sock:
+                sock.close()
+                sock.transport.close()
+        self._sftp = {}
+
     def get_host_sftp(self, host, su):
         asset = self.hosts.get(host)
         system_user = None
@@ -28,18 +40,18 @@ class SFTPServer(paramiko.SFTPServerInterface):
             raise OSError("No asset or system user explicit")
 
         if host not in self._sftp:
-            ssh = SSHConnection(self.server.app)
-            sftp, msg = ssh.get_sftp(asset, system_user)
+            ssh = SSHConnection()
+            sftp, sock, msg = ssh.get_sftp(asset, system_user)
             if sftp:
-                self._sftp[host] = sftp
+                self._sftp[host] = {'sftp': sftp, 'sock': sock}
                 return sftp
             else:
                 raise OSError("Can not connect asset sftp server: {}".format(msg))
         else:
-            return self._sftp[host]
+            return self._sftp[host]['sftp']
 
     def get_perm_hosts(self):
-        assets = self.server.app.service.get_user_assets(
+        assets = app_service.get_user_assets(
             self.server.request.user
         )
         return {asset.hostname: asset for asset in assets}
@@ -89,7 +101,7 @@ class SFTPServer(paramiko.SFTPServerInterface):
             "is_success": is_success,
         }
         for i in range(1, 4):
-            ok = self.server.app.service.create_ftp_log(data)
+            ok = app_service.create_ftp_log(data)
             if ok:
                 break
             else:

coco/utils.py

@@ -396,7 +396,6 @@ def size_of_str_with_zh(s):
     try:
         chinese = find_chinese(s)
     except TypeError:
-        print(type(s))
         raise
     return len(s) + len(chinese)
 
@@ -406,6 +405,8 @@ def item_max_length(_iter, maxi=None, mini=None, key=None):
         _iter = [key(i) for i in _iter]
 
     length = [size_of_str_with_zh(s) for s in _iter]
+    if not length:
+        return 1
     if maxi:
         length.append(maxi)
     length = max(length)

requirements/requirements.txt

@@ -20,7 +20,7 @@ Jinja2==2.10
 jmespath==0.9.3
 jms-es-sdk==0.5.2
 jms-storage==0.0.12
-jumpserver-python-sdk==0.0.41
+jumpserver-python-sdk==0.0.42
 MarkupSafe==1.0
 oss2==2.4.0
 paramiko==2.4.0