Merge pull request #214 from jumpserver/dev

Dev

coco/app.py

@@ -23,7 +23,7 @@ from .session import Session
 from .models import Connection
 
 
-__version__ = '1.4.10'
+__version__ = '1.5.0'
 
 BASE_DIR = os.path.dirname(os.path.dirname(__file__))
 logger = get_logger(__file__)

coco/conf.py

@@ -213,7 +213,7 @@ class Config(dict):
             filename = os.path.join(self.root_path, filename)
         try:
             with open(filename) as f:
-                obj = yaml.load(f)
+                obj = yaml.safe_load(f)
         except IOError as e:
             if silent and e.errno in (errno.ENOENT, errno.EISDIR):
                 return False
@@ -287,22 +287,42 @@ class Config(dict):
             rv[key] = v
         return rv
 
+    def convert_type(self, k, v):
+        default_value = self.defaults.get(k)
+        if default_value is None:
+            return v
+        tp = type(default_value)
+        # 对bool特殊处理
+        if tp is bool and isinstance(v, str):
+            if v in ("true", "True", "1"):
+                return True
+            else:
+                return False
+        if tp in [list, dict] and isinstance(v, str):
+            try:
+                v = json.loads(v)
+                return v
+            except json.JSONDecodeError:
+                return v
+
+        try:
+            v = tp(v)
+        except Exception:
+            pass
+        return v
+
     def __getitem__(self, item):
+        # 先从设置的来
         try:
-            value = super(Config, self).__getitem__(item)
+            value = super().__getitem__(item)
         except KeyError:
             value = None
         if value is not None:
             return value
+        # 其次从环境变量来
         value = os.environ.get(item, None)
         if value is not None:
-            if value.isdigit():
-                value = int(value)
-            elif value.lower() == 'false':
-                value = False
-            elif value.lower() == 'true':
-                value = True
-            return value
+            return self.convert_type(item, value)
         return self.defaults.get(item)
 
     def __getattr__(self, item):
@@ -354,7 +374,8 @@ defaults = {
     'ASSET_LIST_PAGE_SIZE': 'auto',
     'SFTP_ROOT': '/tmp',
     'SFTP_SHOW_HIDDEN_FILE': False,
-    'UPLOAD_FAILED_REPLAY_ON_START': True
+    'UPLOAD_FAILED_REPLAY_ON_START': True,
+    'REUSE_CONNECTION': False,
 }
 
 

coco/connection.py

@@ -24,32 +24,85 @@ AUTO_LOGIN = 'auto'
 
 
 class SSHConnection:
+    connections = {}
+
     @staticmethod
-    def get_system_user_auth(system_user, asset):
+    def make_key(user, asset, system_user):
+        key = "{}_{}_{}".format(user.id, asset.id, system_user.id)
+        return key
+
+    @classmethod
+    def new_connection_from_cache(cls, user, asset, system_user):
+        if not config.REUSE_CONNECTION:
+            return None
+        key = cls.make_key(user, asset, system_user)
+        connection = cls.connections.get(key)
+        if not connection:
+            return None
+        connection.ref += 1
+        return connection
+
+    @classmethod
+    def set_connection_to_cache(cls, conn):
+        if not config.REUSE_CONNECTION:
+            return None
+        key = cls.make_key(conn.user, conn.asset, conn.system_user)
+        cls.connections[key] = conn
+
+    @classmethod
+    def new_connection(cls, user, asset, system_user):
+        connection = cls.new_connection_from_cache(user, asset, system_user)
+
+        if connection:
+            logger.debug("Reuse connection: {}->{}@{}".format(
+                user.username, asset.ip, system_user.username)
+            )
+            return connection
+        connection = cls(user, asset, system_user)
+        cls.set_connection_to_cache(connection)
+        return connection
+
+    @classmethod
+    def remove_ssh_connection(cls, conn):
+        key = "{}_{}_{}".format(conn.user.id, conn.asset.id, conn.system_user.id)
+        cls.connections.pop(key, None)
+
+    def __init__(self, user, asset, system_user):
+        self.user = user
+        self.asset = asset
+        self.system_user = system_user
+        self.client = None
+        self.sock = None
+        self.error = ""
+        self.ref = 1
+
+    def get_system_user_auth(self):
         """
         获取系统用户的认证信息，密码或秘钥
         :return: system user have full info
         """
         password, private_key = \
