mirror of
https://github.com/jumpserver/jumpserver.git
synced 2025-12-25 21:42:37 +00:00
Compare commits
21 Commits
pr@v5@perf
...
v5_refacto
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a474d9be3e | ||
|
|
3a4e93af2f | ||
|
|
9c2ddbba7e | ||
|
|
39129cecbe | ||
|
|
88819bbf26 | ||
|
|
a88e35156a | ||
|
|
22a27946a7 | ||
|
|
4983465a23 | ||
|
|
4d9fc9dfd6 | ||
|
|
c7cb83fa1d | ||
|
|
ee92c72b50 | ||
|
|
6a05fbe0fe | ||
|
|
0284be169a | ||
|
|
a4e9d4f815 | ||
|
|
bbe549696a | ||
|
|
56f720271a | ||
|
|
9755076f7f | ||
|
|
8d7abef191 | ||
|
|
aaa40722c4 | ||
|
|
ca39344937 | ||
|
|
4b9a8227c9 |
@@ -1,4 +1,4 @@
|
|||||||
FROM jumpserver/core-base:20251225_092635 AS stage-build
|
FROM jumpserver/core-base:20251128_025056 AS stage-build
|
||||||
|
|
||||||
ARG VERSION
|
ARG VERSION
|
||||||
|
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ from rest_framework.request import Request
|
|||||||
from assets.models import Node, Platform, Protocol, MyAsset
|
from assets.models import Node, Platform, Protocol, MyAsset
|
||||||
from assets.utils import get_node_from_request, is_query_node_all_assets
|
from assets.utils import get_node_from_request, is_query_node_all_assets
|
||||||
from common.utils import lazyproperty, timeit
|
from common.utils import lazyproperty, timeit
|
||||||
|
from assets.tree.asset_tree import AssetTreeNode
|
||||||
|
|
||||||
|
|
||||||
class SerializeToTreeNodeMixin:
|
class SerializeToTreeNodeMixin:
|
||||||
@@ -19,22 +20,22 @@ class SerializeToTreeNodeMixin:
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
@timeit
|
@timeit
|
||||||
def serialize_nodes(self, nodes: List[Node], with_asset_amount=False):
|
def serialize_nodes(self, nodes: List[AssetTreeNode], with_asset_amount=False, expand_level=1, with_assets=False):
|
||||||
if with_asset_amount:
|
if not nodes:
|
||||||
def _name(node: Node):
|
return []
|
||||||
return '{} ({})'.format(node.value, node.assets_amount)
|
|
||||||
else:
|
|
||||||
def _name(node: Node):
|
|
||||||
return node.value
|
|
||||||
|
|
||||||
def _open(node):
|
def _name(node: AssetTreeNode):
|
||||||
if not self.is_sync:
|
v = node.value
|
||||||
# 异步加载资产树时,默认展开节点
|
if not with_asset_amount:
|
||||||
return True
|
return v
|
||||||
if not node.parent_key:
|
v = f'{v} ({node.assets_amount_total})'
|
||||||
return True
|
return v
|
||||||
|
|
||||||
|
def is_parent(node: AssetTreeNode):
|
||||||
|
if with_assets:
|
||||||
|
return node.assets_amount > 0 or not node.is_leaf
|
||||||
else:
|
else:
|
||||||
return False
|
return not node.is_leaf
|
||||||
|
|
||||||
data = [
|
data = [
|
||||||
{
|
{
|
||||||
@@ -42,15 +43,17 @@ class SerializeToTreeNodeMixin:
|
|||||||
'name': _name(node),
|
'name': _name(node),
|
||||||
'title': _name(node),
|
'title': _name(node),
|
||||||
'pId': node.parent_key,
|
'pId': node.parent_key,
|
||||||
'isParent': True,
|
'isParent': is_parent(node),
|
||||||
'open': _open(node),
|
'open': node.level <= expand_level,
|
||||||
'meta': {
|
'meta': {
|
||||||
|
'type': 'node',
|
||||||
'data': {
|
'data': {
|
||||||
"id": node.id,
|
"id": node.id,
|
||||||
"key": node.key,
|
"key": node.key,
|
||||||
"value": node.value,
|
"value": node.value,
|
||||||
|
"assets_amount": node.assets_amount,
|
||||||
|
"assets_amount_total": node.assets_amount_total,
|
||||||
},
|
},
|
||||||
'type': 'node'
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
for node in nodes
|
for node in nodes
|
||||||
@@ -72,6 +75,9 @@ class SerializeToTreeNodeMixin:
|
|||||||
|
|
||||||
@timeit
|
@timeit
|
||||||
def serialize_assets(self, assets, node_key=None, get_pid=None):
|
def serialize_assets(self, assets, node_key=None, get_pid=None):
|
||||||
|
if not assets:
|
||||||
|
return []
|
||||||
|
|
||||||
if not get_pid and not node_key:
|
if not get_pid and not node_key:
|
||||||
get_pid = lambda asset, platform: getattr(asset, 'parent_key', '')
|
get_pid = lambda asset, platform: getattr(asset, 'parent_key', '')
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
# ~*~ coding: utf-8 ~*~
|
# ~*~ coding: utf-8 ~*~
|
||||||
|
|
||||||
from django.db.models import Q
|
from django.db.models import Q, Count
|
||||||
from django.utils.translation import gettext_lazy as _
|
from django.utils.translation import gettext_lazy as _
|
||||||
from rest_framework.generics import get_object_or_404
|
from rest_framework.generics import get_object_or_404
|
||||||
from rest_framework.response import Response
|
from rest_framework.response import Response
|
||||||
@@ -11,12 +11,16 @@ from common.tree import TreeNodeSerializer
|
|||||||
from common.utils import get_logger
|
from common.utils import get_logger
|
||||||
from orgs.mixins import generics
|
from orgs.mixins import generics
|
||||||
from orgs.utils import current_org
|
from orgs.utils import current_org
|
||||||
|
from orgs.models import Organization
|
||||||
from .mixin import SerializeToTreeNodeMixin
|
from .mixin import SerializeToTreeNodeMixin
|
||||||
from .. import serializers
|
from .. import serializers
|
||||||
from ..const import AllTypes
|
from ..const import AllTypes
|
||||||
from ..models import Node, Platform, Asset
|
from ..models import Node, Platform, Asset
|
||||||
|
from assets.tree.asset_tree import AssetTree
|
||||||
|
|
||||||
|
|
||||||
logger = get_logger(__file__)
|
logger = get_logger(__file__)
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
'NodeChildrenApi',
|
'NodeChildrenApi',
|
||||||
'NodeChildrenAsTreeApi',
|
'NodeChildrenAsTreeApi',
|
||||||
@@ -25,14 +29,13 @@ __all__ = [
|
|||||||
|
|
||||||
|
|
||||||
class NodeChildrenApi(generics.ListCreateAPIView):
|
class NodeChildrenApi(generics.ListCreateAPIView):
|
||||||
"""
|
''' 节点的增删改查 '''
|
||||||
节点的增删改查
|
|
||||||
"""
|
|
||||||
serializer_class = serializers.NodeSerializer
|
serializer_class = serializers.NodeSerializer
|
||||||
search_fields = ('value',)
|
search_fields = ('value',)
|
||||||
|
|
||||||
instance = None
|
instance = None
|
||||||
is_initial = False
|
is_initial = False
|
||||||
|
perm_model = Node
|
||||||
|
|
||||||
def initial(self, request, *args, **kwargs):
|
def initial(self, request, *args, **kwargs):
|
||||||
super().initial(request, *args, **kwargs)
|
super().initial(request, *args, **kwargs)
|
||||||
@@ -65,42 +68,16 @@ class NodeChildrenApi(generics.ListCreateAPIView):
|
|||||||
else:
|
else:
|
||||||
node = Node.org_root()
|
node = Node.org_root()
|
||||||
return node
|
return node
|
||||||
|
|
||||||
if pk:
|
if pk:
|
||||||
node = get_object_or_404(Node, pk=pk)
|
node = get_object_or_404(Node, pk=pk)
|
||||||
else:
|
else:
|
||||||
node = get_object_or_404(Node, key=key)
|
node = get_object_or_404(Node, key=key)
|
||||||
return node
|
return node
|
||||||
|
|
||||||
def get_org_root_queryset(self, query_all):
|
|
||||||
if query_all:
|
|
||||||
return Node.objects.all()
|
|
||||||
else:
|
|
||||||
return Node.org_root_nodes()
|
|
||||||
|
|
||||||
def get_queryset(self):
|
|
||||||
query_all = self.request.query_params.get("all", "0") == "all"
|
|
||||||
|
|
||||||
if self.is_initial and current_org.is_root():
|
|
||||||
return self.get_org_root_queryset(query_all)
|
|
||||||
|
|
||||||
if self.is_initial:
|
|
||||||
with_self = True
|
|
||||||
else:
|
|
||||||
with_self = False
|
|
||||||
|
|
||||||
if not self.instance:
|
|
||||||
return Node.objects.none()
|
|
||||||
|
|
||||||
if query_all:
|
|
||||||
queryset = self.instance.get_all_children(with_self=with_self)
|
|
||||||
else:
|
|
||||||
queryset = self.instance.get_children(with_self=with_self)
|
|
||||||
return queryset
|
|
||||||
|
|
||||||
|
|
||||||
class NodeChildrenAsTreeApi(SerializeToTreeNodeMixin, NodeChildrenApi):
|
class NodeChildrenAsTreeApi(SerializeToTreeNodeMixin, NodeChildrenApi):
|
||||||
"""
|
''' 节点子节点作为树返回,
|
||||||
节点子节点作为树返回,
|
|
||||||
[
|
[
|
||||||
{
|
{
|
||||||
"id": "",
|
"id": "",
|
||||||
@@ -109,51 +86,96 @@ class NodeChildrenAsTreeApi(SerializeToTreeNodeMixin, NodeChildrenApi):
|
|||||||
"meta": ""
|
"meta": ""
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
'''
|
||||||
|
|
||||||
"""
|
|
||||||
model = Node
|
model = Node
|
||||||
|
|
||||||
def filter_queryset(self, queryset):
|
|
||||||
""" queryset is Node queryset """
|
|
||||||
if not self.request.GET.get('search'):
|
|
||||||
return queryset
|
|
||||||
queryset = super().filter_queryset(queryset)
|
|
||||||
queryset = self.model.get_ancestor_queryset(queryset)
|
|
||||||
return queryset
|
|
||||||
|
|
||||||
def get_queryset_for_assets(self):
|
|
||||||
query_all = self.request.query_params.get("all", "0") == "all"
|
|
||||||
include_assets = self.request.query_params.get('assets', '0') == '1'
|
|
||||||
if not self.instance or not include_assets:
|
|
||||||
return Asset.objects.none()
|
|
||||||
if not self.request.GET.get('search') and self.instance.is_org_root():
|
|
||||||
return Asset.objects.none()
|
|
||||||
if query_all:
|
|
||||||
assets = self.instance.get_all_assets()
|
|
||||||
else:
|
|
||||||
assets = self.instance.get_assets()
|
|
||||||
return assets.only(
|
|
||||||
"id", "name", "address", "platform_id",
|
|
||||||
"org_id", "is_active", 'comment'
|
|
||||||
).prefetch_related('platform')
|
|
||||||
|
|
||||||
def filter_queryset_for_assets(self, assets):
|
|
||||||
search = self.request.query_params.get('search')
|
|
||||||
if search:
|
|
||||||
q = Q(name__icontains=search) | Q(address__icontains=search)
|
|
||||||
assets = assets.filter(q)
|
|
||||||
return assets
|
|
||||||
|
|
||||||
def list(self, request, *args, **kwargs):
|
def list(self, request, *args, **kwargs):
|
||||||
nodes = self.filter_queryset(self.get_queryset()).order_by('value')
|
search = request.query_params.get('search')
|
||||||
|
with_assets = request.query_params.get('assets', '0') == '1'
|
||||||
with_asset_amount = request.query_params.get('asset_amount', '1') == '1'
|
with_asset_amount = request.query_params.get('asset_amount', '1') == '1'
|
||||||
nodes = self.serialize_nodes(nodes, with_asset_amount=with_asset_amount)
|
with_asset_amount = True
|
||||||
assets = self.filter_queryset_for_assets(self.get_queryset_for_assets())
|
nodes, assets, expand_level = self.get_nodes_assets(search, with_assets)
|
||||||
node_key = self.instance.key if self.instance else None
|
nodes = self.serialize_nodes(nodes, with_asset_amount=with_asset_amount, expand_level=expand_level)
|
||||||
assets = self.serialize_assets(assets, node_key=node_key)
|
assets = self.serialize_assets(assets)
|
||||||
data = [*nodes, *assets]
|
data = [*nodes, *assets]
|
||||||
return Response(data=data)
|
return Response(data=data)
|
||||||
|
|
||||||
|
def get_nodes_assets(self, search, with_assets):
|
||||||
|
#
|
||||||
|
# 资产管理-节点树
|
||||||
|
#
|
||||||
|
|
||||||
|
# 全局组织: 初始化节点树, 返回所有节点, 不包含资产, 不展开节点
|
||||||
|
# 实体组织: 初始化节点树, 返回所有节点, 不包含资产, 展开一级节点
|
||||||
|
# 前端搜索
|
||||||
|
if not with_assets:
|
||||||
|
if current_org.is_root():
|
||||||
|
orgs = Organization.objects.all()
|
||||||
|
expand_level = 0
|
||||||
|
else:
|
||||||
|
orgs = [current_org]
|
||||||
|
expand_level = 1
|
||||||
|
|
||||||
|
nodes = []
|
||||||
|
assets = []
|
||||||
|
for org in orgs:
|
||||||
|
tree = AssetTree(org=org)
|
||||||
|
org_nodes = tree.get_nodes()
|
||||||
|
nodes.extend(org_nodes)
|
||||||
|
return nodes, assets, expand_level
|
||||||
|
|
||||||
|
#
|
||||||
|
# 权限管理、账号发现、风险检测 - 资产节点树
|
||||||
|
#
|
||||||
|
|
||||||
|
# 全局组织: 搜索资产, 生成资产节点树, 过滤每个组织前 1000 个资产, 展开所有节点
|
||||||
|
# 实体组织: 搜索资产, 生成资产节点树, 过滤前 1000 个资产, 展开所有节点
|
||||||
|
if search:
|
||||||
|
if current_org.is_root():
|
||||||
|
orgs = list(Organization.objects.all())
|
||||||
|
else:
|
||||||
|
orgs = [current_org]
|
||||||
|
nodes = []
|
||||||
|
assets = []
|
||||||
|
assets_q_object = Q(name__icontains=search) | Q(address__icontains=search)
|
||||||
|
with_assets_limit = 1000 / len(orgs)
|
||||||
|
for org in orgs:
|
||||||
|
tree = AssetTree(
|
||||||
|
assets_q_object=assets_q_object, org=org,
|
||||||
|
with_assets=True, with_assets_limit=with_assets_limit, full_tree=False
|
||||||
|
)
|
||||||
|
nodes.extend(tree.get_nodes())
|
||||||
|
assets.extend(tree.get_assets())
|
||||||
|
expand_level = 10000 # search 时展开所有节点
|
||||||
|
return nodes, assets, expand_level
|
||||||
|
|
||||||
|
# 全局组织: 展开某个节点及其资产
|
||||||
|
# 实体组织: 展开某个节点及其资产
|
||||||
|
# 实体组织: 初始化资产节点树, 自动展开根节点及其资产, 所以节点要包含自己 (特殊情况)
|
||||||
|
if self.instance:
|
||||||
|
nodes = []
|
||||||
|
tree = AssetTree(with_assets_node_id=self.instance.id, org=self.instance.org)
|
||||||
|
nodes_with_self = False
|
||||||
|
if not current_org.is_root() and self.instance.is_org_root():
|
||||||
|
nodes_with_self = True
|
||||||
|
nodes = tree.get_node_children(key=self.instance.key, with_self=nodes_with_self)
|
||||||
|
assets = tree.get_assets()
|
||||||
|
expand_level = 1 # 默认只展开第一级
|
||||||
|
return nodes, assets, expand_level
|
||||||
|
|
||||||
|
# 全局组织: 初始化资产节点树, 仅返回各组织根节点, 不展开
|
||||||
|
orgs = Organization.objects.all()
|
||||||
|
nodes = []
|
||||||
|
assets = []
|
||||||
|
for org in orgs:
|
||||||
|
tree = AssetTree(org=org, with_assets=False)
|
||||||
|
if not tree.root:
|
||||||
|
continue
|
||||||
|
nodes.append(tree.root)
|
||||||
|
expand_level = 0 # 默认不展开节点
|
||||||
|
return nodes, assets, expand_level
|
||||||
|
|
||||||
|
|
||||||
class CategoryTreeApi(SerializeToTreeNodeMixin, generics.ListAPIView):
|
class CategoryTreeApi(SerializeToTreeNodeMixin, generics.ListAPIView):
|
||||||
serializer_class = TreeNodeSerializer
|
serializer_class = TreeNodeSerializer
|
||||||
|
|||||||
@@ -63,11 +63,11 @@ class NodeFilterBackend(filters.BaseFilterBackend):
|
|||||||
query_all = is_query_node_all_assets(request)
|
query_all = is_query_node_all_assets(request)
|
||||||
if query_all:
|
if query_all:
|
||||||
return queryset.filter(
|
return queryset.filter(
|
||||||
Q(nodes__key__startswith=f'{node.key}:') |
|
Q(node__key__startswith=f'{node.key}:') |
|
||||||
Q(nodes__key=node.key)
|
Q(node__key=node.key)
|
||||||
).distinct()
|
).distinct()
|
||||||
else:
|
else:
|
||||||
return queryset.filter(nodes__key=node.key).distinct()
|
return queryset.filter(node__key=node.key).distinct()
|
||||||
|
|
||||||
|
|
||||||
class IpInFilterBackend(filters.BaseFilterBackend):
|
class IpInFilterBackend(filters.BaseFilterBackend):
|
||||||
|
|||||||
126
apps/assets/migrations/0020_asset_node.py
Normal file
126
apps/assets/migrations/0020_asset_node.py
Normal file
@@ -0,0 +1,126 @@
|
|||||||
|
# Generated by Django 4.1.13 on 2025-12-16 09:14
|
||||||
|
|
||||||
|
from django.db import migrations, models, transaction
|
||||||
|
import django.db.models.deletion
|
||||||
|
|
||||||
|
|
||||||
|
def log(msg=''):
|
||||||
|
print(f' -> {msg}')
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_asset_single_node(apps, schema_editor):
|
||||||
|
print('')
|
||||||
|
log('Checking that all assets are linked to only one node...')
|
||||||
|
Asset = apps.get_model('assets', 'Asset')
|
||||||
|
Through = Asset.nodes.through
|
||||||
|
|
||||||
|
assets_count_multi_nodes = Through.objects.values('asset_id').annotate(
|
||||||
|
node_count=models.Count('node_id')
|
||||||
|
).filter(node_count__gt=1).count()
|
||||||
|
|
||||||
|
if assets_count_multi_nodes > 0:
|
||||||
|
raise Exception(
|
||||||
|
f'There are {assets_count_multi_nodes} assets associated with more than one node. '
|
||||||
|
'Please ensure each asset is linked to only one node before applying this migration.'
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
log('All assets are linked to only one node. Proceeding with the migration.')
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_asset_has_node(apps, schema_editor):
|
||||||
|
log('Checking that all assets are linked to at least one node...')
|
||||||
|
Asset = apps.get_model('assets', 'Asset')
|
||||||
|
Through = Asset.nodes.through
|
||||||
|
|
||||||
|
asset_count = Asset.objects.count()
|
||||||
|
through_asset_count = Through.objects.values('asset_id').count()
|
||||||
|
|
||||||
|
assets_count_without_node = asset_count - through_asset_count
|
||||||
|
|
||||||
|
if assets_count_without_node > 0:
|
||||||
|
raise Exception(
|
||||||
|
f'Some assets ({assets_count_without_node}) are not associated with any node. '
|
||||||
|
'Please ensure all assets are linked to a node before applying this migration.'
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
log('All assets are linked to a node. Proceeding with the migration.')
|
||||||
|
|
||||||
|
|
||||||
|
def migrate_asset_node_id_field(apps, schema_editor):
|
||||||
|
log('Migrating node_id field for all assets...')
|
||||||
|
|
||||||
|
Asset = apps.get_model('assets', 'Asset')
|
||||||
|
Through = Asset.nodes.through
|
||||||
|
|
||||||
|
assets = Asset.objects.filter(node_id__isnull=True)
|
||||||
|
log (f'Found {assets.count()} assets to migrate.')
|
||||||
|
|
||||||
|
asset_node_mapper = {
|
||||||
|
str(asset_id): str(node_id)
|
||||||
|
for asset_id, node_id in Through.objects.values_list('asset_id', 'node_id')
|
||||||
|
}
|
||||||
|
# 测试
|
||||||
|
asset_node_mapper.pop(None, None) # Remove any entries with None keys
|
||||||
|
|
||||||
|
for asset in assets:
|
||||||
|
node_id = asset_node_mapper.get(str(asset.id))
|
||||||
|
if not node_id:
|
||||||
|
raise Exception(
|
||||||
|
f'Asset (ID: {asset.id}) is not associated with any node. '
|
||||||
|
'Cannot migrate node_id field.'
|
||||||
|
)
|
||||||
|
asset.node_id = node_id
|
||||||
|
|
||||||
|
with transaction.atomic():
|
||||||
|
total = len(assets)
|
||||||
|
batch_size = 5000
|
||||||
|
|
||||||
|
for i in range(0, total, batch_size):
|
||||||
|
batch = assets[i:i+batch_size]
|
||||||
|
start = i + 1
|
||||||
|
end = min(i + batch_size, total)
|
||||||
|
|
||||||
|
for asset in batch:
|
||||||
|
asset.save(update_fields=['node_id'])
|
||||||
|
|
||||||
|
log(f"Migrated {start}-{end}/{total} assets")
|
||||||
|
|
||||||
|
count = Asset.objects.filter(node_id__isnull=True).count()
|
||||||
|
if count > 0:
|
||||||
|
log('Warning: Some assets still have null node_id after migration.')
|
||||||
|
raise Exception('Migration failed: Some assets have null node_id.')
|
||||||
|
|
||||||
|
count = Asset.objects.filter(node_id__isnull=False).count()
|
||||||
|
log(f'Successfully migrated node_id for {count} assets.')
|
||||||
|
|
||||||
|
|
||||||
|
class Migration(migrations.Migration):
|
||||||
|
|
||||||
|
dependencies = [
|
||||||
|
('assets', '0019_alter_asset_connectivity'),
|
||||||
|
]
|
||||||
|
|
||||||
|
operations = [
|
||||||
|
migrations.RunPython(
|
||||||
|
ensure_asset_single_node,
|
||||||
|
reverse_code=migrations.RunPython.noop
|
||||||
|
),
|
||||||
|
migrations.RunPython(
|
||||||
|
ensure_asset_has_node,
|
||||||
|
reverse_code=migrations.RunPython.noop
|
||||||
|
),
|
||||||
|
migrations.AddField(
|
||||||
|
model_name='asset',
|
||||||
|
name='node',
|
||||||
|
field=models.ForeignKey(null=True, on_delete=django.db.models.deletion.PROTECT, related_name='direct_assets', to='assets.node', verbose_name='Node'),
|
||||||
|
),
|
||||||
|
migrations.RunPython(
|
||||||
|
migrate_asset_node_id_field,
|
||||||
|
reverse_code=migrations.RunPython.noop
|
||||||
|
),
|
||||||
|
migrations.AlterField(
|
||||||
|
model_name='asset',
|
||||||
|
name='node',
|
||||||
|
field=models.ForeignKey(on_delete=django.db.models.deletion.PROTECT, related_name='direct_assets', to='assets.node', verbose_name='Node'),
|
||||||
|
),
|
||||||
|
]
|
||||||
@@ -172,6 +172,11 @@ class Asset(NodesRelationMixin, LabeledMixin, AbsConnectivity, JSONFilterMixin,
|
|||||||
"assets.Zone", null=True, blank=True, related_name='assets',
|
"assets.Zone", null=True, blank=True, related_name='assets',
|
||||||
verbose_name=_("Zone"), on_delete=models.SET_NULL
|
verbose_name=_("Zone"), on_delete=models.SET_NULL
|
||||||
)
|
)
|
||||||
|
node = models.ForeignKey(
|
||||||
|
'assets.Node', null=False, blank=False, on_delete=models.PROTECT,
|
||||||
|
related_name='direct_assets', verbose_name=_("Node")
|
||||||
|
)
|
||||||
|
# TODO: 删除完代码中所有使用的地方后,再删除 nodes 字段,并将 node 字段的 related_name 改为 'assets'
|
||||||
nodes = models.ManyToManyField(
|
nodes = models.ManyToManyField(
|
||||||
'assets.Node', default=default_node, related_name='assets', verbose_name=_("Nodes")
|
'assets.Node', default=default_node, related_name='assets', verbose_name=_("Nodes")
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -394,7 +394,7 @@ class NodeAssetsMixin(NodeAllAssetsMappingMixin):
|
|||||||
|
|
||||||
def get_all_assets(self):
|
def get_all_assets(self):
|
||||||
from .asset import Asset
|
from .asset import Asset
|
||||||
q = Q(nodes__key__startswith=f'{self.key}:') | Q(nodes__key=self.key)
|
q = Q(node__key__startswith=f'{self.key}:') | Q(node__key=self.key)
|
||||||
return Asset.objects.filter(q).distinct()
|
return Asset.objects.filter(q).distinct()
|
||||||
|
|
||||||
def get_assets_amount(self):
|
def get_assets_amount(self):
|
||||||
@@ -416,8 +416,8 @@ class NodeAssetsMixin(NodeAllAssetsMappingMixin):
|
|||||||
|
|
||||||
def get_assets(self):
|
def get_assets(self):
|
||||||
from .asset import Asset
|
from .asset import Asset
|
||||||
assets = Asset.objects.filter(nodes=self)
|
assets = Asset.objects.filter(node=self)
|
||||||
return assets.distinct()
|
return assets
|
||||||
|
|
||||||
def get_valid_assets(self):
|
def get_valid_assets(self):
|
||||||
return self.get_assets().valid()
|
return self.get_assets().valid()
|
||||||
@@ -531,6 +531,15 @@ class SomeNodesMixin:
|
|||||||
root_nodes = cls.objects.filter(parent_key='', key__regex=r'^[0-9]+$') \
|
root_nodes = cls.objects.filter(parent_key='', key__regex=r'^[0-9]+$') \
|
||||||
.exclude(key__startswith='-').order_by('key')
|
.exclude(key__startswith='-').order_by('key')
|
||||||
return root_nodes
|
return root_nodes
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def get_or_create_org_root(cls, org):
|
||||||
|
org_root = cls.org_root_nodes().filter(org_id=org.id).first()
|
||||||
|
if org_root:
|
||||||
|
return org_root
|
||||||
|
with tmp_to_org(org):
|
||||||
|
org_root = cls.create_org_root_node()
|
||||||
|
return org_root
|
||||||
|
|
||||||
|
|
||||||
class Node(JMSOrgBaseModel, SomeNodesMixin, FamilyMixin, NodeAssetsMixin):
|
class Node(JMSOrgBaseModel, SomeNodesMixin, FamilyMixin, NodeAssetsMixin):
|
||||||
|
|||||||
1
apps/assets/tree/__init__.py
Normal file
1
apps/assets/tree/__init__.py
Normal file
@@ -0,0 +1 @@
|
|||||||
|
from .asset_tree import *
|
||||||
246
apps/assets/tree/asset_tree.py
Normal file
246
apps/assets/tree/asset_tree.py
Normal file
@@ -0,0 +1,246 @@
|
|||||||
|
from collections import defaultdict
|
||||||
|
from django.db.models import Count, Q
|
||||||
|
|
||||||
|
from orgs.utils import current_org
|
||||||
|
from orgs.models import Organization
|
||||||
|
from assets.models import Asset, Node, Platform
|
||||||
|
from assets.const.category import Category
|
||||||
|
from common.utils import get_logger, timeit, lazyproperty
|
||||||
|
|
||||||
|
from .tree import TreeNode, Tree
|
||||||
|
|
||||||
|
logger = get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ['AssetTree', 'AssetTreeNode']
|
||||||
|
|
||||||
|
|
||||||
|
class AssetTreeNodeAsset:
|
||||||
|
|
||||||
|
def __init__(self, id, node_id, parent_key, name, address,
|
||||||
|
platform_id, is_active, comment, org_id):
|
||||||
|
|
||||||
|
self.id = id
|
||||||
|
self.node_id = node_id
|
||||||
|
self.parent_key = parent_key
|
||||||
|
self.name = name
|
||||||
|
self.address = address
|
||||||
|
self.platform_id = platform_id
|
||||||
|
self.is_active = is_active
|
||||||
|
self.comment = comment
|
||||||
|
self.org_id = org_id
|
||||||
|
|
||||||
|
@lazyproperty
|
||||||
|
def org(self):
|
||||||
|
return Organization.get_instance(self.org_id)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def org_name(self) -> str:
|
||||||
|
return self.org.name
|
||||||
|
|
||||||
|
|
||||||
|
class AssetTreeNode(TreeNode):
|
||||||
|
|
||||||
|
def __init__(self, _id, key, value, assets_amount=0, assets=None):
|
||||||
|
super().__init__(_id, key, value)
|
||||||
|
self.assets_amount = assets_amount
|
||||||
|
self.assets_amount_total = 0
|
||||||
|
self.assets: list[AssetTreeNodeAsset] = []
|
||||||
|
self.init_assets(assets)
|
||||||
|
|
||||||
|
def init_assets(self, assets):
|
||||||
|
if not assets:
|
||||||
|
return
|
||||||
|
for asset in assets:
|
||||||
|
asset['parent_key'] = self.key
|
||||||
|
self.assets.append(AssetTreeNodeAsset(**asset))
|
||||||
|
|
||||||
|
def get_assets(self):
|
||||||
|
return self.assets
|
||||||
|
|
||||||
|
def as_dict(self, simple=True):
|
||||||
|
data = super().as_dict(simple=simple)
|
||||||
|
data.update({
|
||||||
|
'assets_amount_total': self.assets_amount_total,
|
||||||
|
'assets_amount': self.assets_amount,
|
||||||
|
'assets': len(self.assets),
|
||||||
|
})
|
||||||
|
return data
|
||||||
|
|
||||||
|
|
||||||
|
class AssetTree(Tree):
|
||||||
|
|
||||||
|
TreeNode = AssetTreeNode
|
||||||
|
|
||||||
|
def __init__(self, assets_q_object: Q = None, category=None, org=None,
|
||||||
|
with_assets=False, with_assets_node_id=None, with_assets_limit=1000,
|
||||||
|
full_tree=True):
|
||||||
|
'''
|
||||||
|
:param assets_q_object: 只生成这些资产所在的节点树
|
||||||
|
:param category: 只生成该类别资产所在的节点树
|
||||||
|
:param org: 只生成该组织下的资产节点树
|
||||||
|
|
||||||
|
:param with_assets_node_id: 仅指定节点下包含资产
|
||||||
|
:param with_assets: 所有节点都包含资产
|
||||||
|
:param with_assets_limit: 包含资产时, 所有资产的最大数量
|
||||||
|
|
||||||
|
:param full_tree: 完整树包含所有节点,否则只包含节点的资产总数不为0的节点
|
||||||
|
'''
|
||||||
|
|
||||||
|
super().__init__()
|
||||||
|
## 通过资产构建节点树, 支持 Q, category, org 等过滤条件 ##
|
||||||
|
self._assets_q_object: Q = assets_q_object or Q()
|
||||||
|
self._category = self._check_category(category)
|
||||||
|
self._category_platform_ids = set()
|
||||||
|
self._org: Organization = org or current_org
|
||||||
|
|
||||||
|
# org 下全量节点属性映射, 构建资产树时根据完整的节点进行构建
|
||||||
|
self._nodes_attr_mapper = defaultdict(dict)
|
||||||
|
# 节点直接资产数量映射, 用于计算节点下总资产数量
|
||||||
|
self._nodes_assets_amount_mapper = defaultdict(int)
|
||||||
|
# 节点下是否包含资产
|
||||||
|
self._with_assets = with_assets # 所有节点都包含资产
|
||||||
|
self._with_assets_node_id = with_assets_node_id # 仅指定节点下包含资产, 优先级高于 with_assets
|
||||||
|
self._with_assets_limit = with_assets_limit
|
||||||
|
self._node_assets_mapper = defaultdict(dict)
|
||||||
|
|
||||||
|
# 是否包含资产总数量为 0 的节点
|
||||||
|
self._full_tree = full_tree
|
||||||
|
|
||||||
|
# 初始化时构建树
|
||||||
|
self.build()
|
||||||
|
|
||||||
|
def _check_category(self, category):
|
||||||
|
if category is None:
|
||||||
|
return None
|
||||||
|
if category in Category.values:
|
||||||
|
return category
|
||||||
|
logger.warning(f"Invalid category '{category}' for AssetSearchTree.")
|
||||||
|
return None
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def build(self):
|
||||||
|
self._load_nodes_attr_mapper()
|
||||||
|
self._load_category_platforms_if_needed()
|
||||||
|
self._load_nodes_assets_amount()
|
||||||
|
self._load_nodes_assets_if_needed()
|
||||||
|
self._init_tree()
|
||||||
|
self._compute_assets_amount_total()
|
||||||
|
self._remove_nodes_with_zero_assets_if_needed()
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _load_category_platforms_if_needed(self):
|
||||||
|
if self._category is None:
|
||||||
|
return
|
||||||
|
ids = Platform.objects.filter(category=self._category).values_list('id', flat=True)
|
||||||
|
ids = self._uuids_to_string(ids)
|
||||||
|
self._category_platform_ids = ids
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _load_nodes_attr_mapper(self):
|
||||||
|
nodes = Node.objects.filter(org_id=self._org.id).values('id', 'key', 'value')
|
||||||
|
# 保证节点按 key 顺序加载,以便后续构建树时父节点总在子节点前面
|
||||||
|
nodes = sorted(nodes, key=lambda n: [int(i) for i in n['key'].split(':')])
|
||||||
|
for node in list(nodes):
|
||||||
|
node['id'] = str(node['id'])
|
||||||
|
self._nodes_attr_mapper[node['id']] = node
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _load_nodes_assets_amount(self):
|
||||||
|
q = self._make_assets_q_object()
|
||||||
|
nodes_amount = Asset.objects.filter(q).values('node_id').annotate(
|
||||||
|
amount=Count('id')
|
||||||
|
).values('node_id', 'amount')
|
||||||
|
for nc in list(nodes_amount):
|
||||||
|
nid = str(nc['node_id'])
|
||||||
|
self._nodes_assets_amount_mapper[nid] = nc['amount']
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _load_nodes_assets_if_needed(self):
|
||||||
|
need_load = self._with_assets or self._with_assets_node_id
|
||||||
|
if not need_load:
|
||||||
|
return
|
||||||
|
|
||||||
|
q = self._make_assets_q_object()
|
||||||
|
if self._with_assets_node_id:
|
||||||
|
# 仅指定节点下包含资产,优先级高于 with_assets
|
||||||
|
q &= Q(node_id=self._with_assets_node_id)
|
||||||
|
|
||||||
|
assets = Asset.objects.filter(q).values(
|
||||||
|
'node_id', 'id', 'platform_id', 'name', 'address', 'is_active', 'comment', 'org_id'
|
||||||
|
)
|
||||||
|
# 按照 node_key 排序,尽可能保证前面节点的资产较多
|
||||||
|
# 限制资产数量
|
||||||
|
assets = assets.order_by('node__key')[:self._with_assets_limit]
|
||||||
|
for asset in list(assets):
|
||||||
|
nid = asset['node_id'] = str(asset['node_id'])
|
||||||
|
aid = asset['id'] = str(asset['id'])
|
||||||
|
self._node_assets_mapper[nid][aid] = asset
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _make_assets_q_object(self) -> Q:
|
||||||
|
q = Q(org_id=self._org.id)
|
||||||
|
if self._category_platform_ids:
|
||||||
|
q &= Q(platform_id__in=self._category_platform_ids)
|
||||||
|
if self._assets_q_object:
|
||||||
|
q &= self._assets_q_object
|
||||||
|
return q
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _init_tree(self):
|
||||||
|
for nid in self._nodes_attr_mapper.keys():
|
||||||
|
data = self._get_tree_node_data(nid)
|
||||||
|
node = self.TreeNode(**data)
|
||||||
|
self.add_node(node)
|
||||||
|
|
||||||
|
def _get_tree_node_data(self, node_id):
|
||||||
|
attr = self._nodes_attr_mapper[node_id]
|
||||||
|
assets_amount = self._nodes_assets_amount_mapper.get(node_id, 0)
|
||||||
|
data = {
|
||||||
|
'_id': node_id,
|
||||||
|
'key': attr['key'],
|
||||||
|
'value': attr['value'],
|
||||||
|
'assets_amount': assets_amount,
|
||||||
|
}
|
||||||
|
|
||||||
|
assets = self._node_assets_mapper[node_id].values()
|
||||||
|
if assets:
|
||||||
|
assets = list(assets)
|
||||||
|
data.update({ 'assets': assets })
|
||||||
|
return data
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _compute_assets_amount_total(self):
|
||||||
|
for node in reversed(list(self.nodes.values())):
|
||||||
|
total = node.assets_amount
|
||||||
|
for child in node.children:
|
||||||
|
child: AssetTreeNode
|
||||||
|
total += child.assets_amount_total
|
||||||
|
node: AssetTreeNode
|
||||||
|
node.assets_amount_total = total
|
||||||
|
@timeit
|
||||||
|
def _remove_nodes_with_zero_assets_if_needed(self):
|
||||||
|
if self._full_tree:
|
||||||
|
return
|
||||||
|
nodes: list[AssetTreeNode] = list(self.nodes.values())
|
||||||
|
nodes_to_remove = [
|
||||||
|
node for node in nodes if not node.is_root and node.assets_amount_total == 0
|
||||||
|
]
|
||||||
|
for node in nodes_to_remove:
|
||||||
|
self.remove_node(node)
|
||||||
|
|
||||||
|
def get_assets(self):
|
||||||
|
assets = []
|
||||||
|
for node in self.nodes.values():
|
||||||
|
node: AssetTreeNode
|
||||||
|
_assets = node.get_assets()
|
||||||
|
assets.extend(_assets)
|
||||||
|
return assets
|
||||||
|
|
||||||
|
def _uuids_to_string(self, uuids):
|
||||||
|
return [ str(u) for u in uuids ]
|
||||||
|
|
||||||
|
def print(self, count=20, simple=True):
|
||||||
|
print('org_name: ', getattr(self._org, 'name', 'No-org'))
|
||||||
|
print(f'asset_category: {self._category}')
|
||||||
|
super().print(count=count, simple=simple)
|
||||||
164
apps/assets/tree/tree.py
Normal file
164
apps/assets/tree/tree.py
Normal file
@@ -0,0 +1,164 @@
|
|||||||
|
from common.utils import get_logger, lazyproperty
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ['TreeNode', 'Tree']
|
||||||
|
|
||||||
|
logger = get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class TreeNode(object):
|
||||||
|
|
||||||
|
def __init__(self, _id, key, value):
|
||||||
|
self.id = _id
|
||||||
|
self.key = key
|
||||||
|
self.value = value
|
||||||
|
self.children = []
|
||||||
|
self.parent = None
|
||||||
|
|
||||||
|
@lazyproperty
|
||||||
|
def parent_key(self):
|
||||||
|
if self.is_root:
|
||||||
|
return None
|
||||||
|
return ':'.join(self.key.split(':')[:-1])
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_root(self):
|
||||||
|
return self.key.isdigit()
|
||||||
|
|
||||||
|
def add_child(self, child_node: 'TreeNode'):
|
||||||
|
child_node.parent = self
|
||||||
|
self.children.append(child_node)
|
||||||
|
|
||||||
|
def remove_child(self, child_node: 'TreeNode'):
|
||||||
|
self.children.remove(child_node)
|
||||||
|
child_node.parent = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_leaf(self):
|
||||||
|
return len(self.children) == 0
|
||||||
|
|
||||||
|
@lazyproperty
|
||||||
|
def level(self):
|
||||||
|
return self.key.count(':') + 1
|
||||||
|
|
||||||
|
@property
|
||||||
|
def children_count(self):
|
||||||
|
return len(self.children)
|
||||||
|
|
||||||
|
def as_dict(self, simple=True):
|
||||||
|
data = {
|
||||||
|
'key': self.key,
|
||||||
|
}
|
||||||
|
if simple:
|
||||||
|
return data
|
||||||
|
|
||||||
|
data.update({
|
||||||
|
'id': self.id,
|
||||||
|
'value': self.value,
|
||||||
|
'level': self.level,
|
||||||
|
'children_count': self.children_count,
|
||||||
|
'is_root': self.is_root,
|
||||||
|
'is_leaf': self.is_leaf,
|
||||||
|
})
|
||||||
|
return data
|
||||||
|
|
||||||
|
def print(self, simple=True, is_print_keys=False):
|
||||||
|
def info_as_string(_info):
|
||||||
|
return ' | '.join(s.ljust(25) for s in _info)
|
||||||
|
|
||||||
|
if is_print_keys:
|
||||||
|
info_keys = [k for k in self.as_dict(simple=simple).keys()]
|
||||||
|
info_keys_string = info_as_string(info_keys)
|
||||||
|
print('-' * len(info_keys_string))
|
||||||
|
print(info_keys_string)
|
||||||
|
print('-' * len(info_keys_string))
|
||||||
|
|
||||||
|
info_values = [str(v) for v in self.as_dict(simple=simple).values()]
|
||||||
|
info_values_as_string = info_as_string(info_values)
|
||||||
|
print(info_values_as_string)
|
||||||
|
print('-' * len(info_values_as_string))
|
||||||
|
|
||||||
|
|
||||||
|
class Tree(object):
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.root = None
|
||||||
|
# { key -> TreeNode }
|
||||||
|
self.nodes: dict[TreeNode] = {}
|
||||||
|
|
||||||
|
@property
|
||||||
|
def size(self):
|
||||||
|
return len(self.nodes)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_empty(self):
|
||||||
|
return self.size == 0
|
||||||
|
|
||||||
|
@property
|
||||||
|
def depth(self):
|
||||||
|
" 返回树的最大深度,以及对应的节点key "
|
||||||
|
if self.is_empty:
|
||||||
|
return 0, 0
|
||||||
|
node = max(self.nodes.values(), key=lambda node: node.level)
|
||||||
|
node: TreeNode
|
||||||
|
print(f"max_depth_node_key: {node.key}")
|
||||||
|
return node.level
|
||||||
|
|
||||||
|
@property
|
||||||
|
def width(self):
|
||||||
|
" 返回树的最大宽度,以及对应的层级数 "
|
||||||
|
if self.is_empty:
|
||||||
|
return 0, 0
|
||||||
|
node = max(self.nodes.values(), key=lambda node: node.children_count)
|
||||||
|
node: TreeNode
|
||||||
|
print(f"max_width_level: {node.level + 1}")
|
||||||
|
return node.children_count
|
||||||
|
|
||||||
|
def add_node(self, node: TreeNode):
|
||||||
|
if node.is_root:
|
||||||
|
self.root = node
|
||||||
|
self.nodes[node.key] = node
|
||||||
|
return
|
||||||
|
parent = self.get_node(node.parent_key)
|
||||||
|
if not parent:
|
||||||
|
error = f""" Cannot add node {node.key}: parent key {node.parent_key} not found.
|
||||||
|
Please ensure parent nodes are added before child nodes."""
|
||||||
|
raise ValueError(error)
|
||||||
|
parent.add_child(node)
|
||||||
|
self.nodes[node.key] = node
|
||||||
|
|
||||||
|
def get_node(self, key: str) -> TreeNode:
|
||||||
|
return self.nodes.get(key)
|
||||||
|
|
||||||
|
def remove_node(self, node: TreeNode):
|
||||||
|
if node.is_root:
|
||||||
|
self.root = None
|
||||||
|
else:
|
||||||
|
parent: TreeNode = node.parent
|
||||||
|
parent.remove_child(node)
|
||||||
|
self.nodes.pop(node.key, None)
|
||||||
|
|
||||||
|
def get_nodes(self):
|
||||||
|
return list(self.nodes.values())
|
||||||
|
|
||||||
|
def get_node_children(self, key, with_self=False):
|
||||||
|
node = self.get_node(key)
|
||||||
|
if not node:
|
||||||
|
return []
|
||||||
|
nodes = []
|
||||||
|
if with_self:
|
||||||
|
nodes.append(node)
|
||||||
|
nodes.extend(node.children)
|
||||||
|
return nodes
|
||||||
|
|
||||||
|
def print(self, count=10, simple=True):
|
||||||
|
print('tree_root_key: ', getattr(self.root, 'key', 'No-root'))
|
||||||
|
print('tree_size: ', self.size)
|
||||||
|
print('tree_depth: ', self.depth)
|
||||||
|
print('tree_width: ', self.width)
|
||||||
|
|
||||||
|
is_print_key = True
|
||||||
|
for n in list(self.nodes.values())[:count]:
|
||||||
|
n: TreeNode
|
||||||
|
n.print(simple=simple, is_print_keys=is_print_key)
|
||||||
|
is_print_key = False
|
||||||
@@ -12,7 +12,7 @@ from urllib.parse import urlparse
|
|||||||
|
|
||||||
from django.conf import settings
|
from django.conf import settings
|
||||||
from django.core.exceptions import SuspiciousOperation
|
from django.core.exceptions import SuspiciousOperation
|
||||||
from django.utils.encoding import force_bytes
|
from django.utils.encoding import force_bytes, smart_bytes
|
||||||
from jwkest import JWKESTException
|
from jwkest import JWKESTException
|
||||||
from jwkest.jwk import KEYS
|
from jwkest.jwk import KEYS
|
||||||
from jwkest.jws import JWS
|
from jwkest.jws import JWS
|
||||||
@@ -58,7 +58,7 @@ def _get_jwks_keys(shared_key):
|
|||||||
# Adds the shared key (which can correspond to the client_secret) as an oct key so it can be
|
# Adds the shared key (which can correspond to the client_secret) as an oct key so it can be
|
||||||
# used for HMAC signatures.
|
# used for HMAC signatures.
|
||||||
logger.debug(log_prompt.format('Add key'))
|
logger.debug(log_prompt.format('Add key'))
|
||||||
jwks_keys.add({'key': force_bytes(shared_key), 'kty': 'oct'})
|
jwks_keys.add({'key': smart_bytes(shared_key), 'kty': 'oct'})
|
||||||
logger.debug(log_prompt.format('End'))
|
logger.debug(log_prompt.format('End'))
|
||||||
return jwks_keys
|
return jwks_keys
|
||||||
|
|
||||||
|
|||||||
@@ -28,7 +28,7 @@ is_available:
|
|||||||
sample: true
|
sample: true
|
||||||
'''
|
'''
|
||||||
|
|
||||||
import telnetlib3
|
import telnetlib
|
||||||
|
|
||||||
from ansible.module_utils.basic import AnsibleModule
|
from ansible.module_utils.basic import AnsibleModule
|
||||||
|
|
||||||
@@ -57,9 +57,9 @@ def main():
|
|||||||
port = module.params['login_port']
|
port = module.params['login_port']
|
||||||
timeout = module.params['timeout']
|
timeout = module.params['timeout']
|
||||||
try:
|
try:
|
||||||
client = telnetlib3.Telnet(host, port, timeout=timeout)
|
client = telnetlib.Telnet(host, port, timeout=timeout)
|
||||||
client.close()
|
client.close()
|
||||||
except Exception as err: # noqa
|
except Exception as err: # noqa
|
||||||
result['is_available'] = False
|
result['is_available'] = False
|
||||||
module.fail_json(msg='Unable to connect to asset: %s' % err)
|
module.fail_json(msg='Unable to connect to asset: %s' % err)
|
||||||
|
|
||||||
|
|||||||
73
apps/perms/tree.py
Normal file
73
apps/perms/tree.py
Normal file
@@ -0,0 +1,73 @@
|
|||||||
|
from collections import defaultdict
|
||||||
|
from django.db.models import Q, Count
|
||||||
|
|
||||||
|
from common.utils import get_logger
|
||||||
|
from users.models import User
|
||||||
|
from assets.tree.asset_tree import AssetTree, AssetTreeNode
|
||||||
|
from perms.utils.utils import UserPermUtil
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ['UserPermTree']
|
||||||
|
|
||||||
|
|
||||||
|
logger = get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class PermTreeNode(AssetTreeNode):
|
||||||
|
|
||||||
|
class Type:
|
||||||
|
# Neither a permission node nor a node with direct permission assets
|
||||||
|
BRIDGE = 'bridge'
|
||||||
|
# Node with direct permission
|
||||||
|
DN = 'dn'
|
||||||
|
# Node with only direct permission assets
|
||||||
|
DA = 'da'
|
||||||
|
|
||||||
|
def __init__(self, tp, _id, key, value, assets_count=0, assets=None):
|
||||||
|
super().__init__(_id, key, value, assets_count)
|
||||||
|
self.type = tp or self.Type.BRIDGE
|
||||||
|
|
||||||
|
def as_dict(self, simple=True):
|
||||||
|
data = super().as_dict(simple=simple)
|
||||||
|
data.update({
|
||||||
|
'type': self.type,
|
||||||
|
})
|
||||||
|
return data
|
||||||
|
|
||||||
|
|
||||||
|
class UserPermTree(AssetTree):
|
||||||
|
|
||||||
|
TreeNode = PermTreeNode
|
||||||
|
|
||||||
|
def __init__(self, user=None, assets_q_object=None, category=None, org=None, with_assets=False):
|
||||||
|
super().__init__(
|
||||||
|
assets_q_object=assets_q_object,
|
||||||
|
category=category,
|
||||||
|
org=org,
|
||||||
|
with_assets=with_assets,
|
||||||
|
full_tree=False
|
||||||
|
)
|
||||||
|
self._user: User = user
|
||||||
|
self._util = UserPermUtil(user, org=self._org)
|
||||||
|
|
||||||
|
def _make_assets_q_object(self):
|
||||||
|
q = super()._make_assets_q_object()
|
||||||
|
q_perm_assets = Q(id__in=self._util._user_direct_asset_ids)
|
||||||
|
q_perm_nodes = Q(node_id__in=self._util._user_direct_node_all_children_ids)
|
||||||
|
q = q & (q_perm_assets | q_perm_nodes)
|
||||||
|
return q
|
||||||
|
|
||||||
|
def _get_tree_node_data(self, node_id):
|
||||||
|
data = super()._get_tree_node_data(node_id)
|
||||||
|
if node_id in self._util._user_direct_node_all_children_ids:
|
||||||
|
tp = PermTreeNode.Type.DN
|
||||||
|
elif self._nodes_assets_count_mapper.get(node_id, 0) > 0:
|
||||||
|
tp = PermTreeNode.Type.DA
|
||||||
|
else:
|
||||||
|
tp = PermTreeNode.Type.BRIDGE
|
||||||
|
data.update({ 'tp': tp })
|
||||||
|
return data
|
||||||
|
|
||||||
|
def print(self, simple=True, count=10):
|
||||||
|
self._util.print()
|
||||||
|
super().print(simple=simple, count=count)
|
||||||
138
apps/perms/utils/utils.py
Normal file
138
apps/perms/utils/utils.py
Normal file
@@ -0,0 +1,138 @@
|
|||||||
|
|
||||||
|
from django.db.models import Q
|
||||||
|
|
||||||
|
from common.utils import timeit, lazyproperty, get_logger, is_uuid
|
||||||
|
from orgs.utils import current_org
|
||||||
|
from users.models import User
|
||||||
|
from assets.models import Node, Asset
|
||||||
|
from perms.models import AssetPermission
|
||||||
|
|
||||||
|
|
||||||
|
logger = get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ['UserPermUtil']
|
||||||
|
|
||||||
|
|
||||||
|
class UserPermUtil(object):
|
||||||
|
|
||||||
|
UserGroupThrough = User.groups.through
|
||||||
|
PermUserThrough = AssetPermission.users.through
|
||||||
|
PermUserGroupThrough = AssetPermission.user_groups.through
|
||||||
|
PermAssetThrough = AssetPermission.assets.through
|
||||||
|
PermNodeThrough = AssetPermission.nodes.through
|
||||||
|
|
||||||
|
def __init__(self, user, org=None):
|
||||||
|
self._user: User = user
|
||||||
|
self._org = org or current_org
|
||||||
|
self._user_permission_ids = set()
|
||||||
|
self._user_group_ids = set()
|
||||||
|
self._user_group_permission_ids = set()
|
||||||
|
self._user_all_permission_ids = set()
|
||||||
|
self._user_direct_asset_ids = set()
|
||||||
|
self._user_direct_node_ids = set()
|
||||||
|
self._user_direct_node_all_children_ids = set()
|
||||||
|
self._init()
|
||||||
|
|
||||||
|
def _init(self):
|
||||||
|
self._load_user_permission_ids()
|
||||||
|
self._load_user_group_ids()
|
||||||
|
self._load_user_group_permission_ids()
|
||||||
|
self._load_user_direct_asset_ids()
|
||||||
|
self._load_user_direct_node_ids()
|
||||||
|
self._load_user_direct_node_all_children_ids()
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _load_user_permission_ids(self):
|
||||||
|
perm_ids = self.PermUserThrough.objects.filter(
|
||||||
|
user_id=self._user.id
|
||||||
|
).distinct('assetpermission_id').values_list('assetpermission_id', flat=True)
|
||||||
|
perm_ids = self._uuids_to_string(perm_ids)
|
||||||
|
self._user_permission_ids.update(perm_ids)
|
||||||
|
self._user_all_permission_ids.update(perm_ids)
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _load_user_group_ids(self):
|
||||||
|
group_ids = self.UserGroupThrough.objects.filter(
|
||||||
|
user_id=self._user.id
|
||||||
|
).distinct('usergroup_id').values_list('usergroup_id', flat=True)
|
||||||
|
group_ids = self._uuids_to_string(group_ids)
|
||||||
|
self._user_group_ids.update(group_ids)
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _load_user_group_permission_ids(self):
|
||||||
|
perm_ids = self.PermUserGroupThrough.objects.filter(
|
||||||
|
usergroup_id__in=self._user_group_ids
|
||||||
|
).distinct('assetpermission_id').values_list('assetpermission_id', flat=True)
|
||||||
|
perm_ids = self._uuids_to_string(perm_ids)
|
||||||
|
self._user_group_permission_ids.update(perm_ids)
|
||||||
|
self._user_all_permission_ids.update(perm_ids)
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _load_user_direct_asset_ids(self):
|
||||||
|
asset_ids = self.PermAssetThrough.objects.filter(
|
||||||
|
assetpermission_id__in=self._user_all_permission_ids
|
||||||
|
).distinct('asset_id').values_list('asset_id', flat=True)
|
||||||
|
asset_ids = self._uuids_to_string(asset_ids)
|
||||||
|
self._user_direct_asset_ids.update(asset_ids)
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _load_user_direct_node_ids(self):
|
||||||
|
node_ids = self.PermNodeThrough.objects.filter(
|
||||||
|
assetpermission_id__in=self._user_all_permission_ids
|
||||||
|
).distinct('node_id').values_list('node_id', flat=True)
|
||||||
|
node_ids = self._uuids_to_string(node_ids)
|
||||||
|
self._user_direct_node_ids.update(node_ids)
|
||||||
|
|
||||||
|
@timeit
|
||||||
|
def _load_user_direct_node_all_children_ids(self):
|
||||||
|
nid_key_pairs = Node.objects.filter(org_id=self._org.id).values_list('id', 'key')
|
||||||
|
nid_key_mapper = { str(nid): key for nid, key in nid_key_pairs }
|
||||||
|
|
||||||
|
dn_keys = [ nid_key_mapper[nid] for nid in self._user_direct_node_ids ]
|
||||||
|
|
||||||
|
def has_ancestor_in_direct_nodes(key: str) -> bool:
|
||||||
|
ancestor_keys = [ ':'.join(key.split(':')[:i]) for i in range(1, key.count(':') + 1) ]
|
||||||
|
return bool(set(ancestor_keys) & set(dn_keys))
|
||||||
|
|
||||||
|
dn_children_ids = [ nid for nid, key in nid_key_mapper.items() if has_ancestor_in_direct_nodes(key) ]
|
||||||
|
|
||||||
|
self._user_direct_node_all_children_ids.update(self._user_direct_node_ids)
|
||||||
|
self._user_direct_node_all_children_ids.update(dn_children_ids)
|
||||||
|
|
||||||
|
def get_node_assets(self, node: Node):
|
||||||
|
''' 获取节点下授权的直接资产, Luna 页面展开时需要 '''
|
||||||
|
q = Q(node_id=node.id)
|
||||||
|
if str(node.id) not in self._user_direct_node_all_children_ids:
|
||||||
|
q &= Q(id__in=self._user_direct_asset_ids)
|
||||||
|
assets = Asset.objects.filter(q)
|
||||||
|
return assets
|
||||||
|
|
||||||
|
def get_node_all_assets(self, node: Node):
|
||||||
|
''' 获取节点及其子节点下所有授权资产, 测试时需要 '''
|
||||||
|
if str(node.id) in self._user_direct_node_all_children_ids:
|
||||||
|
assets = node.get_all_assets()
|
||||||
|
return assets
|
||||||
|
|
||||||
|
children_ids = node.get_all_children(with_self=True).values_list('id', flat=True)
|
||||||
|
children_ids = self._uuids_to_string(children_ids)
|
||||||
|
dn_all_nodes_ids = set(children_ids) & self._user_direct_node_all_children_ids
|
||||||
|
other_nodes_ids = set(children_ids) - dn_all_nodes_ids
|
||||||
|
|
||||||
|
q = Q(node_id__in=dn_all_nodes_ids)
|
||||||
|
q |= Q(node_id__in=other_nodes_ids) & Q(id__in=self._user_direct_asset_ids)
|
||||||
|
assets = Asset.objects.filter(q)
|
||||||
|
return assets
|
||||||
|
|
||||||
|
def _uuids_to_string(self, uuids):
|
||||||
|
return [ str(u) for u in uuids ]
|
||||||
|
|
||||||
|
def print(self):
|
||||||
|
print('user_perm_tree:', self._user.username)
|
||||||
|
print('user_permission_ids_count:', len(self._user_permission_ids))
|
||||||
|
print('user_group_ids_count:', len(self._user_group_ids))
|
||||||
|
print('user_group_permission_ids_count:', len(self._user_permission_ids) - len(self._user_group_ids))
|
||||||
|
print('user_all_permission_ids_count:', len(self._user_all_permission_ids))
|
||||||
|
print('user_direct_asset_ids_count:', len(self._user_direct_asset_ids))
|
||||||
|
print('user_direct_node_ids_count:', len(self._user_direct_node_ids))
|
||||||
|
print('user_direct_node_all_children_ids_count:', len(self._user_direct_node_all_children_ids))
|
||||||
@@ -2,8 +2,7 @@
|
|||||||
#
|
#
|
||||||
import asyncio
|
import asyncio
|
||||||
import socket
|
import socket
|
||||||
|
import telnetlib
|
||||||
import telnetlib3
|
|
||||||
|
|
||||||
from settings.utils import generate_ips
|
from settings.utils import generate_ips
|
||||||
|
|
||||||
@@ -13,7 +12,7 @@ PROMPT_REGEX = r'[\<|\[](.*)[\>|\]]'
|
|||||||
async def telnet(dest_addr, port_number=23, timeout=10):
|
async def telnet(dest_addr, port_number=23, timeout=10):
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
try:
|
try:
|
||||||
connection = await loop.run_in_executor(None, telnetlib3.Telnet, dest_addr, port_number, timeout)
|
connection = await loop.run_in_executor(None, telnetlib.Telnet, dest_addr, port_number, timeout)
|
||||||
except asyncio.TimeoutError:
|
except asyncio.TimeoutError:
|
||||||
return False, 'Timeout'
|
return False, 'Timeout'
|
||||||
except (ConnectionRefusedError, socket.timeout, socket.gaierror) as e:
|
except (ConnectionRefusedError, socket.timeout, socket.gaierror) as e:
|
||||||
|
|||||||
@@ -40,11 +40,11 @@ dependencies = [
|
|||||||
'pyyaml==6.0.1',
|
'pyyaml==6.0.1',
|
||||||
'requests==2.32.4',
|
'requests==2.32.4',
|
||||||
'simplejson==3.19.1',
|
'simplejson==3.19.1',
|
||||||
'six==1.17.0',
|
'six==1.16.0',
|
||||||
'sshtunnel==0.4.0',
|
'sshtunnel==0.4.0',
|
||||||
'sshpubkeys==3.3.1',
|
'sshpubkeys==3.3.1',
|
||||||
'uritemplate==4.1.1',
|
'uritemplate==4.1.1',
|
||||||
'vine==5.1.0',
|
'vine==5.0.0',
|
||||||
'werkzeug==3.0.6',
|
'werkzeug==3.0.6',
|
||||||
'unicodecsv==0.14.1',
|
'unicodecsv==0.14.1',
|
||||||
'httpsig==1.3.0',
|
'httpsig==1.3.0',
|
||||||
@@ -68,32 +68,32 @@ dependencies = [
|
|||||||
'ipip-ipdb==1.6.1',
|
'ipip-ipdb==1.6.1',
|
||||||
'pywinrm==0.4.3',
|
'pywinrm==0.4.3',
|
||||||
'python-nmap==0.7.1',
|
'python-nmap==0.7.1',
|
||||||
'django==5.2.9',
|
'django==4.1.13',
|
||||||
'django-bootstrap3==23.4',
|
'django-bootstrap3==23.4',
|
||||||
'django-filter==24.3',
|
'django-filter==23.2',
|
||||||
'django-formtools==2.5.1',
|
'django-formtools==2.5.1',
|
||||||
'django-ranged-response==0.2.0',
|
'django-ranged-response==0.2.0',
|
||||||
'django-simple-captcha==0.5.18',
|
'django-simple-captcha==0.5.18',
|
||||||
'django-timezone-field==7.1',
|
'django-timezone-field==5.1',
|
||||||
'djangorestframework==3.15.0',
|
'djangorestframework==3.14.0',
|
||||||
'djangorestframework-bulk==0.2.1',
|
'djangorestframework-bulk==0.2.1',
|
||||||
'django-simple-history==3.6.0',
|
'django-simple-history==3.6.0',
|
||||||
'django-private-storage==3.1.3',
|
'django-private-storage==3.1',
|
||||||
'drf-nested-routers==0.93.4',
|
'drf-nested-routers==0.93.4',
|
||||||
'drf-writable-nested==0.7.0',
|
'drf-writable-nested==0.7.0',
|
||||||
'rest-condition==1.0.3',
|
'rest-condition==1.0.3',
|
||||||
'drf-spectacular==0.29.0',
|
'drf-spectacular==0.28.0',
|
||||||
'pillow==12.0.0',
|
'pillow==10.2.0',
|
||||||
'pytz==2025.2',
|
'pytz==2025.2',
|
||||||
'django-proxy==1.2.2',
|
'django-proxy==1.2.2',
|
||||||
'python-daemon==3.0.1',
|
'python-daemon==3.0.1',
|
||||||
'eventlet==0.40.3',
|
'eventlet==0.40.3',
|
||||||
'greenlet==3.2.4',
|
'greenlet==3.2.4',
|
||||||
'gunicorn==23.0.0',
|
'gunicorn==23.0.0',
|
||||||
'celery==5.6.0',
|
'celery==5.3.1',
|
||||||
'flower==2.0.1',
|
'flower==2.0.1',
|
||||||
'django-celery-beat==2.8.1',
|
'django-celery-beat==2.6.0',
|
||||||
'kombu==5.6.0',
|
'kombu==5.3.5',
|
||||||
'uvicorn==0.22.0',
|
'uvicorn==0.22.0',
|
||||||
'websockets==11.0.3',
|
'websockets==11.0.3',
|
||||||
'python-ldap==3.4.5',
|
'python-ldap==3.4.5',
|
||||||
@@ -103,10 +103,10 @@ dependencies = [
|
|||||||
'python-cas==1.6.0',
|
'python-cas==1.6.0',
|
||||||
'django-auth-ldap==4.4.0',
|
'django-auth-ldap==4.4.0',
|
||||||
'mysqlclient==2.2.4',
|
'mysqlclient==2.2.4',
|
||||||
'pymssql==2.3.10',
|
'pymssql==2.3.4',
|
||||||
'django-redis==5.4.0',
|
'django-redis==5.3.0',
|
||||||
'python-redis-lock==4.0.0',
|
'python-redis-lock==4.0.0',
|
||||||
'pyopenssl==24.0.0',
|
'pyopenssl==23.2.0',
|
||||||
'redis',
|
'redis',
|
||||||
'pymongo==4.6.3',
|
'pymongo==4.6.3',
|
||||||
'pyfreerdp==0.0.2',
|
'pyfreerdp==0.0.2',
|
||||||
@@ -153,7 +153,6 @@ dependencies = [
|
|||||||
'pdf2image==1.17.0',
|
'pdf2image==1.17.0',
|
||||||
'drf-spectacular-sidecar==2025.8.1',
|
'drf-spectacular-sidecar==2025.8.1',
|
||||||
"django-oauth-toolkit==2.4.0",
|
"django-oauth-toolkit==2.4.0",
|
||||||
"telnetlib3==2.0.8",
|
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.urls]
|
[project.urls]
|
||||||
|
|||||||
@@ -0,0 +1,358 @@
|
|||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import django
|
||||||
|
import random
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
if os.path.exists('../../apps'):
|
||||||
|
sys.path.insert(0, '../../apps')
|
||||||
|
if os.path.exists('../apps'):
|
||||||
|
sys.path.insert(0, '../apps')
|
||||||
|
elif os.path.exists('./apps'):
|
||||||
|
sys.path.insert(0, './apps')
|
||||||
|
|
||||||
|
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "jumpserver.settings")
|
||||||
|
django.setup()
|
||||||
|
|
||||||
|
|
||||||
|
from assets.models import Asset, Node
|
||||||
|
from orgs.models import Organization
|
||||||
|
from django.db.models import Count
|
||||||
|
|
||||||
|
OUTPUT_FILE = 'report_cleanup_and_keep_one_node_for_multi_parent_nodes_assets.txt'
|
||||||
|
|
||||||
|
# Special organization IDs and names
|
||||||
|
SPECIAL_ORGS = {
|
||||||
|
'00000000-0000-0000-0000-000000000000': 'GLOBAL',
|
||||||
|
'00000000-0000-0000-0000-000000000002': 'DEFAULT',
|
||||||
|
'00000000-0000-0000-0000-000000000004': 'SYSTEM',
|
||||||
|
}
|
||||||
|
|
||||||
|
try:
|
||||||
|
AssetNodeThrough = Asset.nodes.through
|
||||||
|
except Exception as e:
|
||||||
|
print("Failed to get AssetNodeThrough model. Check Asset.nodes field definition.")
|
||||||
|
raise e
|
||||||
|
|
||||||
|
|
||||||
|
def log(msg=''):
|
||||||
|
"""Print log with timestamp to console"""
|
||||||
|
print(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}")
|
||||||
|
|
||||||
|
|
||||||
|
def write_report(content):
|
||||||
|
"""Write content to report file"""
|
||||||
|
with open(OUTPUT_FILE, 'a', encoding='utf-8') as f:
|
||||||
|
f.write(content)
|
||||||
|
|
||||||
|
|
||||||
|
def get_org_name(org_id, orgs_map):
|
||||||
|
"""Get organization name, check special orgs first, then orgs_map"""
|
||||||
|
# Check if it's a special organization
|
||||||
|
org_id_str = str(org_id)
|
||||||
|
if org_id_str in SPECIAL_ORGS:
|
||||||
|
return SPECIAL_ORGS[org_id_str]
|
||||||
|
|
||||||
|
# Try to get from orgs_map
|
||||||
|
org = orgs_map.get(org_id)
|
||||||
|
if org:
|
||||||
|
return org.name
|
||||||
|
|
||||||
|
return 'Unknown'
|
||||||
|
|
||||||
|
|
||||||
|
def find_and_cleanup_multi_parent_assets():
|
||||||
|
"""Find and cleanup assets with multiple parent nodes"""
|
||||||
|
|
||||||
|
log("Searching for assets with multiple parent nodes...")
|
||||||
|
|
||||||
|
# Find all asset_ids that belong to multiple node_ids
|
||||||
|
multi_parent_assets = AssetNodeThrough.objects.values('asset_id').annotate(
|
||||||
|
node_count=Count('node_id', distinct=True)
|
||||||
|
).filter(node_count__gt=1).order_by('-node_count')
|
||||||
|
|
||||||
|
total_count = multi_parent_assets.count()
|
||||||
|
log(f"Found {total_count:,} assets with multiple parent nodes\n")
|
||||||
|
|
||||||
|
if total_count == 0:
|
||||||
|
log("✓ All assets already have single parent node")
|
||||||
|
return {}
|
||||||
|
|
||||||
|
# Collect all asset_ids and node_ids
|
||||||
|
asset_ids = [item['asset_id'] for item in multi_parent_assets]
|
||||||
|
|
||||||
|
# Get all through records
|
||||||
|
all_through_records = AssetNodeThrough.objects.filter(asset_id__in=asset_ids)
|
||||||
|
node_ids = list(set(through.node_id for through in all_through_records))
|
||||||
|
|
||||||
|
# Batch fetch all objects
|
||||||
|
log("Batch loading Asset objects...")
|
||||||
|
assets_map = {asset.id: asset for asset in Asset.objects.filter(id__in=asset_ids)}
|
||||||
|
|
||||||
|
log("Batch loading Node objects...")
|
||||||
|
nodes_map = {node.id: node for node in Node.objects.filter(id__in=node_ids)}
|
||||||
|
|
||||||
|
# Batch fetch all Organization objects
|
||||||
|
org_ids = list(set(asset.org_id for asset in assets_map.values())) + \
|
||||||
|
list(set(node.org_id for node in nodes_map.values()))
|
||||||
|
org_ids = list(set(org_ids))
|
||||||
|
|
||||||
|
log("Batch loading Organization objects...")
|
||||||
|
orgs_map = {org.id: org for org in Organization.objects.filter(id__in=org_ids)}
|
||||||
|
|
||||||
|
# Build mapping of asset_id -> list of through_records
|
||||||
|
asset_nodes_map = {}
|
||||||
|
for through in all_through_records:
|
||||||
|
if through.asset_id not in asset_nodes_map:
|
||||||
|
asset_nodes_map[through.asset_id] = []
|
||||||
|
asset_nodes_map[through.asset_id].append(through)
|
||||||
|
|
||||||
|
# Organize by organization
|
||||||
|
org_cleanup_data = {} # org_id -> { asset_id -> { keep_node_id, remove_node_ids } }
|
||||||
|
|
||||||
|
for item in multi_parent_assets:
|
||||||
|
asset_id = item['asset_id']
|
||||||
|
|
||||||
|
# Get Asset object
|
||||||
|
asset = assets_map.get(asset_id)
|
||||||
|
if not asset:
|
||||||
|
log(f"⚠ Asset {asset_id} not found in map, skipping")
|
||||||
|
continue
|
||||||
|
|
||||||
|
org_id = asset.org_id
|
||||||
|
|
||||||
|
# Initialize org data if not exists
|
||||||
|
if org_id not in org_cleanup_data:
|
||||||
|
org_cleanup_data[org_id] = {}
|
||||||
|
|
||||||
|
# Get all nodes for this asset
|
||||||
|
through_records = asset_nodes_map.get(asset_id, [])
|
||||||
|
|
||||||
|
if len(through_records) < 2:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Randomly select one node to keep
|
||||||
|
keep_through = random.choice(through_records)
|
||||||
|
remove_throughs = [t for t in through_records if t.id != keep_through.id]
|
||||||
|
|
||||||
|
org_cleanup_data[org_id][asset_id] = {
|
||||||
|
'asset_name': asset.name,
|
||||||
|
'keep_node_id': keep_through.node_id,
|
||||||
|
'keep_node': nodes_map.get(keep_through.node_id),
|
||||||
|
'remove_records': remove_throughs,
|
||||||
|
'remove_nodes': [nodes_map.get(t.node_id) for t in remove_throughs]
|
||||||
|
}
|
||||||
|
|
||||||
|
return org_cleanup_data
|
||||||
|
|
||||||
|
|
||||||
|
def perform_cleanup(org_cleanup_data, dry_run=False):
|
||||||
|
"""Perform the actual cleanup - delete extra node relationships"""
|
||||||
|
|
||||||
|
if dry_run:
|
||||||
|
log("DRY RUN: Simulating cleanup process (no data will be deleted)...")
|
||||||
|
else:
|
||||||
|
log("\nStarting cleanup process...")
|
||||||
|
|
||||||
|
total_deleted = 0
|
||||||
|
|
||||||
|
for org_id in org_cleanup_data.keys():
|
||||||
|
for asset_id, cleanup_info in org_cleanup_data[org_id].items():
|
||||||
|
# Delete the extra relationships
|
||||||
|
for through_record in cleanup_info['remove_records']:
|
||||||
|
if not dry_run:
|
||||||
|
through_record.delete()
|
||||||
|
total_deleted += 1
|
||||||
|
|
||||||
|
return total_deleted
|
||||||
|
|
||||||
|
|
||||||
|
def verify_cleanup():
|
||||||
|
"""Verify that there are no more assets with multiple parent nodes"""
|
||||||
|
log("\n" + "="*80)
|
||||||
|
log("VERIFICATION: Checking for remaining assets with multiple parent nodes...")
|
||||||
|
log("="*80)
|
||||||
|
|
||||||
|
# Find all asset_ids that belong to multiple node_ids
|
||||||
|
multi_parent_assets = AssetNodeThrough.objects.values('asset_id').annotate(
|
||||||
|
node_count=Count('node_id', distinct=True)
|
||||||
|
).filter(node_count__gt=1).order_by('-node_count')
|
||||||
|
|
||||||
|
remaining_count = multi_parent_assets.count()
|
||||||
|
|
||||||
|
if remaining_count == 0:
|
||||||
|
log(f"✓ Verification successful: No assets with multiple parent nodes remaining\n")
|
||||||
|
return True
|
||||||
|
else:
|
||||||
|
log(f"✗ Verification failed: Found {remaining_count:,} assets still with multiple parent nodes\n")
|
||||||
|
# Show some details
|
||||||
|
for item in multi_parent_assets[:10]:
|
||||||
|
asset_id = item['asset_id']
|
||||||
|
node_count = item['node_count']
|
||||||
|
try:
|
||||||
|
asset = Asset.objects.get(id=asset_id)
|
||||||
|
log(f" - Asset: {asset.name} ({asset_id}) has {node_count} parent nodes")
|
||||||
|
except:
|
||||||
|
log(f" - Asset ID: {asset_id} has {node_count} parent nodes")
|
||||||
|
|
||||||
|
if remaining_count > 10:
|
||||||
|
log(f" ... and {remaining_count - 10} more")
|
||||||
|
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def generate_report(org_cleanup_data, total_deleted):
|
||||||
|
"""Generate and write report to file"""
|
||||||
|
# Clear previous report
|
||||||
|
if os.path.exists(OUTPUT_FILE):
|
||||||
|
os.remove(OUTPUT_FILE)
|
||||||
|
|
||||||
|
# Write header
|
||||||
|
write_report(f"Multi-Parent Assets Cleanup Report\n")
|
||||||
|
write_report(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
|
||||||
|
write_report(f"{'='*80}\n\n")
|
||||||
|
|
||||||
|
# Get all organizations
|
||||||
|
all_org_ids = list(set(org_id for org_id in org_cleanup_data.keys()))
|
||||||
|
all_orgs = {org.id: org for org in Organization.objects.filter(id__in=all_org_ids)}
|
||||||
|
|
||||||
|
# Calculate statistics
|
||||||
|
total_orgs = Organization.objects.count()
|
||||||
|
orgs_processed = len(org_cleanup_data)
|
||||||
|
orgs_no_issues = total_orgs - orgs_processed
|
||||||
|
total_assets_cleaned = sum(len(assets) for assets in org_cleanup_data.values())
|
||||||
|
|
||||||
|
# Overview
|
||||||
|
write_report("OVERVIEW\n")
|
||||||
|
write_report(f"{'-'*80}\n")
|
||||||
|
write_report(f"Total organizations: {total_orgs:,}\n")
|
||||||
|
write_report(f"Organizations processed: {orgs_processed:,}\n")
|
||||||
|
write_report(f"Organizations without issues: {orgs_no_issues:,}\n")
|
||||||
|
write_report(f"Total assets cleaned: {total_assets_cleaned:,}\n")
|
||||||
|
total_relationships = AssetNodeThrough.objects.count()
|
||||||
|
write_report(f"Total relationships (through records): {total_relationships:,}\n")
|
||||||
|
write_report(f"Total relationships deleted: {total_deleted:,}\n\n")
|
||||||
|
|
||||||
|
# Summary by organization
|
||||||
|
write_report("Summary by Organization:\n")
|
||||||
|
for org_id in sorted(org_cleanup_data.keys()):
|
||||||
|
org_name = get_org_name(org_id, all_orgs)
|
||||||
|
asset_count = len(org_cleanup_data[org_id])
|
||||||
|
write_report(f" - {org_name} ({org_id}): {asset_count:,} assets cleaned\n")
|
||||||
|
|
||||||
|
write_report(f"\n{'='*80}\n\n")
|
||||||
|
|
||||||
|
# Detailed cleanup information grouped by organization
|
||||||
|
for org_id in sorted(org_cleanup_data.keys()):
|
||||||
|
org_name = get_org_name(org_id, all_orgs)
|
||||||
|
asset_count = len(org_cleanup_data[org_id])
|
||||||
|
|
||||||
|
write_report(f"ORGANIZATION: {org_name} ({org_id})\n")
|
||||||
|
write_report(f"Total assets cleaned: {asset_count:,}\n")
|
||||||
|
write_report(f"{'-'*80}\n\n")
|
||||||
|
|
||||||
|
for asset_id, cleanup_info in org_cleanup_data[org_id].items():
|
||||||
|
write_report(f"Asset: {cleanup_info['asset_name']} ({asset_id})\n")
|
||||||
|
|
||||||
|
# Kept node
|
||||||
|
keep_node = cleanup_info['keep_node']
|
||||||
|
if keep_node:
|
||||||
|
write_report(f" ✓ Kept: {keep_node.name} (key: {keep_node.key}) (id: {keep_node.id})\n")
|
||||||
|
else:
|
||||||
|
write_report(f" ✓ Kept: Unknown (id: {cleanup_info['keep_node_id']})\n")
|
||||||
|
|
||||||
|
# Removed nodes
|
||||||
|
write_report(f" ✗ Removed: {len(cleanup_info['remove_nodes'])} node(s)\n")
|
||||||
|
for node in cleanup_info['remove_nodes']:
|
||||||
|
if node:
|
||||||
|
write_report(f" - {node.name} (key: {node.key}) (id: {node.id})\n")
|
||||||
|
else:
|
||||||
|
write_report(f" - Unknown\n")
|
||||||
|
|
||||||
|
write_report(f"\n")
|
||||||
|
|
||||||
|
write_report(f"{'='*80}\n\n")
|
||||||
|
|
||||||
|
log(f"✓ Report written to {OUTPUT_FILE}")
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
try:
|
||||||
|
# Display warning banner
|
||||||
|
warning_message = """
|
||||||
|
╔══════════════════════════════════════════════════════════════════════════════╗
|
||||||
|
║ ⚠️ WARNING ⚠️ ║
|
||||||
|
║ ║
|
||||||
|
║ This script is designed for TEST/FAKE DATA ONLY! ║
|
||||||
|
║ DO NOT run this script in PRODUCTION environment! ║
|
||||||
|
║ ║
|
||||||
|
║ This script will DELETE asset-node relationships from the database. ║
|
||||||
|
║ Use only for data cleanup in development/testing environments. ║
|
||||||
|
║ ║
|
||||||
|
╚══════════════════════════════════════════════════════════════════════════════╝
|
||||||
|
"""
|
||||||
|
print(warning_message)
|
||||||
|
|
||||||
|
# Ask user to confirm before proceeding
|
||||||
|
confirm = input("Do you understand the warning and want to continue? (yes/no): ").strip().lower()
|
||||||
|
if confirm not in ['yes', 'y']:
|
||||||
|
log("✗ Operation cancelled by user")
|
||||||
|
sys.exit(0)
|
||||||
|
|
||||||
|
log("✓ Proceeding with operation\n")
|
||||||
|
|
||||||
|
org_cleanup_data = find_and_cleanup_multi_parent_assets()
|
||||||
|
|
||||||
|
if not org_cleanup_data:
|
||||||
|
log("✓ Cleanup complete, no assets to process")
|
||||||
|
sys.exit(0)
|
||||||
|
|
||||||
|
total_assets = sum(len(assets) for assets in org_cleanup_data.values())
|
||||||
|
log(f"\nProcessing {total_assets:,} assets across {len(org_cleanup_data):,} organizations...")
|
||||||
|
|
||||||
|
# First, do a dry-run to show what will be deleted
|
||||||
|
log("\n" + "="*80)
|
||||||
|
log("PREVIEW: Simulating cleanup process...")
|
||||||
|
log("="*80)
|
||||||
|
total_deleted_preview = perform_cleanup(org_cleanup_data, dry_run=True)
|
||||||
|
log(f"✓ Dry-run complete: {total_deleted_preview:,} relationships would be deleted\n")
|
||||||
|
|
||||||
|
# Generate preview report
|
||||||
|
generate_report(org_cleanup_data, total_deleted_preview)
|
||||||
|
log(f"✓ Preview report written to {OUTPUT_FILE}\n")
|
||||||
|
|
||||||
|
# Ask for confirmation 3 times before actual deletion
|
||||||
|
log("="*80)
|
||||||
|
log("FINAL CONFIRMATION: Do you want to proceed with actual cleanup?")
|
||||||
|
log("="*80)
|
||||||
|
confirmation_count = 3
|
||||||
|
for attempt in range(1, confirmation_count + 1):
|
||||||
|
response = input(f"Confirm cleanup (attempt {attempt}/{confirmation_count})? (yes/no): ").strip().lower()
|
||||||
|
if response not in ['yes', 'y']:
|
||||||
|
log(f"✗ Cleanup cancelled by user at attempt {attempt}")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
log("✓ All confirmations received, proceeding with actual cleanup")
|
||||||
|
|
||||||
|
# Perform cleanup
|
||||||
|
total_deleted = perform_cleanup(org_cleanup_data)
|
||||||
|
log(f"✓ Deleted {total_deleted:,} relationships")
|
||||||
|
|
||||||
|
# Generate final report
|
||||||
|
generate_report(org_cleanup_data, total_deleted)
|
||||||
|
|
||||||
|
# Verify cleanup by checking for remaining multi-parent assets
|
||||||
|
verify_cleanup()
|
||||||
|
|
||||||
|
log(f"✓ Cleanup complete: processed {total_assets:,} assets")
|
||||||
|
sys.exit(0)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
log(f"✗ Error occurred: {str(e)}")
|
||||||
|
import traceback
|
||||||
|
traceback.print_exc()
|
||||||
|
sys.exit(2)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,270 @@
|
|||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import django
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
if os.path.exists('../../apps'):
|
||||||
|
sys.path.insert(0, '../../apps')
|
||||||
|
if os.path.exists('../apps'):
|
||||||
|
sys.path.insert(0, '../apps')
|
||||||
|
elif os.path.exists('./apps'):
|
||||||
|
sys.path.insert(0, './apps')
|
||||||
|
|
||||||
|
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "jumpserver.settings")
|
||||||
|
django.setup()
|
||||||
|
|
||||||
|
|
||||||
|
from assets.models import Asset, Node
|
||||||
|
from orgs.models import Organization
|
||||||
|
from django.db.models import Count
|
||||||
|
|
||||||
|
OUTPUT_FILE = 'report_find_multi_parent_nodes_assets.txt'
|
||||||
|
|
||||||
|
# Special organization IDs and names
|
||||||
|
SPECIAL_ORGS = {
|
||||||
|
'00000000-0000-0000-0000-000000000000': 'GLOBAL',
|
||||||
|
'00000000-0000-0000-0000-000000000002': 'DEFAULT',
|
||||||
|
'00000000-0000-0000-0000-000000000004': 'SYSTEM',
|
||||||
|
}
|
||||||
|
|
||||||
|
try:
|
||||||
|
AssetNodeThrough = Asset.nodes.through
|
||||||
|
except Exception as e:
|
||||||
|
print("Failed to get AssetNodeThrough model. Check Asset.nodes field definition.")
|
||||||
|
raise e
|
||||||
|
|
||||||
|
|
||||||
|
def log(msg=''):
|
||||||
|
"""Print log with timestamp"""
|
||||||
|
print(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}")
|
||||||
|
|
||||||
|
|
||||||
|
def get_org_name(org_id, orgs_map):
|
||||||
|
"""Get organization name, check special orgs first, then orgs_map"""
|
||||||
|
# Check if it's a special organization
|
||||||
|
org_id_str = str(org_id)
|
||||||
|
if org_id_str in SPECIAL_ORGS:
|
||||||
|
return SPECIAL_ORGS[org_id_str]
|
||||||
|
|
||||||
|
# Try to get from orgs_map
|
||||||
|
org = orgs_map.get(org_id)
|
||||||
|
if org:
|
||||||
|
return org.name
|
||||||
|
|
||||||
|
return 'Unknown'
|
||||||
|
|
||||||
|
|
||||||
|
def write_report(content):
|
||||||
|
"""Write content to report file"""
|
||||||
|
with open(OUTPUT_FILE, 'a', encoding='utf-8') as f:
|
||||||
|
f.write(content)
|
||||||
|
|
||||||
|
|
||||||
|
def find_assets_multiple_parents():
|
||||||
|
"""Find assets belonging to multiple node_ids organized by organization"""
|
||||||
|
|
||||||
|
log("Searching for assets with multiple parent nodes...")
|
||||||
|
|
||||||
|
# Find all asset_ids that belong to multiple node_ids
|
||||||
|
multi_parent_assets = AssetNodeThrough.objects.values('asset_id').annotate(
|
||||||
|
node_count=Count('node_id', distinct=True)
|
||||||
|
).filter(node_count__gt=1).order_by('-node_count')
|
||||||
|
|
||||||
|
total_count = multi_parent_assets.count()
|
||||||
|
log(f"Found {total_count:,} assets with multiple parent nodes\n")
|
||||||
|
|
||||||
|
if total_count == 0:
|
||||||
|
log("✓ All assets belong to only one node")
|
||||||
|
return {}
|
||||||
|
|
||||||
|
# Collect all asset_ids and node_ids that need to be fetched
|
||||||
|
asset_ids = [item['asset_id'] for item in multi_parent_assets]
|
||||||
|
|
||||||
|
# Get all through records for these assets
|
||||||
|
all_through_records = AssetNodeThrough.objects.filter(asset_id__in=asset_ids)
|
||||||
|
node_ids = list(set(through.node_id for through in all_through_records))
|
||||||
|
|
||||||
|
# Batch fetch all Asset and Node objects
|
||||||
|
log("Batch loading Asset objects...")
|
||||||
|
assets_map = {asset.id: asset for asset in Asset.objects.filter(id__in=asset_ids)}
|
||||||
|
|
||||||
|
log("Batch loading Node objects...")
|
||||||
|
nodes_map = {node.id: node for node in Node.objects.filter(id__in=node_ids)}
|
||||||
|
|
||||||
|
# Batch fetch all Organization objects
|
||||||
|
org_ids = list(set(asset.org_id for asset in assets_map.values())) + \
|
||||||
|
list(set(node.org_id for node in nodes_map.values()))
|
||||||
|
org_ids = list(set(org_ids)) # Remove duplicates
|
||||||
|
|
||||||
|
log("Batch loading Organization objects...")
|
||||||
|
orgs_map = {org.id: org for org in Organization.objects.filter(id__in=org_ids)}
|
||||||
|
|
||||||
|
# Build mapping of asset_id -> list of through_records
|
||||||
|
asset_nodes_map = {}
|
||||||
|
for through in all_through_records:
|
||||||
|
if through.asset_id not in asset_nodes_map:
|
||||||
|
asset_nodes_map[through.asset_id] = []
|
||||||
|
asset_nodes_map[through.asset_id].append(through)
|
||||||
|
|
||||||
|
# Organize by organization first, then by node count, then by asset
|
||||||
|
org_assets_data = {} # org_id -> { node_count -> [asset_data] }
|
||||||
|
|
||||||
|
for item in multi_parent_assets:
|
||||||
|
asset_id = item['asset_id']
|
||||||
|
node_count = item['node_count']
|
||||||
|
|
||||||
|
# Get Asset object from map
|
||||||
|
asset = assets_map.get(asset_id)
|
||||||
|
if not asset:
|
||||||
|
log(f"⚠ Asset {asset_id} not found in map, skipping")
|
||||||
|
continue
|
||||||
|
|
||||||
|
org_id = asset.org_id
|
||||||
|
|
||||||
|
# Initialize org data if not exists
|
||||||
|
if org_id not in org_assets_data:
|
||||||
|
org_assets_data[org_id] = {}
|
||||||
|
|
||||||
|
# Get all nodes for this asset
|
||||||
|
through_records = asset_nodes_map.get(asset_id, [])
|
||||||
|
|
||||||
|
node_details = []
|
||||||
|
for through in through_records:
|
||||||
|
# Get Node object from map
|
||||||
|
node = nodes_map.get(through.node_id)
|
||||||
|
if not node:
|
||||||
|
log(f"⚠ Node {through.node_id} not found in map, skipping")
|
||||||
|
continue
|
||||||
|
|
||||||
|
node_details.append({
|
||||||
|
'id': node.id,
|
||||||
|
'name': node.name,
|
||||||
|
'key': node.key,
|
||||||
|
'path': node.full_value if hasattr(node, 'full_value') else ''
|
||||||
|
})
|
||||||
|
|
||||||
|
if not node_details:
|
||||||
|
continue
|
||||||
|
|
||||||
|
if node_count not in org_assets_data[org_id]:
|
||||||
|
org_assets_data[org_id][node_count] = []
|
||||||
|
|
||||||
|
org_assets_data[org_id][node_count].append({
|
||||||
|
'asset_id': asset.id,
|
||||||
|
'asset_name': asset.name,
|
||||||
|
'nodes': node_details
|
||||||
|
})
|
||||||
|
|
||||||
|
return org_assets_data
|
||||||
|
|
||||||
|
|
||||||
|
def generate_report(org_assets_data):
|
||||||
|
"""Generate and write report to file organized by organization"""
|
||||||
|
# Clear previous report
|
||||||
|
if os.path.exists(OUTPUT_FILE):
|
||||||
|
os.remove(OUTPUT_FILE)
|
||||||
|
|
||||||
|
# Write header
|
||||||
|
write_report(f"Multi-Parent Assets Report\n")
|
||||||
|
write_report(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
|
||||||
|
write_report(f"{'='*80}\n\n")
|
||||||
|
|
||||||
|
# Get all organizations
|
||||||
|
all_org_ids = list(set(org_id for org_id in org_assets_data.keys()))
|
||||||
|
all_orgs = {org.id: org for org in Organization.objects.filter(id__in=all_org_ids)}
|
||||||
|
|
||||||
|
# Calculate statistics
|
||||||
|
total_orgs = Organization.objects.count()
|
||||||
|
orgs_with_issues = len(org_assets_data)
|
||||||
|
orgs_without_issues = total_orgs - orgs_with_issues
|
||||||
|
total_assets_with_issues = sum(
|
||||||
|
len(assets)
|
||||||
|
for org_id in org_assets_data
|
||||||
|
for assets in org_assets_data[org_id].values()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Overview
|
||||||
|
write_report("OVERVIEW\n")
|
||||||
|
write_report(f"{'-'*80}\n")
|
||||||
|
write_report(f"Total organizations: {total_orgs:,}\n")
|
||||||
|
write_report(f"Organizations with multiple-parent assets: {orgs_with_issues:,}\n")
|
||||||
|
write_report(f"Organizations without issues: {orgs_without_issues:,}\n")
|
||||||
|
write_report(f"Total assets with multiple parent nodes: {total_assets_with_issues:,}\n\n")
|
||||||
|
|
||||||
|
# Summary by organization
|
||||||
|
write_report("Summary by Organization:\n")
|
||||||
|
for org_id in sorted(org_assets_data.keys()):
|
||||||
|
org_name = get_org_name(org_id, all_orgs)
|
||||||
|
|
||||||
|
org_asset_count = sum(
|
||||||
|
len(assets)
|
||||||
|
for assets in org_assets_data[org_id].values()
|
||||||
|
)
|
||||||
|
write_report(f" - {org_name} ({org_id}): {org_asset_count:,} assets\n")
|
||||||
|
|
||||||
|
write_report(f"\n{'='*80}\n\n")
|
||||||
|
|
||||||
|
# Detailed sections grouped by organization, then node count
|
||||||
|
for org_id in sorted(org_assets_data.keys()):
|
||||||
|
org_name = get_org_name(org_id, all_orgs)
|
||||||
|
|
||||||
|
org_asset_count = sum(
|
||||||
|
len(assets)
|
||||||
|
for assets in org_assets_data[org_id].values()
|
||||||
|
)
|
||||||
|
|
||||||
|
write_report(f"ORGANIZATION: {org_name} ({org_id})\n")
|
||||||
|
write_report(f"Total assets with issues: {org_asset_count:,}\n")
|
||||||
|
write_report(f"{'-'*80}\n\n")
|
||||||
|
|
||||||
|
# Group by node count within this organization
|
||||||
|
for node_count in sorted(org_assets_data[org_id].keys(), reverse=True):
|
||||||
|
assets = org_assets_data[org_id][node_count]
|
||||||
|
|
||||||
|
write_report(f" Section: {node_count} Parent Nodes ({len(assets):,} assets)\n")
|
||||||
|
write_report(f" {'-'*76}\n\n")
|
||||||
|
|
||||||
|
for asset in assets:
|
||||||
|
write_report(f" {asset['asset_name']} ({asset['asset_id']})\n")
|
||||||
|
|
||||||
|
for node in asset['nodes']:
|
||||||
|
write_report(f" {node['name']} ({node['key']}) ({node['path']}) ({node['id']})\n")
|
||||||
|
|
||||||
|
write_report(f"\n")
|
||||||
|
|
||||||
|
write_report(f"\n")
|
||||||
|
|
||||||
|
write_report(f"{'='*80}\n\n")
|
||||||
|
|
||||||
|
log(f"✓ Report written to {OUTPUT_FILE}")
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
try:
|
||||||
|
org_assets_data = find_assets_multiple_parents()
|
||||||
|
|
||||||
|
if not org_assets_data:
|
||||||
|
log("✓ Detection complete, no issues found")
|
||||||
|
sys.exit(0)
|
||||||
|
|
||||||
|
total_assets = sum(
|
||||||
|
len(assets)
|
||||||
|
for org_id in org_assets_data
|
||||||
|
for assets in org_assets_data[org_id].values()
|
||||||
|
)
|
||||||
|
log(f"Generating report for {total_assets:,} assets across {len(org_assets_data):,} organizations...")
|
||||||
|
|
||||||
|
generate_report(org_assets_data)
|
||||||
|
|
||||||
|
log(f"✗ Detected {total_assets:,} assets with multiple parent nodes")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
log(f"✗ Error occurred: {str(e)}")
|
||||||
|
import traceback
|
||||||
|
traceback.print_exc()
|
||||||
|
sys.exit(2)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -3,7 +3,7 @@ from random import choice
|
|||||||
|
|
||||||
import forgery_py
|
import forgery_py
|
||||||
|
|
||||||
from assets.const import AllTypes
|
from assets.const import AllTypes, Category
|
||||||
from assets.models import *
|
from assets.models import *
|
||||||
from .base import FakeDataGenerator
|
from .base import FakeDataGenerator
|
||||||
|
|
||||||
@@ -48,12 +48,12 @@ class AssetsGenerator(FakeDataGenerator):
|
|||||||
|
|
||||||
def pre_generate(self):
|
def pre_generate(self):
|
||||||
self.node_ids = list(Node.objects.all().values_list('id', flat=True))
|
self.node_ids = list(Node.objects.all().values_list('id', flat=True))
|
||||||
self.platform_ids = list(Platform.objects.filter(category='host').values_list('id', flat=True))
|
self.platform_ids = list(Platform.objects.filter(category=Category.DATABASE).values_list('id', flat=True))
|
||||||
|
|
||||||
def set_assets_nodes(self, assets):
|
def set_assets_nodes(self, assets):
|
||||||
for asset in assets:
|
for asset in assets:
|
||||||
nodes_id_add_to = random.sample(self.node_ids, 3)
|
nodes_id_add_to = random.choice(self.node_ids)
|
||||||
asset.nodes.add(*nodes_id_add_to)
|
asset.node_id = nodes_id_add_to
|
||||||
|
|
||||||
def do_generate(self, batch, batch_size):
|
def do_generate(self, batch, batch_size):
|
||||||
assets = []
|
assets = []
|
||||||
|
|||||||
Reference in New Issue
Block a user