[Update] 优化分页展示逻辑，展示最后一页后退出分页交互 (#142)

* [Update] 优化分页展示逻辑，展示最后一页后退出分页交互

* [Update] 修复展示资产时直接代理后，会再次显示所有资产

* [Update] 当显示所有资产时, 一次性显示

* [Update] 修复bug，不分页时P展示资产也会卡住

* [Update] 优化代码

coco/interactive.py

@@ -24,7 +24,6 @@ PAGE_DOWN = 'down'
 PAGE_UP = 'up'
 BACK = 'back'
 PROXY = 'proxy'
-PAGE_SIZE_ALL = 100000000
 
 
 class InteractiveServer:
@@ -33,50 +32,51 @@ class InteractiveServer:
     def __init__(self, client):
         self.client = client
         self.closed = False
-        self._search_result = None
+        self._results = None
         self.nodes = None
         self.offset = 0
         self.limit = 100
-        self.assets_list = []
+        self.assets = []
         self.finish = False
         self.page = 1
-        self.total_assets = 0
-        self.total_count = 0  # 分页展示中用来存放数目总条数
-        self.nodes_tree = None  # 授权节点树
+        self.total_assets = 0   # 用户被授权的所有资产
+        self.total_count = 0  # 分页展示中的资产总数量
+        self.node_tree = None  # 授权节点树
         self.get_user_assets_paging_async()
         self.get_user_nodes_async()
 
     @property
     def page_size(self):
         _page_size = config['ASSET_LIST_PAGE_SIZE']
-
         if _page_size.isdigit():
             return int(_page_size)
-        elif _page_size == 'auto':
-            return self.client.request.meta['height'] - 8
+        elif _page_size == 'all':
+            return self.total_count
         else:
-            return PAGE_SIZE_ALL
+            return self.client.request.meta['height'] - 8
+
+    @property
+    def total_pages(self):
+        return math.ceil(self.total_count/self.page_size)
+
+    @property
+    def need_paging(self):
+        return config['ASSET_LIST_PAGE_SIZE'] != 'all'
 
     @property
-    def search_result(self):
-        if self._search_result:
-            return self._search_result
+    def results(self):
+        if self._results:
+            return self._results
         else:
             return []
 
-    @search_result.setter
-    def search_result(self, value):
-        self._search_result = value
+    @results.setter
+    def results(self, value):
+        self._results = value
 
-    def display_logo(self):
-        logo_path = os.path.join(config['ROOT_PATH'], "logo.txt")
-        if not os.path.isfile(logo_path):
-            return
-        with open(logo_path, 'rb') as f:
-            for i in f:
-                if i.decode('utf-8').startswith('#'):
-                    continue
-                self.client.send(i.decode('utf-8').replace('\n', '\r\n'))
+    #
+    # Display banner
+    #
 
     def display_banner(self):
         self.client.send(char.CLEAR_CHAR)
@@ -102,15 +102,25 @@ class InteractiveServer:
                 T='\t', R='\r\n\r'
             ))
 
+    def display_logo(self):
+        logo_path = os.path.join(config['ROOT_PATH'], "logo.txt")
+        if not os.path.isfile(logo_path):
+            return
+        with open(logo_path, 'rb') as f:
+            for i in f:
+                if i.decode('utf-8').startswith('#'):
+                    continue
+                self.client.send(i.decode('utf-8').replace('\n', '\r\n'))
+
     def dispatch(self, opt):
         if opt is None:
             return self._sentinel
         elif opt.startswith("/"):
-            self.search_and_display(opt.lstrip("/"))
+            self.search_and_display_assets(opt.lstrip("/"))
         elif opt in ['p', 'P', '']:
             self.display_assets()
         elif opt in ['g', 'G']:
-            self.display_nodes_tree()
+            self.display_nodes_as_tree()
         elif opt.startswith("g") and opt.lstrip("g").isdigit():
             self.display_node_assets(int(opt.lstrip("g")))
         elif opt in ['q', 'Q', 'exit', 'quit']:
@@ -121,95 +131,165 @@ class InteractiveServer:
         elif opt in ['h', 'H']:
             self.display_banner()
         else:
-            self.search_and_proxy(opt)
+            self.search_and_proxy_assets(opt)
+
+    #
+    # Search assets
+    #
+
+    def search_and_display_assets(self, q):
+        assets = self.search_assets(q)
+        self.display_assets(assets)
+
+    def search_and_proxy_assets(self, opt):
+        assets = self.search_assets(opt)
+        if assets and len(assets) == 1:
+            asset = assets[0]
+            if asset.protocol == "rdp" \
+                    or asset.platform.lower().startswith("windows"):
+                self.client.send(warning(
+                    _("Terminal does not support login rdp, "
+                      "please use web terminal to access"))
+                )
+                return
+            self.proxy(asset)
+        else:
+            self.display_assets(assets)
 
     def search_assets(self, q):
