Stash

apps/perms/utils/asset_permission.py

@@ -17,8 +17,8 @@ from orgs.utils import set_to_root_org
 from common.utils import get_logger
 from common.tree import TreeNode
 from .. import const
-from ..models import AssetPermission, Action
-from ..hands import Node
+from ..models import AssetPermission, Action, ActionFlag
+from ..hands import Node, Asset
 from assets.utils import NodeUtil
 
 logger = get_logger(__file__)
@@ -31,17 +31,57 @@ __all__ = [
 ]
 
 
+class TreeNodeCounter(NodeUtil):
+    def __init__(self, nodes):
+        self.__nodes = nodes
+        super().__init__(with_assets_amount=True)
+
+    def get_queryset(self):
+        return self.__nodes
+
+
+def timeit(func):
+    def wrapper(*args, **kwargs):
+        logger.debug("Start call: {}".format(func.__name__))
+        now = time.time()
+        result = func(*args, **kwargs)
+        using = time.time() - now
+        logger.debug("Call {} end, using: {:.2}s".format(func.__name__, using))
+        return result
+    return wrapper
+
+
 class GenerateTree:
     def __init__(self):
         """
-        nodes: {"node_instance": {
-            "asset_instance": set("system_user")
+        nodes = {
+          "<node1>": {
+            "system_users": {
+              "system_user": action,
+              "system_user2": action,
+            },
+            "assets": set([<asset_instance>]),
+          }
+        }
+        assets = {
+           "<asset_instance2>": {
+             "system_user": action,
+             "system_user2": action,
+           },
         }
         """
-        self.node_util = NodeUtil()
-        self.nodes = defaultdict(dict)
+        self._node_util = None
+        self.nodes = defaultdict(lambda: {"system_users": defaultdict(int), "assets": set(), "assets_amount": 0})
+        self.assets = defaultdict(lambda: defaultdict(int))
         self._root_node = None
         self._ungroup_node = None
+        self._nodes_with_assets = None
+
+    @property
+    def node_util(self):
+        if not self._node_util:
+            self._node_util = NodeUtil()
+        return self._node_util
 
     @property
     def root_node(self):
@@ -66,35 +106,79 @@ class GenerateTree:
             node_key = '0:0'
         node_value = _("Default")
         node = Node(id=node_id, key=node_key, value=node_value)
-        self.add_node(node)
+        self.add_node(node, {})
         self._ungroup_node = node
         return node
 
-    def add_asset(self, asset, system_users):
+    @timeit
+    def add_assets_without_system_users(self, assets):
+        for asset in assets:
+            self.add_asset(asset, {})
+
+    @timeit
+    def add_assets(self, assets):
+        for asset, system_users in assets.items():
+            self.add_asset(asset, system_users)
+
+    @timeit
+    def add_asset(self, asset, system_users=None):
         nodes = asset.nodes.all()
-        for node in nodes:
-            if node in self.nodes:
-                self.nodes[node][asset].update(system_users)
-            else:
-                self.nodes[self.ungrouped_node][asset].update(system_users)
+        nodes = self.node_util.get_nodes_by_queryset(nodes)
+        if not system_users:
+            system_users = defaultdict(int)
+        else:
+            system_users = {k: v for k, v in system_users.items()}
+        _system_users = self.assets[asset]
+        for system_user, action in _system_users.items():
+            system_users[system_user] |= action
 
-    def get_nodes(self):
-        for node in self.nodes:
-            assets = set(self.nodes.get(node).keys())
-            for n in self.nodes.keys():
-                if n.key.startswith(node.key + ':'):
-                    assets.update(set(self.nodes[n].keys()))
-            node.assets_amount = len(assets)
-        return self.nodes
+        # 获取父节点们
+        parents = self.node_util.get_nodes_parents(nodes, with_self=True)
+        for node in parents:
+            _system_users = self.nodes[node]["system_users"]
+            self.nodes[node]["assets_amount"] += 1
+            for system_user, action in _system_users.items():
+                system_users[system_user] |= action
 
-    def add_node(self, node):
-        self.nodes[node] = defaultdict(set)
+        # 过滤系统用户的协议
+        system_users = {s: v for s, v in system_users.items() if asset.has_protocol(s.protocol)}
+        self.assets[asset] = system_users
+
+        in_nodes = set(self.nodes.keys()) & set(nodes)
+        if not in_nodes:
+            self.nodes[self.ungrouped_node]["assets_amount"] += 1
+            self.nodes[self.ungrouped_node]["assets"].add(system_users)
+            return
+
+        for node in in_nodes:
+            self.nodes[node]["assets"].add(asset)
+
+    def add_node(self, node, system_users=None):
+        if not system_users:
+            system_users = defaultdict(int)
+        self.nodes[node]["system_users"] = system_users
 
     # 添加树节点
+    @timeit
     def add_nodes(self, nodes):
