[Update] 优化树结构

apps/assets/api/asset_user.py

@@ -148,6 +148,7 @@ class AssetUserTestConnectiveApi(generics.RetrieveAPIView):
     Test asset users connective
     """
     permission_classes = (IsOrgAdminOrAppUser,)
+    serializer_class = serializers.TaskIDSerializer
 
     def get_asset_users(self):
         username = self.request.GET.get('username')

apps/assets/api/node.py

@@ -26,6 +26,7 @@ from ..hands import IsOrgAdmin
 from ..models import Node
 from ..tasks import update_assets_hardware_info_util, test_asset_connectivity_util
 from .. import serializers
+from ..utils import NodeUtil
 
 
 logger = get_logger(__file__)
@@ -79,12 +80,10 @@ class NodeListAsTreeApi(generics.ListAPIView):
     serializer_class = TreeNodeSerializer
 
     def get_queryset(self):
-        queryset = [node.as_tree_node() for node in Node.objects.all()]
-        return queryset
-
-    def filter_queryset(self, queryset):
-        if self.request.query_params.get('refresh', '0') == '1':
-            queryset = self.refresh_nodes(queryset)
+        queryset = Node.objects.all()
+        util = NodeUtil()
+        nodes = util.get_nodes_by_queryset(queryset)
+        queryset = [node.as_tree_node() for node in nodes]
         return queryset
 
     @staticmethod
@@ -114,15 +113,11 @@ class NodeChildrenAsTreeApi(generics.ListAPIView):
 
     def get_queryset(self):
         node_key = self.request.query_params.get('key')
-        if node_key:
-            self.node = Node.objects.get(key=node_key)
-            queryset = self.node.get_children(with_self=False)
-        else:
-            self.is_root = True
-            self.node = Node.root()
-            queryset = list(self.node.get_children(with_self=True))
-            nodes_invalid = Node.objects.exclude(key__startswith=self.node.key)
-            queryset.extend(list(nodes_invalid))
+        util = NodeUtil()
+        if not node_key:
+            node_key = Node.root().key
+        self.node = util.get_node_by_key(node_key)
+        queryset = self.node.get_children(with_self=True)
         queryset = [node.as_tree_node() for node in queryset]
         queryset = sorted(queryset)
         return queryset

apps/assets/models/asset.py

@@ -46,12 +46,6 @@ class AssetQuerySet(models.QuerySet):
         return self.active()
 
 
-class AssetManager(OrgManager):
-    def get_queryset(self):
-        queryset = super().get_queryset().prefetch_related("nodes", "protocols")
-        return queryset
-
-
 class Protocol(models.Model):
     PROTOCOL_SSH = 'ssh'
     PROTOCOL_RDP = 'rdp'
@@ -131,7 +125,7 @@ class Asset(OrgModelMixin):
     date_created = models.DateTimeField(auto_now_add=True, null=True, blank=True, verbose_name=_('Date created'))
     comment = models.TextField(max_length=128, default='', blank=True, verbose_name=_('Comment'))
 
-    objects = AssetManager.from_queryset(AssetQuerySet)()
+    objects = OrgManager.from_queryset(AssetQuerySet)()
 
     def __str__(self):
         return '{0.hostname}({0.ip})'.format(self)
@@ -300,15 +294,20 @@ class Asset(OrgModelMixin):
     @classmethod
     def generate_fake(cls, count=100):
         from random import seed, choice
-        import forgery_py
         from django.db import IntegrityError
         from .node import Node
+        from orgs.utils import get_current_org
+        from orgs.models import Organization
+        org = get_current_org()
+        if not org or not org.is_real():
+            Organization.default().change_to()
+
         nodes = list(Node.objects.all())
         seed()
         for i in range(count):
             ip = [str(i) for i in random.sample(range(255), 4)]
             asset = cls(ip='.'.join(ip),
-                        hostname=forgery_py.internet.user_name(True),
+                        hostname='.'.join(ip),
                         admin_user=choice(AdminUser.objects.all()),
                         created_by='Fake')
             try:

apps/assets/utils.py

 # ~*~ coding: utf-8 ~*~
 #
-from django.utils.translation import ugettext_lazy as _
-from django.core.cache import cache
-from django.utils import timezone
+from django.db.models import Prefetch
 
-from common.utils import get_object_or_none
-from .models import SystemUser, Label
+from common.utils import get_object_or_none, get_logger
+from common.struct import Stack
+from .models import SystemUser, Label, Node, Asset
 
 
-def get_assets_by_id_list(id_list):
-    return Asset.objects.filter(id__in=id_list).filter(is_active=True)
-
-
-def get_system_users_by_id_list(id_list):
-    return SystemUser.objects.filter(id__in=id_list)
+logger = get_logger(__file__)
 
 
 def get_system_user_by_name(name):
@@ -47,4 +41,154 @@ class LabelFilter:
         return queryset
 
 
+class NodeUtil:
+    def __init__(self, with_assets_amount=False, debug=False):
+        self.stack = Stack()
+        self._nodes = {}
+        self.with_assets_amount = with_assets_amount
+        self._debug = debug
+        self.init()
+
+    @staticmethod
+    def sorted_by(node):
+        return [int(i) for i in node.key.split(':')]
+
+    def get_all_nodes(self):
+        all_nodes = Node.objects.all()
+        if self.with_assets_amount:
+            all_nodes = all_nodes.prefetch_related(
+                Prefetch('assets', queryset=Asset.objects.all().only('id'))
+            )
+            for node in all_nodes:
+                node._assets = set(node.assets.all())
+        all_nodes = sorted(all_nodes, key=self.sorted_by)
+
+        guarder = Node(key='', value='Guarder')
+        guarder._assets = []
+        all_nodes.append(guarder)
+        return all_nodes
+
+    def push_to_stack(self, node):
+        # 入栈之前检查
+        # 如果栈是空的，证明是一颗树的根部
+        if self.stack.is_empty():
+            node._full_value = node.value
+            node._parents = []
+        else:
+            # 如果不是根节点,
+            # 该节点的祖先应该是父节点的祖先加上父节点
+            # 该节点的名字是父节点的名字+自己的名字
+            node._parents = [self.stack.top] + self.stack.top._parents
+            node._full_value = ' / '.join(
+                [self.stack.top._full_value, node.value]
+            )
+        node._children = []
+        node._all_children = []
+        self.debug("入栈: {}".format(node.key))
+        self.stack.push(node)
+
+    # 出栈
+    def pop_from_stack(self):
+        _node = self.stack.pop()
+        self.debug("出栈: {} 栈顶: {}".format(_node.key, self.stack.top.key if self.stack.top else None))
+        self._nodes[_node.key] = _node
+        if not self.stack.top:
+            return
+        if self.with_assets_amount:
+            self.stack.top._assets.update(_node._assets)
+            _node._assets_amount = len(_node._assets)
+            delattr(_node, '_assets')
+        self.stack.top._children.append(_node)
+        self.stack.top._all_children.extend([_node] + _node._children)
+
+    def init(self):
+        all_nodes = self.get_all_nodes()
+        for node in all_nodes:
+            self.debug("准备: {} 栈顶: {}".format(node.key, self.stack.top.key if self.stack.top else None))
+            # 入栈之前检查，该节点是不是栈顶节点的子节点
+            # 如果不是，则栈顶出栈
+            while self.stack.top and not self.stack.top.is_children(node):
+                self.pop_from_stack()
+            self.push_to_stack(node)
+        # 出栈最后一个
+        self.debug("剩余: {}".format(', '.join([n.key for n in self.stack])))
+
+    def get_nodes_by_queryset(self, queryset):
+        nodes = []
+        for n in queryset:
+            node = self._nodes.get(n.key)
+            if not node:
+                continue
+            nodes.append(nodes)
+        return [self]
+
+    def get_node_by_key(self, key):
+        return self._nodes.get(key)
+
+    def debug(self, msg):
+        self._debug and logger.debug(msg)
+
+    def set_assets_amount(self):
+        for node in self._nodes.values():
+            node.assets_amount = node._assets_amount
+
+    def set_full_value(self):
+        for node in self._nodes.values():
+            node.full_value = node._full_value
+
+    @property
+    def nodes(self):
+        return list(self._nodes.values())
+
+    # 使用给定节点生成一颗树
+    # 找到他们的祖先节点
+    # 可选找到他们的子孙节点
+    def get_family(self, nodes, with_children=False):
+        tree_nodes = set()
+        for n in nodes:
+            node = self.get_node_by_key(n.key)
+            if not node:
+                continue
+            tree_nodes.update(node._parents)
+            tree_nodes.add(node)
+            if with_children:
+                tree_nodes.update(node._children)
+        for n in tree_nodes:
+            delattr(n, '_children')
+            delattr(n, '_parents')
+        return list(tree_nodes)
+
+
+def test_node_tree():
+    tree = NodeUtil()
+    for node in tree._nodes.values():
+        print("Check {}".format(node.key))
+        children_wanted = node.get_all_children().count()
+        children = len(node._children)
+        if children != children_wanted:
+            print("{} children not equal: {} != {}".format(node.key, children, children_wanted))
+
+        assets_amount_wanted = node.get_all_assets().count()
+        if node._assets_amount != assets_amount_wanted:
+            print("{} assets amount not equal: {} != {}".format(
+                node.key, node._assets_amount, assets_amount_wanted)
+            )
+
+        full_value_wanted = node.full_value
+        if node._full_value != full_value_wanted:
+            print("{} full value not equal: {} != {}".format(
+                node.key, node._full_value, full_value_wanted)
+            )
+
+        parents_wanted = node.get_ancestor().count()
+        parents = len(node._parents)
+        if parents != parents_wanted:
+            print("{} parents count not equal: {} != {}".format(
+                node.key, parents, parents_wanted)
+            )
+
+
+
+
+
 

apps/common/struct.py

+# -*- coding: utf-8 -*-
+#
+
+
+class Stack(list):
+    def is_empty(self):
+        return len(self) == 0
+
+    @property
+    def top(self):
+        if self.is_empty():
+            return None
+        return self[-1]
+
+    @property
+    def bottom(self):
+        if self.is_empty():
+            return None
+        return self[0]
+
+    def size(self):
+        return len(self)
+
+    def push(self, item):
+        self.append(item)

apps/jumpserver/urls.py

@@ -7,7 +7,7 @@ from django.conf.urls.static import static
 from django.conf.urls.i18n import i18n_patterns
 from django.views.i18n import JavaScriptCatalog
 
-from .views import IndexView, LunaView, I18NView
+from .views import IndexView, LunaView, I18NView, HealthCheckView
 from .swagger import get_swagger_view
 
 api_v1 = [
@@ -63,6 +63,7 @@ urlpatterns = [
     path('', IndexView.as_view(), name='index'),
     path('', include(api_v2_patterns)),
     path('', include(api_v1_patterns)),
+    path('api/health/', HealthCheckView.as_view(), name="health"),
     path('luna/', LunaView.as_view(), name='luna-view'),
     path('i18n/<str:lang>/', I18NView.as_view(), name='i18n-switch'),
     path('settings/', include('settings.urls.view_urls', namespace='settings')),

apps/jumpserver/views.py

 import datetime
 import re
+import time
 
 from django.http import HttpResponse, HttpResponseRedirect
 from django.conf import settings
@@ -9,6 +10,7 @@ from django.utils.translation import ugettext_lazy as _
 from django.db.models import Count
 from django.shortcuts import redirect
 from rest_framework.response import Response
+from rest_framework.views import APIView
 from django.views.decorators.csrf import csrf_exempt
 from django.http import HttpResponse
 from django.utils.encoding import iri_to_uri
@@ -222,3 +224,10 @@ def redirect_format_api(request, *args, **kwargs):
         return HttpResponseTemporaryRedirect(_path)
     else:
         return Response({"msg": "Redirect url failed: {}".format(_path)}, status=404)
+
+
+class HealthCheckView(APIView):
+    permission_classes = ()
+
+    def get(self, request):
+        return Response({"status": 1, "time": int(time.time())})

apps/ops/inventory.py

@@ -2,7 +2,6 @@
 #
 
 from .ansible.inventory import BaseInventory
-from assets.utils import get_assets_by_id_list, get_system_user_by_id
 
 from common.utils import get_logger
 

apps/orgs/mixins.py

 # -*- coding: utf-8 -*-
 #
-
+import traceback
 from django.db import models
 from django.utils.translation import ugettext_lazy as _
 from django.shortcuts import redirect, get_object_or_404
@@ -33,8 +33,8 @@ class OrgManager(models.Manager):
     def get_queryset(self):
         queryset = super(OrgManager, self).get_queryset()
         kwargs = {}
-        _current_org = get_current_org()
 
+        _current_org = get_current_org()
         if _current_org is None:
             kwargs['id'] = None
         elif _current_org.is_real():
@@ -42,12 +42,17 @@ class OrgManager(models.Manager):
         elif _current_org.is_default():
             queryset = queryset.filter(org_id="")
 
+        # lines = traceback.format_stack()
+        # print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>")
+        # for line in lines[-10:-5]:
+        #     print(line)
+        # print("<<<<<<<<<<<<<<<<<<<<<<<<<<<<")
+
         queryset = queryset.filter(**kwargs)
         return queryset
 
     def all(self):
-        _current_org = get_current_org()
-        if _current_org is None:
+        if not current_org:
             msg = 'You can `objects.set_current_org(org).all()` then run it'
             return self
         else:

apps/perms/api/user_permission.py

@@ -258,7 +258,9 @@ class UserGrantedNodesWithAssetsAsTreeApi(UserPermissionCacheMixin, ListAPIView)
             util.filter_permissions(
                 system_users=self.system_user_id
             )
+        print("111111111111")
         nodes = util.get_nodes_with_assets()
+        print("22222222222222")
         for node, assets in nodes.items():
             data = parse_node_to_tree_node(node)
             queryset.append(data)

apps/perms/forms/asset_permission.py

@@ -7,7 +7,7 @@ from django.utils.translation import ugettext_lazy as _
 from orgs.mixins import OrgModelForm
 from orgs.utils import current_org
 from perms.models import AssetPermission
-from assets.models import Asset
+from assets.models import Asset, Node
 
 __all__ = [
     'AssetPermissionForm',

apps/settings/api.py

@@ -242,22 +242,3 @@ class CommandStorageDeleteAPI(APIView):
         storage_name = str(request.data.get('name'))
         Setting.delete_storage('TERMINAL_COMMAND_STORAGE', storage_name)
         return Response({"msg": _('Delete succeed')}, status=200)
-
-
-class DjangoSettingsAPI(APIView):
-    def get(self, request):
-        if not settings.DEBUG:
-            return Response("Not in debug mode")
-
-        data = {}
-        for i in [settings, getattr(settings, '_wrapped')]:
-            if not i:
-                continue
-            for k, v in i.__dict__.items():
-                if k and k.isupper():
-                    try:
-                        json.dumps(v)
-                        data[k] = v
-                    except (json.JSONDecodeError, TypeError):
-                        data[k] = str(v)
-        return Response(data)
\ No newline at end of file

apps/settings/urls/api_urls.py

@@ -15,5 +15,4 @@ urlpatterns = [
     path('terminal/replay-storage/delete/', api.ReplayStorageDeleteAPI.as_view(), name='replay-storage-delete'),
     path('terminal/command-storage/create/', api.CommandStorageCreateAPI.as_view(), name='command-storage-create'),
     path('terminal/command-storage/delete/', api.CommandStorageDeleteAPI.as_view(), name='command-storage-delete'),
-    path('django-settings/', api.DjangoSettingsAPI.as_view(), name='django-settings'),
 ]