-        if not self.finish:
+        if self.finish:
+            assets = self.search_assets_from_local(q)
+        else:
+            assets = self.search_assets_from_core(q)
+        return assets
+
+    def search_assets_from_core(self, q):
         assets = app_service.get_search_user_granted_assets(self.client.user, q)
         assets = self.filter_system_users(assets)
         return assets
-        assets = self.assets_list
+
+    def search_assets_from_local(self, q):
         result = []
 
         # 所有的
         if q in ('', None):
-            result = assets
+            result = self.assets
 
         # 全匹配到则直接返回全匹配的
         if len(result) == 0:
-            _result = [asset for asset in assets
-                       if is_obj_attr_eq(asset, q)]
+            _result = [asset for asset in self.assets if is_obj_attr_eq(asset, q)]
             if len(_result) == 1:
                 result = _result
 
         # 最后模糊匹配
         if len(result) == 0:
-            result = [asset for asset in assets
-                      if is_obj_attr_has(asset, q)]
+            result = [asset for asset in self.assets if is_obj_attr_has(asset, q)]
 
         return result
 
-    def display_assets(self):
-        """
-        Display user all assets
-        :return:
-        """
-        self.display_result_paging(self.assets_list)
+    #
+    # Display assets
+    #
 
-    def display_nodes(self):
-        if self.nodes is None:
-            self.get_user_nodes()
+    def display_assets(self, assets=None):
+        if assets is None:
+            assets = self.assets
+        self.display_assets_paging(assets)
 
-        if len(self.nodes) == 0:
-            self.client.send(warning(_("No")))
+    def display_assets_paging(self, assets):
+
+        if len(assets) == 0:
+            self.client.send(wr(_("No Assets"), before=0))
             return
 
-        id_length = max(len(str(len(self.nodes))), 5)
-        name_length = item_max_length(self.nodes, 15, key=lambda x: x.name)
-        amount_length = item_max_length(self.nodes, 10, key=lambda x: x.assets_amount)
-        size_list = [id_length, name_length, amount_length]
-        fake_data = ['ID', _("Name"), _("Assets")]
+        self.total_count = self.total_assets if assets is self.assets else len(assets)
 
-        self.client.send(wr(title(format_with_zh(size_list, *fake_data))))
-        for index, node in enumerate(self.nodes, 1):
-            data = [index, node.name, node.assets_amount]
-            self.client.send(wr(format_with_zh(size_list, *data)))
-        self.client.send(wr(_("Total: {}").format(len(self.nodes)), before=1))
+        action = None
+        gen = self._page_generator(assets)
+        while True:
+            try:
+                page, _assets = gen.send(action)
+            except StopIteration as e:
+                if None not in e.value:
+                    page, _assets = e.value
+                    self.display_a_page_assets(page, _assets)
+                break
+            else:
+                self.display_a_page_assets(page, _assets)
+                self.display_page_bottom_prompt()
+                action = self.get_user_action()
 
-    def display_nodes_tree(self):
-        if self.nodes is None:
-            self.get_user_nodes()
+    def _page_generator(self, assets):
+        start, page = 0, 1
+        while True:
+            _assets = assets[start:start+self.page_size]
+            # 等待加载
+            if (assets is self.assets) and (not self.finish) and (not self.need_paging):
+                time.sleep(1)
+                continue
+            # 最后一页
+            elif _assets and (page == self.total_pages) and (
+                    assets is not self.assets
+                    or (assets is self.assets and self.finish)):
+                return page, _assets
+            # 执行动作
+            else:
+                action = yield page, _assets
 
-        if not self.nodes:
-            self.client.send(wr(_('No Nodes'), before=0))
-            return
+                # 退出
+                if action == BACK:
+                    return None, None
+                # 不分页, 不对页码和下标做更改
+                elif not self.need_paging:
+                    continue
+                # 上一页
+                elif action == PAGE_UP:
+                    if page <= 1:
+                        page = 1
+                        start = 0
+                    else:
+                        page -= 1
+                        start -= self.page_size
+                # 下一页
+                else:
+                    page += 1
+                    start += len(_assets)
 