-            app_service.get_system_user_auth_info(system_user, asset)
-        system_user.password = password
-        system_user.private_key = private_key
+            app_service.get_system_user_auth_info(self.system_user, self.asset)
+        self.system_user.password = password
+        self.system_user.private_key = private_key
 
-    def get_ssh_client(self, asset, system_user):
+    def get_ssh_client(self):
         ssh = paramiko.SSHClient()
         ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
         sock = None
         error = ''
 
-        if not system_user.password and not system_user.private_key:
-            self.get_system_user_auth(system_user, asset)
+        if not self.system_user.password and not self.system_user.private_key:
+            self.get_system_user_auth()
 
-        if asset.domain:
-            sock = self.get_proxy_sock_v2(asset)
+        if self.asset.domain:
+            sock = self.get_proxy_sock_v2(self.asset)
             if not sock:
                 error = 'Connect gateway failed.'
                 logger.error(error)
 
+        asset = self.asset
+        system_user = self.system_user
         try:
             try:
                 ssh.connect(
@@ -86,30 +139,53 @@ class SSHConnection:
                 password_short, key_fingerprint,
             ))
             error += '\r\n' + str(e) if error else str(e)
-            return None, None, error
-        return ssh, sock, None
+            ssh, sock, error = None, None, error
+        self.client = ssh
+        self.sock = ssh
+        self.error = error
 
-    def get_transport(self, asset, system_user):
-        ssh, sock, msg = self.get_ssh_client(asset, system_user)
-        if ssh:
-            return ssh.get_transport(), sock, None
+    def get_transport(self):
+        if not self.client:
+            self.get_ssh_client()
+        if not self.client:
+            return self.client.get_transport()
         else:
-            return None, None, msg
+            return None
 
-    def get_channel(self, asset, system_user, term="xterm", width=80, height=24):
-        ssh, sock, msg = self.get_ssh_client(asset, system_user)
-        if ssh:
-            chan = ssh.invoke_shell(term, width=width, height=height)
-            return chan, sock, None
+    def get_channel(self, term="xterm", width=80, height=24):
+        if not self.client:
+            self.get_ssh_client()
+        if self.client:
+            chan = self.client.invoke_shell(term, width=width, height=height)
+            return chan
         else:
-            return None, sock, msg
+            return None
 
-    def get_sftp(self, asset, system_user):
-        ssh, sock, msg = self.get_ssh_client(asset, system_user)
-        if ssh:
-            return ssh.open_sftp(), sock, None
+    def get_sftp(self):
+        if not self.client:
+            self.get_ssh_client()
+        if self.client:
+            return self.client.open_sftp()
         else:
-            return None, sock, msg
+            return None
+
+    def close(self):
+        if self.ref > 1:
+            self.ref -= 1
+            logger.debug("Connection ref -1: {}->{}@{}".format(
+                self.user.username, self.asset.hostname, self.system_user.username)
+            )
+            return
+        self.__class__.remove_ssh_connection(self)
+        try:
+            self.client.close()
+            if self.sock:
+                self.sock.close()
+        except Exception as e:
+            logger.error("Close connection error: ", e)
+        logger.debug("Close connection: {}->{}@{}".format(
+            self.user.username, self.asset.ip, self.system_user.username)
+        )
 
     @staticmethod
     def get_proxy_sock_v2(asset):

coco/models.py