-        need_nodes = self.node_util.get_family(nodes, with_children=True)
-        for node in need_nodes:
-            self.add_node(node)
+        _nodes = nodes.keys()
+        family = self.node_util.get_family(_nodes, with_children=True)
+        for node in family:
+            self.add_node(node, nodes.get(node, {}))
+
+    def get_assets(self):
+        return dict(self.assets)
+
+    @timeit
+    def get_nodes_with_assets(self):
+        if self._nodes_with_assets:
+            return self._nodes_with_assets
+        nodes = {}
+        for node, values in self.nodes.items():
+            node._assets_amount = values["assets_amount"]
+            nodes[node] = {asset: self.assets.get(asset, {}) for asset in values["assets"]}
+        self._nodes_with_assets = nodes
+        return dict(nodes)
 
 
 def get_user_permissions(user, include_group=True):
@@ -131,17 +215,6 @@ def get_system_user_permissions(system_user):
     )
 
 
-def timeit(func):
-    def wrapper(*args, **kwargs):
-        logger.debug("Start call: {}".format(func.__name__))
-        now = time.time()
-        result = func(*args, **kwargs)
-        using = time.time() - now
-        logger.debug("Call {} end, using: {:.2}".format(func.__name__, using))
-        return result
-    return wrapper
-
-
 class AssetPermissionCacheMixin:
     CACHE_KEY_PREFIX = '_ASSET_PERM_CACHE_'
     CACHE_META_KEY_PREFIX = '_ASSET_PERM_META_KEY_'
@@ -216,6 +289,16 @@ class AssetPermissionCacheMixin:
             cached = cache.get(self.system_key)
         return cached
 
+    def get_assets(self):
+        if self._is_not_using_cache():
+            return self.get_assets_from_cache()
+        elif self._is_refresh_cache():
+            self.expire_cache()
+            return self.get_assets_from_cache()
+        else:
+            self.expire_cache()
+            return self.get_assets_without_cache()
+
     def get_system_users(self):
         if self._is_using_cache():
             return self.get_system_user_from_cache()
@@ -282,57 +365,6 @@ class AssetPermissionCacheMixin:
         cache.delete_pattern(key)
 
 