-        self.nodes_tree.show(key=lambda node: node.identifier)
-        self.client.send(wr(title(_("Node: [ ID.Name(Asset amount) ]")), before=0))
-        self.client.send(wr(self.nodes_tree._reader.replace('\n', '\r\n'), before=0))
-        prompt = _("Tips: Enter g+NodeID to display the host under the node, such as g1")
-        self.client.send(wr(title(prompt), before=1))
+    def display_a_page_assets(self, page, assets):
+        self.client.send(char.CLEAR_CHAR)
+        self.page = page
+        self.results = assets
+        self.display_results()
 
-    def display_node_assets(self, _id):
-        if self.nodes is None:
-            self.get_user_nodes()
-        if _id > len(self.nodes) or _id <= 0:
-            msg = wr(warning(_("There is no matched node, please re-enter")))
-            self.client.send(msg)
-            self.display_nodes_tree()
-            return
+    def display_page_bottom_prompt(self):
+        self.client.send(wr(_('Tips: Enter the asset ID and log directly into the asset.'), before=1))
+        prompt_page_up = _("Page up: P/p")
+        prompt_page_down = _("Page down: Enter|N/n")
+        prompt_back = _("BACK: b/q")
+        prompts = [prompt_page_up, prompt_page_down, prompt_back]
+        prompt = '\t'.join(prompts)
+        self.client.send(wr(prompt, before=1))
 
-        assets = self.nodes[_id-1].assets_granted
-        self.display_result_paging(assets)
+    def get_user_action(self):
+        opt = net_input(self.client, prompt=':')
+        if opt in ('p', 'P'):
+            return PAGE_UP
+        elif opt in ('b', 'q'):
+            return BACK
+        elif opt.isdigit() and self.results and 0 < int(opt) <= len(self.results):
+            self.proxy(self.results[int(opt)-1])
+            return BACK
+        else:
+            return PAGE_DOWN
 
-    def display_search_result(self):
+    def display_results(self):
         sort_by = config["ASSET_LIST_SORT_BY"]
-        self.search_result = sort_assets(self.search_result, sort_by)
+        self.results = sort_assets(self.results, sort_by)
         fake_data = [_("ID"), _("Hostname"), _("IP"), _("LoginAs")]
-        id_length = max(len(str(len(self.search_result))), 4)
-        hostname_length = item_max_length(self.search_result, 15,
+        id_length = max(len(str(len(self.results))), 4)
+        hostname_length = item_max_length(self.results, 15,
                                           key=lambda x: x.hostname)
-        sysuser_length = item_max_length(self.search_result,
+        sysuser_length = item_max_length(self.results,
                                          key=lambda x: x.system_users_name_list)
         size_list = [id_length, hostname_length, 16, sysuser_length]
         header_without_comment = format_with_zh(size_list, *fake_data)
@@ -221,44 +301,100 @@ class InteractiveServer:
         size_list.append(comment_length)
         fake_data.append(_("Comment"))
         self.client.send(wr(title(format_with_zh(size_list, *fake_data))))
-        for index, asset in enumerate(self.search_result, 1):
+        for index, asset in enumerate(self.results, 1):
             data = [
                 index, asset.hostname, asset.ip,
                 asset.system_users_name_list, asset.comment
             ]
             self.client.send(wr(format_with_zh(size_list, *data)))
 
-        total_page = math.ceil(self.total_count/self.page_size)
         self.client.send(wr(title(_("Page: {}, Count: {}, Total Page: {}, Total Count: {}").format(
-            self.page, len(self.search_result), total_page, self.total_count)), before=1)
+            self.page, len(self.results), self.total_pages, self.total_count)), before=1)
         )
 
-    def search_and_display(self, q):
-        assets = self.search_assets(q)
-        self.display_result_paging(assets)
+    #
+    # Get assets
+    #
+
+    def get_user_assets_paging_async(self):
+        thread = threading.Thread(target=self.get_user_assets_paging)
+        thread.start()
+
+    def get_user_assets_paging(self):
+        while not self.closed:
+            assets, total = app_service.get_user_assets_paging(
+                self.client.user, offset=self.offset, limit=self.limit
+            )
+            if not assets:
+                logger.info('Get user assets paging async finished.')
+                self.finish = True
+                break
+
+            logger.info('Get user assets paging async: {}'.format(len(assets)))
+            assets = self.filter_system_users(assets)
+            self.total_assets = total
+            self.assets.extend(assets)
+            self.offset += self.limit
+
+    #
+    # Nodes
+    #
+
+    def get_user_nodes_async(self):
+        thread = threading.Thread(target=self.get_user_nodes)
+        thread.start()
 
     def get_user_nodes(self):
