[Feature] 添加asset display

.gitignore

@@ -5,3 +5,4 @@
 *.log
 logs/*
 conf.py
+host_rsa_key

coco/app.py

@@ -61,17 +61,25 @@ class Coco:
     def make_logger(self):
         create_logger(self)
 
-    def prepare(self):
+    def bootstrap(self):
         self.make_logger()
         self.sshd = SSHServer(self)
         self.httpd = HttpServer(self)
         self.initial_service()
+        self.heartbeat()
 
     def heartbeat(self):
-        pass
+        def _heartbeat():
+            while not self.stop_evt.is_set():
+                self.service.terminal_heartbeat()
+                time.sleep(self.config["HEARTBEAT_INTERVAL"])
+
+        thread = threading.Thread(target=_heartbeat)
+        thread.daemon = True
+        thread.start()
 
     def run_forever(self):
-        self.prepare()
+        self.bootstrap()
         print(time.ctime())
         print('Coco version %s, more see https://www.jumpserver.org' % __version__)
         print('Quit the server with CONTROL-C.')

coco/forward.py

@@ -19,7 +19,7 @@ class ProxyServer:
     def __init__(self, app, client):
         self.app = app
         self.client = client
-        self.request = client.do
+        self.request = client.request
         self.server = None
         self.connecting = True
 

coco/interactive.py

@@ -7,8 +7,9 @@ import threading
 from jms.models import Asset, SystemUser
 
 from . import char
-from .utils import TtyIOParser, wrap_with_line_feed as wr, \
-    wrap_with_primary as primary, wrap_with_warning as warning
+from .utils import wrap_with_line_feed as wr, wrap_with_title as title, \
+    wrap_with_primary as primary, wrap_with_warning as warning, \
+    is_obj_attr_has, is_obj_attr_eq, sort_assets, TtyIOParser
 from .forward import ProxyServer
 from .session import Session
 
@@ -20,22 +21,27 @@ class InteractiveServer:
     def __init__(self, app, client):
         self.app = app
         self.client = client
-        self.request = client.do
+        self.request = client.request
+        self.assets = None
+        self.search_result = None
+        self.asset_groups = None
+        self.get_user_assets_async()
 
     def display_banner(self):
         self.client.send(char.CLEAR_CHAR)
 
-        banner = u"""\n\033[1;32m  %s, 欢迎使用Jumpserver开源跳板机系统  \033[0m\r\n\r
-        1) 输入 \033[32mID\033[0m 直接登录 或 输入\033[32m部分 IP,主机名,备注\033[0m 进行搜索登录(如果唯一).\r
-        2) 输入 \033[32m/\033[0m + \033[32mIP, 主机名 or 备注 \033[0m搜索. 如: /ip\r
-        3) 输入 \033[32mP/p\033[0m 显示您有权限的主机.\r
-        4) 输入 \033[32mG/g\033[0m 显示您有权限的主机组.\r
-        5) 输入 \033[32mG/g\033[0m\033[0m + \033[32m组ID\033[0m 显示该组下主机. 如: g1\r
-        6) 输入 \033[32mE/e\033[0m 批量执行命令.(未完成)\r
-        7) 输入 \033[32mU/u\033[0m 批量上传文件.(未完成)\r
-        8) 输入 \033[32mD/d\033[0m 批量下载文件.(未完成)\r
-        9) 输入 \033[32mH/h\033[0m 帮助.\r
-        0) 输入 \033[32mQ/q\033[0m 退出.\r\n""" % self.request.user
+        banner = u"""\n {title} {user}, 欢迎使用Jumpserver开源跳板机系统  {end}\r\n\r
+        1) 输入 {green}ID{end} 直接登录 或 输入{green}部分 IP,主机名,备注{end} 进行搜索登录(如果唯一).\r
+        2) 输入 {green}/{end} + {green}IP, 主机名 or 备注 {end}搜索. 如: /ip\r
+        3) 输入 {green}P/p{end} 显示您有权限的主机.\r
+        4) 输入 {green}G/g{end} 显示您有权限的主机组.\r
+        5) 输入 {green}G/g{end}{end} + {green}组ID{end} 显示该组下主机. 如: g1\r
+        6) 输入 {green}E/e{end} 批量执行命令.(未完成)\r
+        7) 输入 {green}U/u{end} 批量上传文件.(未完成)\r
+        8) 输入 {green}D/d{end} 批量下载文件.(未完成)\r
+        9) 输入 {green}H/h{end} 帮助.\r
+        0) 输入 {green}Q/q{end} 退出.\r\n""".format(
+            title="\033[1;32m", green="\033[32m", end="\033[0m", user=self.client.user)
         self.client.send(banner.encode('utf-8'))
 
     def get_choice(self, prompt=b'Opt> '):
@@ -97,33 +103,81 @@ class InteractiveServer:
         else:
             self.search_and_proxy(opt)
 
-    def search_assets(self, opt, from_result=False):
-        pass
+    def search_assets(self, q, from_result=False):
+        if self.assets is None:
+            self.get_user_assets()
+        if from_result:
+            source = self.search_result
+        else:
+            source = self.assets
+
+        result = []
+        # 用户输入的是数字，可能想使用id唯一键搜索
+        if q.isdigit():
+            if len(source) > int(q):
+                result = [source[q-1]]
+
+        # 全匹配则直接返回
+        if len(result) == 0:
+            _result = [asset for asset in source if is_obj_attr_eq(asset, q)]
+            if len(_result) == 1:
+                result = _result
+
+        # 最后模糊匹配
+        if len(result) == 0:
+            result = [asset for asset in source if is_obj_attr_has(asset, q)]
+
+        self.search_result = sort_assets(result, self.app.config["ASSET_LIST_SORT_BY"])
 
     def display_assets(self):
         """
         Display user all assets
         :return:
         """
-        pass
-
-    def display_asset_groups(self):
-        pass
+        self.search_and_display('')
 
     def display_group_assets(self):
         pass
 
     def display_search_result(self):
-        pass
-
-    def search_and_display(self, opt):
-        self.search_assets(opt=opt)
+        if len(self.search_result) == 0:
+            self.client.send(warning("Nothing match"))
+            return
+
+        fake_asset = Asset(hostname="Hostname", ip="IP", system_users_join="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_join) for asset in self.search_result + [fake_asset]])
+        header = '{1:%d} {0.hostname:%d} {0.ip:15} {0.system_users_join:%d} ' % \
+                 (id_max_length, hostname_max_length, sysuser_max_length)
+        comment_length = self.request.meta["width"] - len(header.format(fake_asset, id_max_length))
+        line = header + '{0.comment:.%d}' % (comment_length / 2)  # comment中可能有中文
+        header += '{0.comment:%s}' % comment_length
+        self.client.send(wr(title(header.format(fake_asset, "ID").encode("utf-8"))))
+        self.client.send(b'\r\n')
+        for index, asset in enumerate(self.search_result, 1):
+            self.client.send(line.format(asset, index).encode('utf-8'))
+            self.client.send(b'\r\n')
+
+    def search_and_display(self, q):
+        self.search_assets(q)
         self.display_search_result()
 
     def get_user_asset_groups(self):
         pass
 
+    def get_user_asset_groups_async(self):
+        thread = threading.Thread(target=self.get_user_asset_groups)
+        thread.start()
+
     def get_user_assets(self):
+        self.assets = self.app.service.get_user_assets(self.client.user)
+
+    def get_user_assets_async(self):
+        thread = threading.Thread(target=self.get_user_assets)
+        thread.start()
+
+    def display_asset_groups(self):
         pass
 
     def choose_system_user(self, system_users):

coco/interface.py

@@ -34,8 +34,13 @@ class SSHInterface(paramiko.ServerInterface):
         return False
 
     def get_allowed_auths(self, username):
-        # Todo: Check with server settings or self config
-        return ",".join(["password", "publickkey"])
+        supported = []
+        if self.app.config["SSH_PASSWORD_AUTH"]:
+            supported.append("password")
+        if self.app.config["SSH_PUBLIC_KEY_AUTH"]:
+            supported.append("publickey")
+
+        return ",".join(supported)
 
     def check_auth_none(self, username):
         return paramiko.AUTH_FAILED
@@ -44,12 +49,20 @@ class SSHInterface(paramiko.ServerInterface):
         return self.validate_auth(username, password=password)
 
     def check_auth_publickey(self, username, key):
+        key = key.get_base64()
         return self.validate_auth(username, key=key)
 
     def validate_auth(self, username, password="", key=""):
-        # Todo: Implement it
-        self.request.user = "guang"
-        return paramiko.AUTH_SUCCESSFUL
+        user, _ = self.app.service.authenticate(
+            username, password=password, pubkey=key,
+            remote_addr=self.request.remote_ip, login_type="ST"
+        )
+
+        if user:
+            self.request.user = user
+            return paramiko.AUTH_SUCCESSFUL
+        else:
+            return paramiko.AUTH_FAILED
 
     def check_channel_direct_tcpip_request(self, chanid, origin, destination):
         logger.debug("Check channel direct tcpip request: %d %s %s" %

coco/sshd.py

@@ -86,7 +86,7 @@ class SSHServer:
         self.dispatch(client)
 
     def dispatch(self, client):
-        request_type = client.do.type
+        request_type = client.request.type
         if request_type == 'pty':
             InteractiveServer(self.app, client).activate()
         elif request_type == 'exec':

coco/utils.py

@@ -18,7 +18,6 @@ import pytz
 from email.utils import formatdate
 
 
-
 try:
     from Queue import Queue, Empty
 except ImportError:
@@ -190,6 +189,30 @@ class TtyIOParser(object):
         return command
 
 
+def is_obj_attr_has(obj, val, attrs=("hostname", "ip", "comment")):
+    if not attrs:
+        vals = [val for val in obj.__dict__.values() if isinstance(val, (str, int))]
+    else:
+        vals = [getattr(obj, attr) for attr in attrs if hasattr(obj, attr) and isinstance(hasattr(obj, attr), (str, int))]
+
+    for v in vals:
+        if str(v).find(val) != -1:
+            return True
+    return False
+
+
+def is_obj_attr_eq(obj, val, attrs=("id", "hostname", "ip")):
+    if not attrs:
+        vals = [val for val in obj.__dict__.values() if isinstance(val, (str, int))]
+    else:
+        vals = [getattr(obj, attr) for attr in attrs if hasattr(obj, attr)]
+
+    for v in vals:
+        if str(v) == str(val):
+            return True
+    return False
+
+
 def wrap_with_line_feed(s, before=0, after=1):
     if isinstance(s, bytes):
         return b'\r\n' * before + s + b'\r\n' * after
@@ -230,10 +253,15 @@ def wrap_with_color(text, color='white', background=None,
         wrap_with.append(background_map.get(background, ''))
     wrap_with.append(color_map.get(color, ''))
 
+    is_bytes = True if isinstance(text, bytes) else False
+
+    if is_bytes:
+        text = text.decode("utf-8")
     data = '\033[' + ';'.join(wrap_with) + 'm' + text + '\033[0m'
-    if isinstance(text, bytes):
+    if is_bytes:
         return data.encode('utf-8')
-    return data
+    else:
+        return data
 
 
 def wrap_with_warning(text, bolder=False):
@@ -303,10 +331,6 @@ class PKey(object):
                 return None
 
 
-def from_string(cls, key_string):
-    return cls(key_string=key_string).pkey
-
-
 def timestamp_to_datetime_str(ts):
     datetime_format = '%Y-%m-%dT%H:%M:%S.%fZ'
     dt = datetime.datetime.fromtimestamp(ts, tz=pytz.timezone('UTC'))