@@ -398,20 +398,15 @@ class Server(BaseServer):
     """
 
     # Todo: Server name is not very suitable
-    def __init__(self, chan, sock, asset, system_user):
-        self.sock = sock
+    def __init__(self, chan, connection, asset, system_user):
+        self.connection = connection
         self.asset = asset
         self.system_user = system_user
         super(Server, self).__init__(chan=chan)
 
     def close(self):
         super(Server, self).close()
-        for i in range(5):
-            if not self.chan.transport.is_alive():
-                break
-            self.chan.transport.close()
-        if self.sock:
-            self.sock.transport.close()
+        self.connection.close()
 
 
 class WSProxy(object):

coco/proxy.py

@@ -64,9 +64,9 @@ class ProxyServer:
     def proxy(self):
         if not self.check_protocol():
             return
-        self.get_system_user_username_if_need()
-        self.get_system_user_auth_or_manual_set()
-        self.server = self.get_server_conn()
+        self.server = self.get_server_conn_from_cache()
+        if not self.server:
+            self.server = self.get_server_conn()
         if self.server is None:
             return
         if self.client.closed:
@@ -102,16 +102,25 @@ class ProxyServer:
         }
         return app_service.validate_user_asset_permission(**kwargs)
 
+    def get_server_conn_from_cache(self):
+        server = None
+        if self.system_user.protocol == 'ssh':
+            server = self.get_ssh_server_conn(cache=True)
+        return server
+
     def get_server_conn(self):
-        logger.info("Connect to {}:{} ...".format(self.asset.hostname, self.asset.port))
+        # 与获取连接
+        self.get_system_user_username_if_need()
+        self.get_system_user_auth_or_manual_set()
         self.send_connecting_message()
+        logger.info("Connect to {}:{} ...".format(self.asset.hostname, self.asset.port))
         if not self.validate_permission():
             msg = _('No permission')
             self.client.send_unicode(warning(wr(msg, before=2, after=0)))
             server = None
-        elif self.system_user.protocol == self.asset.protocol == 'telnet':
+        elif self.system_user.protocol == 'telnet':
             server = self.get_telnet_server_conn()
-        elif self.system_user.protocol == self.asset.protocol == 'ssh':
+        elif self.system_user.protocol == 'ssh':
             server = self.get_ssh_server_conn()
         else:
             server = None
@@ -129,21 +138,28 @@ class ProxyServer:
             server = TelnetServer(sock, self.asset, self.system_user)
         return server
 
-    def get_ssh_server_conn(self):
+    def get_ssh_server_conn(self, cache=False):
         request = self.client.request
         term = request.meta.get('term', 'xterm')
         width = request.meta.get('width', 80)
         height = request.meta.get('height', 24)
-        ssh = SSHConnection()
-        chan, sock, msg = ssh.get_channel(
-            self.asset, self.system_user, term=term,
-            width=width, height=height
-        )
+
+        if cache:
+            conn = SSHConnection.new_connection_from_cache(
+                self.client.user, self.asset, self.system_user
+            )
+            if not conn:
+                return None
+        else:
+            conn = SSHConnection.new_connection(
+                self.client.user, self.asset, self.system_user
+            )
+        chan = conn.get_channel(term=term, width=width, height=height)
         if not chan:
-            self.client.send_unicode(warning(wr(msg, before=1, after=0)))
+            self.client.send_unicode(warning(wr(conn.error, before=1, after=0)))
             server = None
         else:
-            server = Server(chan, sock, self.asset, self.system_user)
+            server = Server(chan, conn, self.asset, self.system_user)
         return server
 
     def send_connecting_message(self):

coco/sftp.py

@@ -77,7 +77,7 @@ class SFTPServer(paramiko.SFTPServerInterface):
     def get_permed_hosts(self):
         hosts = {}
         assets = app_service.get_user_assets(
-            self.server.connection.user
+            self.server.connection.user, cache_policy='1',
         )
         for asset in assets:
             if asset.protocol != 'ssh':
@@ -87,9 +87,9 @@ class SFTPServer(paramiko.SFTPServerInterface):
             if asset.org_id:
                 key = "{}.{}".format(asset.hostname, asset.org_name)
             value['asset'] = asset
-            value['system_users'] = {su.name: su
+            value['system_users'] = {
+                su.name: su
                 for su in asset.system_users_granted
-                if su.protocol == "ssh" and su.login_mode == 'auto'
             }
             hosts[key] = value
         return hosts
@@ -99,17 +99,9 @@ class SFTPServer(paramiko.SFTPServerInterface):
         super(SFTPServer, self).session_ended()
         for _, v in self._sftp.items():
             sftp = v['client']
-            proxy = v.get('proxy')
-            chan = sftp.get_channel()
-            trans = chan.get_transport()
+            conn = v.get('connection')
             sftp.close()
-
-            active_channels = [c for c in trans._channels.values() if not c.closed]
-            if not active_channels:
-                trans.close()
-                if proxy:
-                    proxy.close()
-                    proxy.transport.close()
+            conn.close()
         self._sftp = {}
 
     def get_host_sftp(self, host, su):
@@ -121,17 +113,18 @@ class SFTPServer(paramiko.SFTPServerInterface):
 
         cache_key = '{}@{}'.format(su, host)
         if cache_key not in self._sftp:
-            ssh = SSHConnection()
-            __sftp, proxy, msg = ssh.get_sftp(asset, system_user)
+            conn = SSHConnection.new_connection(self.server.connection.user,
+                                                asset, system_user)
+            __sftp = conn.get_sftp()
             if __sftp:
                 sftp = {
-                    'client': __sftp, 'proxy': proxy,
+                    'client': __sftp, 'connection': conn,
                     'home': __sftp.normalize('')
                 }
                 self._sftp[cache_key] = sftp
                 return sftp
             else:
-                raise OSError("Can not connect asset sftp server: {}".format(msg))
+                raise OSError("Can not connect asset sftp server: {}".format(conn.error))
         else:
             return self._sftp[cache_key]
 

coco/utils.py

@@ -161,7 +161,7 @@ def is_obj_attr_has(obj, val, attrs=("hostname", "ip", "comment")):
                 hasattr(obj, attr) and isinstance(hasattr(obj, attr), (str, int))]
 
     for v in vals:
-        if str(v).find(val) != -1:
+        if str(v).lower().find(val.lower()) != -1:
             return True
     return False
 
@@ -173,7 +173,7 @@ def is_obj_attr_eq(obj, val, attrs=("id", "hostname", "ip")):
         vals = [getattr(obj, attr) for attr in attrs if hasattr(obj, attr)]
 
     for v in vals:
-        if str(v) == str(val):
+        if str(v).lower() == str(val).lower():
             return True
     return False
 

config_example.yml

@@ -57,3 +57,6 @@ BOOTSTRAP_TOKEN: <PleasgeChangeSameWithJumpserver>
 
 # SFTP是否显示隐藏文件
 # SFTP_SHOW_HIDDEN_FILE: false
+
+# 是否复用和用户后端资产已建立的连接(用户不会复用其他用户的连接)
+# REUSE_CONNECTION: false

requirements/requirements.txt

@@ -16,13 +16,13 @@ Flask==1.0.2
 Flask-SocketIO==2.9.2
 idna==2.6
 itsdangerous==0.24
-Jinja2==2.10
+Jinja2==2.10.1
 jmespath==0.9.3
 jms-storage==0.0.22
 jumpserver-python-sdk==0.0.58
 MarkupSafe==1.0
 oss2==2.4.0
-paramiko==2.4.1
+paramiko==2.4.2
 psutil==5.4.1
 pyasn1==0.4.2
 pycparser==2.19
@@ -33,14 +33,14 @@ python-engineio==2.1.0
 python-gssapi==0.6.4
 python-socketio==1.8.3
 pytz==2018.3
-requests==2.18.4
+requests==2.22.0
 s3transfer==0.1.13
 simplejson==3.13.2
 six==1.11.0
 tornado==4.5.2
-urllib3==1.22
+urllib3==1.25.2
 wcwidth==0.1.7
 eventlet==0.24.1
 Werkzeug==0.14.1
 treelib==1.5.3     
-pyyaml==3.13
+PyYAML==5.1