-        self.nodes = app_service.get_user_asset_groups(self.client.user)
-        self.filter_nodes_assets_system_user()
-        self.sort_nodes()
-        self.construct_nodes_tree()
+        nodes = app_service.get_user_asset_groups(self.client.user)
+        nodes = sorted(nodes, key=lambda node: node.key)
+        self.nodes = self.filter_system_users_of_assets_under_nodes(nodes)
+        self._construct_node_tree()
 
-    def sort_nodes(self):
-        self.nodes = sorted(self.nodes, key=lambda node: node.key)
+    def filter_system_users_of_assets_under_nodes(self, nodes):
+        for node in nodes:
+            node.assets_granted = self.filter_system_users(node.assets_granted)
+        return nodes
 
-    def construct_nodes_tree(self):
-        self.nodes_tree = Tree()
+    def _construct_node_tree(self):
+        self.node_tree = Tree()
         root = 'ROOT_ALL_ORG_NODE'
-        self.nodes_tree.create_node(tag='', identifier=root, parent=None)
+        self.node_tree.create_node(tag='', identifier=root, parent=None)
         for index, node in enumerate(self.nodes):
             tag = "{}.{}({})".format(index+1, node.name, node.assets_amount)
             key = node.key
             parent_key = key[:node.key.rfind(':')] or root
-            self.nodes_tree.create_node(tag=tag, identifier=key, data=node, parent=parent_key)
+            self.node_tree.create_node(tag=tag, identifier=key, data=node, parent=parent_key)
 
-    def get_user_nodes_async(self):
-        thread = threading.Thread(target=self.get_user_nodes)
-        thread.start()
+    def display_nodes_as_tree(self):
+        if self.nodes is None:
+            self.get_user_nodes()
+
+        if not self.nodes:
+            self.client.send(wr(_('No Nodes'), before=0))
+            return
+
+        self.node_tree.show(key=lambda node: node.identifier)
+        self.client.send(wr(title(_("Node: [ ID.Name(Asset amount) ]")), before=0))
+        self.client.send(wr(self.node_tree._reader.replace('\n', '\r\n'), before=0))
+        prompt = _("Tips: Enter g+NodeID to display the host under the node, such as g1")
+        self.client.send(wr(title(prompt), before=1))
+
+    def display_node_assets(self, _id):
+        if self.nodes is None:
+            self.get_user_nodes()
+
+        if _id > len(self.nodes) or _id <= 0:
+            msg = wr(warning(_("There is no matched node, please re-enter")))
+            self.client.send(msg)
+            self.display_nodes_as_tree()
+            return
+
+        assets = self.nodes[_id-1].assets_granted
+        self.display_assets(assets)
+
+    #
+    # System users
+    #
 
     @staticmethod
     def filter_system_users(assets):
@@ -271,31 +407,6 @@ class InteractiveServer:
             asset.system_users_granted = system_users_cleaned
         return assets
 
-    def filter_nodes_assets_system_user(self):
-        for node in self.nodes:
-            node.assets_granted = self.filter_system_users(node.assets_granted)
-
-    def get_user_assets_paging(self):
-        while not self.closed:
-            assets, total = app_service.get_user_assets_paging(
-                self.client.user, offset=self.offset, limit=self.limit
-            )
-            logger.info('Get user assets paging async: {}'.format(len(assets)))
-            if not assets:
-                logger.info('Get user assets paging async finished.')
-                self.finish = True
-                return
-            if not self.total_assets:
-                self.total_assets = total
-                self.total_count = total
-            assets = self.filter_system_users(assets)
-            self.assets_list.extend(assets)
-            self.offset += self.limit
-
-    def get_user_assets_paging_async(self):
-        thread = threading.Thread(target=self.get_user_assets_paging)
-        thread.start()
-
     def choose_system_user(self, system_users):
         if len(system_users) == 1:
             return system_users[0]
@@ -319,127 +430,9 @@ class InteractiveServer:
         for index, system_user in enumerate(system_users):
             self.client.send(wr("{} {}".format(index, system_user.name)))
 