-class FlatPermissionQueryset(set):
-    def add_many(self, assets_or_nodes, system_users, action, rtp="asset"):
-        print("Add many: {}-{}-{}".format(len(assets_or_nodes), len(system_users), action))
-        if not any([assets_or_nodes, system_users, action]):
-            return
-
-        iterable = itertools.product(assets_or_nodes, system_users, [action])
-        for source, sysuser, action in iterable:
-            permission = FlatPermission(source, sysuser, action, rtp=rtp)
-            print("ADDDDDDDDDDDDDDDd")
-            self.add(permission)
-
-    def group_by_resource(self):
-        resources = defaultdict(lambda: defaultdict(int))
-        for i in self:
-            resources[i.resource][i.system_user] |= i.action
-        return resources
-
-
-class FlatPermission:
-    def __init__(self, assets_or_node, system_user, action, rtp="asset"):
-        self.id = "{}_{}_{}".format(assets_or_node.id, system_user.id, action)
-        self.resource = assets_or_node
-        self.resource_type = rtp
-        self.system_user = system_user
-        self.action = action
-
-    def __eq__(self, other):
-        if self.id == other.id:
-            return True
-        # 资产不同
-        if self.resource_type == "asset" and self.id != other.id:
-            return False
-        # 不是子节点
-        elif self.resource_type == "node" and not other.resource.key.startswith(self.resource.key):
-            return False
-        # 系统用户优先级大于后者，则相同
-        if self.system_user.priority > self.system_user.priority:
-            return True
-        # 如果系统用户不同，则不同
-        elif self.system_user != other.system_user:
-            return False
-        # 如果action为与后的结果则相同
-        if self.action == self.action | other.action:
-            return True
-        return False
-
-    def __hash__(self):
-        return hash(self.id)
-
-
 class AssetPermissionUtil(AssetPermissionCacheMixin):
     get_permissions_map = {
         "User": get_user_permissions,
@@ -353,6 +385,7 @@ class AssetPermissionUtil(AssetPermissionCacheMixin):
         self.tree = GenerateTree()
         self.change_org_if_need()
         self.nodes = None
+        self._nodes = None
 
     @staticmethod
     def change_org_if_need():
@@ -380,29 +413,32 @@ class AssetPermissionUtil(AssetPermissionCacheMixin):
         返回用户/组授权规则直接关联的节点
         :return: {node1: {system_user1: {'actions': set()},}}
         """
-        queryset = FlatPermissionQueryset()
+        nodes = defaultdict(lambda: defaultdict(int))
         for perm in self.permissions:
-            actions = perm.action
+            actions = [perm.actions]
             system_users = perm.system_users.all()
-            nodes = perm.nodes.all()
-            queryset.add_many(nodes, system_users, actions, rtp="nodes")
-        print(queryset)
-        return queryset.group_by_resource()
+            _nodes = perm.nodes.all()
+            for node, system_user, action in itertools.product(_nodes, system_users, actions):
+                nodes[node][system_user] |= action
+        self.tree.add_nodes(nodes)
+        return nodes
 
     @timeit
     def get_assets_direct(self):
         """
         返回用户授权规则直接关联的资产
-        :return: {asset1: {system_user1: {'actions': set()},}}
+        :return: {asset1: {system_user1: 1,}}
         """
-        queryset = FlatPermissionQueryset()
+        assets = defaultdict(lambda: defaultdict(int))
         for perm in self.permissions:
-            action = perm.action
-            assets = perm.assets.all()
+            actions = [perm.actions]
+            _assets = perm.assets.all().prefetch_related('nodes', 'protocols')
             system_users = perm.system_users.all()
-            queryset.add_many(assets, system_users, action, rtp="assets")
-        print(queryset)
-        return queryset.group_by_resource()
+            iterable = itertools.product(_assets, system_users, actions)
+            for asset, system_user, action in iterable:
+                assets[asset][system_user] |= action
+        self.tree.add_assets(assets)
+        return assets
 
     @timeit
     def get_assets_without_cache(self):
@@ -411,24 +447,34 @@ class AssetPermissionUtil(AssetPermissionCacheMixin):
         """
         if self._assets:
             return self._assets
-        assets = self.get_assets_direct()
         nodes = self.get_nodes_direct()
-        print("++++++++++++++++++++++")
-        print(assets)
-        print("---------------------")
-        print(nodes)
+        pattern = set()
+        for node in nodes:
+            pattern.add(r'^{0}$|^{0}:'.format(node.key))
+        pattern = '|'.join(list(pattern))
+        now = time.time()
+        print("Get node assets start")
+        if pattern:
+            assets = Asset.objects.filter(nodes__key__regex=pattern)\
+                .prefetch_related('nodes', "protocols").only('id', 'hostname', 'ip').distinct()
+        else:
+            assets = []
+        assets = list(assets)
+        print("Get node assets end, using: {}".format(time.time() - now))
+        self.tree.add_assets_without_system_users(assets)
+        assets = self.tree.get_assets()
+        self._assets = assets
+        return assets
 
     @timeit
     def get_nodes_with_assets_without_cache(self):
         """
         返回节点并且包含资产
-        {"node": {"assets": set("system_user")}}
+        {"node": {"asset": {"system_user": 1})}}
         :return:
         """
-        assets = self.get_assets_without_cache()
-        for asset, system_users in assets.items():
-            self.tree.add_asset(asset, system_users)
-        return self.tree.get_nodes()
+        self.get_assets_without_cache()
+        return self.tree.get_nodes_with_assets()
 
     def get_system_user_without_cache(self):
         system_users = set()
@@ -460,9 +506,7 @@ def sort_assets(assets, order_by='hostname', reverse=False):
 
 
 def parse_node_to_tree_node(node):
-    from .. import serializers
     name = '{} ({})'.format(node.value, node.assets_amount)
-    node_serializer = serializers.GrantedNodeSerializer(node)
     data = {
         'id': node.key,
         'name': name,
@@ -471,7 +515,11 @@ def parse_node_to_tree_node(node):
         'isParent': True,
         'open': node.is_root(),
         'meta': {
-            'node': node_serializer.data,
+            'node': {
+                "id": node.id,
+                "key": node.key,
+                "value": node.value,
+            },
             'type': 'node'
         }
     }
@@ -480,23 +528,21 @@ def parse_node_to_tree_node(node):
 
 
 def parse_asset_to_tree_node(node, asset, system_users):
-    system_users_protocol_matched = [s for s in system_users if asset.has_protocol(s.protocol)]
     icon_skin = 'file'
     if asset.platform.lower() == 'windows':
         icon_skin = 'windows'
     elif asset.platform.lower() == 'linux':
         icon_skin = 'linux'
-    system_users = []
-    for system_user in system_users_protocol_matched:
-        system_users.append({
+    _system_users = []
+    for system_user, action in system_users.items():
+        _system_users.append({
             'id': system_user.id,
             'name': system_user.name,
             'username': system_user.username,
             'protocol': system_user.protocol,
             'priority': system_user.priority,
             'login_mode': system_user.login_mode,
-            'actions': [action.name for action in system_user.actions],
-            'comment': system_user.comment,
+            'actions': [ActionFlag.value_to_choices(action)],
         })
     data = {
         'id': str(asset.id),
@@ -507,7 +553,7 @@ def parse_asset_to_tree_node(node, asset, system_users):
         'open': False,
         'iconSkin': icon_skin,
         'meta': {
-            'system_users': system_users,
+            'system_users': _system_users,
             'type': 'asset',
             'asset': {
                 'id': asset.id,