-    def search_and_proxy(self, opt):
-        assets = self.search_assets(opt)
-        if assets and len(assets) == 1:
-            asset = assets[0]
-            self.search_result = None
-            if asset.protocol == "rdp" or asset.platform.lower().startswith("windows"):
-                self.client.send(warning(
-                    _("Terminal does not support login rdp, "
-                      "please use web terminal to access"))
-                )
-                return
-            self.proxy(asset)
-        else:
-            self.display_result_paging(assets)
-
-    def display_result_paging(self, result_list):
-
-        if result_list is self.assets_list:
-            self.total_count = self.total_assets
-        else:
-            if len(result_list) == 0:
-                return
-            self.total_count = len(result_list)
-
-        action = PAGE_DOWN
-        gen_result = self.get_result_page_down_or_up(result_list)
-        while True:
-            try:
-                page, result = gen_result.send(action)
-            except TypeError:
-                try:
-                    page, result = next(gen_result)
-                except StopIteration:
-                    logger.info('No Assets')
-                    # self.display_banner()
-                    self.client.send(wr(_("No Assets"), before=0))
-                    return None
-            except StopIteration:
-                logger.info('Back display result paging.')
-                # self.display_banner()
-                return None
-            self.display_result_of_page(page, result)
-            action = self.get_user_action()
-
-    def get_result_page_down_or_up(self, result_list):
-        left = 0
-        page = 1
-        page_up_size = 0  # 记录上一页大小
-        while True:
-            right = left + self.page_size
-            result = result_list[left:right]
-
-            if not result and (result_list is self.assets_list) and self.finish and self.total_assets == 0:
-                # 无授权资产
-                return None, None
-
-            elif not result and (result_list is self.assets_list) and self.finish:
-                # 上一页是最后一页
-                left -= page_up_size
-                page -= 1
-                continue
-            elif not result and (result_list is self.assets_list) and not self.finish:
-                # 还有下一页(暂时没有加载完)，需要等待
-                time.sleep(1)
-                continue
-            elif not result and (result_list is not self.assets_list):
-                # 上一页是最后一页
-                left -= page_up_size
-                page -= 1
-                continue
-            else:
-                # 其他4中情况，返回assets
-                action = yield (page, result)
-
-                if action == BACK:
-                    return None, None
-                elif action == PAGE_UP:
-                    if page <= 1:
-                        # 已经是第一页了
-                        page = 1
-                        left = 0
-                    else:
-                        page -= 1
-                        left -= self.page_size
-                else:
-                    if self.page_size == PAGE_SIZE_ALL:
-                        # 如果全部显示左下标不做修改
-                        continue
-                    # PAGE_DOWN
-                    page += 1
-                    left += len(result)
-                    page_up_size = len(result)
-
-    def display_result_of_page(self, page, result):
-        self.client.send(char.CLEAR_CHAR)
-        self.page = page
-        self.search_result = result
-        self.display_search_result()
-        self.display_prompt_of_page()
-
-    def display_prompt_of_page(self):
-        self.client.send(wr(_('Tips: Enter the asset ID and log directly into the asset.'), before=1))
-        prompt_page_up = _("Page up: P/p")
-        prompt_page_down = _("Page down: Enter|N/n")
-        prompt_back = _("BACK: b/q")
-        prompts = [prompt_page_up, prompt_page_down, prompt_back]
-        prompt = '\t'.join(prompts)
-        self.client.send(wr(prompt, before=1))
-
-    def get_user_action(self):
-        opt = net_input(self.client, prompt=':')
-        if opt in ('p', 'P'):
-            return PAGE_UP
-        elif opt in ('b', 'q'):
-            return BACK
-        elif opt.isdigit() and self.search_result and 0 < int(opt) <= len(self.search_result):
-            self.proxy(self.search_result[int(opt)-1])
-            return BACK
-        else:
-            # PAGE_DOWN
-            return PAGE_DOWN
+    #
+    # Proxy
+    #
 
     def proxy(self, asset):
         system_user = self.choose_system_user(asset.system_users_granted)
@@ -449,6 +442,10 @@ class InteractiveServer:
         forwarder = ProxyServer(self.client, asset, system_user)
         forwarder.proxy()
 
+    #
+    # Entrance
+    #
+
     def interact(self):
         self.display_banner()
         while not self.closed:
@@ -462,15 +459,16 @@ class InteractiveServer:
                 break
         self.close()
 
-    def interact_async(self):
-        thread = threading.Thread(target=self.interact)
-        thread.daemon = True
-        thread.start()
-
     def close(self):
         logger.debug("Interactive server server close: {}".format(self))
         self.closed = True
         # current_app.remove_client(self.client)
 
+    def interact_async(self):
+        # 目前没用
+        thread = threading.Thread(target=self.interact)
+        thread.daemon = True
+        thread.start()
+
     # def __del__(self):
     #     print("GC: Interactive class been gc")