diff --git a/apps/assets/api/asset.py b/apps/assets/api/asset.py index 16ce3de59..1227b01c9 100644 --- a/apps/assets/api/asset.py +++ b/apps/assets/api/asset.py @@ -40,7 +40,9 @@ class AssetViewSet(IDInFilterMixin, LabelFilter, BulkModelViewSet): permission_classes = (IsSuperUserOrAppUser,) def get_queryset(self): - queryset = super().get_queryset() + queryset = super().get_queryset()\ + .prefetch_related('labels', 'nodes')\ + .select_related('admin_user') admin_user_id = self.request.query_params.get('admin_user_id') node_id = self.request.query_params.get("node_id") show_current_asset = self.request.query_params.get("show_current_asset") @@ -66,7 +68,6 @@ class AssetViewSet(IDInFilterMixin, LabelFilter, BulkModelViewSet): queryset = queryset.filter( nodes__key__regex='^{}(:[0-9]+)*$'.format(node.key), ).distinct() - return queryset diff --git a/apps/assets/models/asset.py b/apps/assets/models/asset.py index 7477cbc7f..6e6b6c678 100644 --- a/apps/assets/models/asset.py +++ b/apps/assets/models/asset.py @@ -59,42 +59,70 @@ class Asset(models.Model): ('Other', 'Other'), ) id = models.UUIDField(default=uuid.uuid4, primary_key=True) - ip = models.GenericIPAddressField(max_length=32, verbose_name=_('IP'), db_index=True) - hostname = models.CharField(max_length=128, unique=True, verbose_name=_('Hostname')) + ip = models.GenericIPAddressField(max_length=32, verbose_name=_('IP'), + db_index=True) + hostname = models.CharField(max_length=128, unique=True, + verbose_name=_('Hostname')) port = models.IntegerField(default=22, verbose_name=_('Port')) - platform = models.CharField(max_length=128, choices=PLATFORM_CHOICES, default='Linux', verbose_name=_('Platform')) - domain = models.ForeignKey("assets.Domain", null=True, blank=True, related_name='assets', verbose_name=_("Domain"), on_delete=models.SET_NULL) - nodes = models.ManyToManyField('assets.Node', default=default_node, related_name='assets', verbose_name=_("Nodes")) + platform = models.CharField(max_length=128, choices=PLATFORM_CHOICES, + default='Linux', verbose_name=_('Platform')) + domain = models.ForeignKey("assets.Domain", null=True, blank=True, + related_name='assets', verbose_name=_("Domain"), + on_delete=models.SET_NULL) + nodes = models.ManyToManyField('assets.Node', default=default_node, + related_name='assets', + verbose_name=_("Nodes")) is_active = models.BooleanField(default=True, verbose_name=_('Is active')) # Auth - admin_user = models.ForeignKey('assets.AdminUser', on_delete=models.PROTECT, null=True, verbose_name=_("Admin user")) + admin_user = models.ForeignKey('assets.AdminUser', on_delete=models.PROTECT, + null=True, verbose_name=_("Admin user")) # Some information - public_ip = models.GenericIPAddressField(max_length=32, blank=True, null=True, verbose_name=_('Public IP')) - number = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Asset number')) + public_ip = models.GenericIPAddressField(max_length=32, blank=True, + null=True, + verbose_name=_('Public IP')) + number = models.CharField(max_length=32, null=True, blank=True, + verbose_name=_('Asset number')) # Collect - vendor = models.CharField(max_length=64, null=True, blank=True, verbose_name=_('Vendor')) - model = models.CharField(max_length=54, null=True, blank=True, verbose_name=_('Model')) - sn = models.CharField(max_length=128, null=True, blank=True, verbose_name=_('Serial number')) - - cpu_model = models.CharField(max_length=64, null=True, blank=True, verbose_name=_('CPU model')) + vendor = models.CharField(max_length=64, null=True, blank=True, + verbose_name=_('Vendor')) + model = models.CharField(max_length=54, null=True, blank=True, + verbose_name=_('Model')) + sn = models.CharField(max_length=128, null=True, blank=True, + verbose_name=_('Serial number')) + + cpu_model = models.CharField(max_length=64, null=True, blank=True, + verbose_name=_('CPU model')) cpu_count = models.IntegerField(null=True, verbose_name=_('CPU count')) cpu_cores = models.IntegerField(null=True, verbose_name=_('CPU cores')) - memory = models.CharField(max_length=64, null=True, blank=True, verbose_name=_('Memory')) - disk_total = models.CharField(max_length=1024, null=True, blank=True, verbose_name=_('Disk total')) - disk_info = models.CharField(max_length=1024, null=True, blank=True, verbose_name=_('Disk info')) - - os = models.CharField(max_length=128, null=True, blank=True, verbose_name=_('OS')) - os_version = models.CharField(max_length=16, null=True, blank=True, verbose_name=_('OS version')) - os_arch = models.CharField(max_length=16, blank=True, null=True, verbose_name=_('OS arch')) - hostname_raw = models.CharField(max_length=128, blank=True, null=True, verbose_name=_('Hostname raw')) - - labels = models.ManyToManyField('assets.Label', blank=True, related_name='assets', verbose_name=_("Labels")) - created_by = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Created by')) - date_created = models.DateTimeField(auto_now_add=True, null=True, blank=True, verbose_name=_('Date created')) - comment = models.TextField(max_length=128, default='', blank=True, verbose_name=_('Comment')) + memory = models.CharField(max_length=64, null=True, blank=True, + verbose_name=_('Memory')) + disk_total = models.CharField(max_length=1024, null=True, blank=True, + verbose_name=_('Disk total')) + disk_info = models.CharField(max_length=1024, null=True, blank=True, + verbose_name=_('Disk info')) + + os = models.CharField(max_length=128, null=True, blank=True, + verbose_name=_('OS')) + os_version = models.CharField(max_length=16, null=True, blank=True, + verbose_name=_('OS version')) + os_arch = models.CharField(max_length=16, blank=True, null=True, + verbose_name=_('OS arch')) + hostname_raw = models.CharField(max_length=128, blank=True, null=True, + verbose_name=_('Hostname raw')) + + labels = models.ManyToManyField('assets.Label', blank=True, + related_name='assets', + verbose_name=_("Labels")) + created_by = models.CharField(max_length=32, null=True, blank=True, + verbose_name=_('Created by')) + date_created = models.DateTimeField(auto_now_add=True, null=True, + blank=True, + verbose_name=_('Date created')) + comment = models.TextField(max_length=128, default='', blank=True, + verbose_name=_('Comment')) objects = AssetManager() @@ -121,6 +149,22 @@ class Asset(models.Model): nodes = self.nodes.all() or [Node.root()] return nodes + @property + def nodes_cache_key(self): + key = "NODES_OF_{}".format(str(self.id)) + return key + + def get_nodes_or_cache(self): + cached = cache.get(self.nodes_cache_key) + if cached is not None: + return cached + nodes = list(self.get_nodes()) + cache.set(self.nodes_cache_key, nodes, 3600) + return nodes + + def expire_nodes_cache(self): + cache.delete(self.nodes_cache_key) + @property def hardware_info(self): if self.cpu_count: diff --git a/apps/assets/models/node.py b/apps/assets/models/node.py index 7c2fe7dcc..cb01cc3d8 100644 --- a/apps/assets/models/node.py +++ b/apps/assets/models/node.py @@ -13,9 +13,6 @@ __all__ = ['Node'] class Node(models.Model): id = models.UUIDField(default=uuid.uuid4, primary_key=True) key = models.CharField(unique=True, max_length=64, verbose_name=_("Key")) # '1:1:1:1' - # value = models.CharField( - # max_length=128, unique=True, verbose_name=_("Value") - # ) value = models.CharField(max_length=128, verbose_name=_("Value")) child_mark = models.IntegerField(default=0) date_create = models.DateTimeField(auto_now_add=True) @@ -31,10 +28,11 @@ class Node(models.Model): @property def full_value(self): - if self == self.__class__.root(): + ancestor = [a.value for a in self.ancestor] + if self.is_root(): return self.value - else: - return '{} / {}'.format(self.parent.full_value, self.value) + ancestor.append(self.value) + return ' / '.join(ancestor) @property def level(self): @@ -108,7 +106,6 @@ class Node(models.Model): def parent(self): if self.key == "0" or not self.key.startswith("0"): return self.__class__.root() - parent_key = ":".join(self.key.split(":")[:-1]) try: parent = self.__class__.objects.get(key=parent_key) @@ -132,16 +129,17 @@ class Node(models.Model): @property def ancestor(self): - _key = self.key.split(':') - ancestor_keys = [] - if self.is_root(): - return [self.__class__.root()] - - for i in range(len(_key)-1): - _key.pop() - ancestor_keys.append(':'.join(_key)) - return self.__class__.objects.filter(key__in=ancestor_keys) + ancestor = self.__class__.objects.filter(key='0') + else: + _key = self.key.split(':') + ancestor_keys = [] + for i in range(len(_key)-1): + _key.pop() + ancestor_keys.append(':'.join(_key)) + ancestor = self.__class__.objects.filter(key__in=ancestor_keys) + ancestor = list(ancestor) + return ancestor @property def ancestor_with_self(self): diff --git a/apps/assets/serializers/asset.py b/apps/assets/serializers/asset.py index 2c5aa253c..ac666e3a7 100644 --- a/apps/assets/serializers/asset.py +++ b/apps/assets/serializers/asset.py @@ -12,34 +12,11 @@ __all__ = [ ] -class NodeTMPSerializer(serializers.ModelSerializer): - parent = serializers.SerializerMethodField() - assets_amount = serializers.SerializerMethodField() - - class Meta: - model = Node - fields = ['id', 'key', 'value', 'parent', 'assets_amount', 'is_node'] - list_serializer_class = BulkListSerializer - - @staticmethod - def get_parent(obj): - return obj.parent.id - - @staticmethod - def get_assets_amount(obj): - return obj.get_all_assets().count() - - def get_fields(self): - fields = super().get_fields() - field = fields["key"] - field.required = False - return fields - - class AssetSerializer(BulkSerializerMixin, serializers.ModelSerializer): """ 资产的数据结构 """ + nodes = serializers.SerializerMethodField() class Meta: model = Asset @@ -54,6 +31,10 @@ class AssetSerializer(BulkSerializerMixin, serializers.ModelSerializer): ]) return fields + @staticmethod + def get_nodes(obj): + return [n.id for n in obj.get_nodes_or_cache()] + class AssetGrantedSerializer(serializers.ModelSerializer): """ diff --git a/apps/assets/signals_handler.py b/apps/assets/signals_handler.py index 06cd9f63e..16459c786 100644 --- a/apps/assets/signals_handler.py +++ b/apps/assets/signals_handler.py @@ -63,22 +63,31 @@ def on_system_user_assets_change(sender, instance=None, **kwargs): @receiver(m2m_changed, sender=Asset.nodes.through) def on_asset_node_changed(sender, instance=None, **kwargs): - if isinstance(instance, Asset) and kwargs['action'] == 'post_add': - logger.debug("Asset node change signal received") - nodes = kwargs['model'].objects.filter(pk__in=kwargs['pk_set']) - system_users_assets = defaultdict(set) - system_users = SystemUser.objects.filter(nodes__in=nodes) - for system_user in system_users: - system_users_assets[system_user].update({instance}) - for system_user, assets in system_users_assets.items(): - system_user.assets.add(*tuple(assets)) + if isinstance(instance, Asset): + instance.expire_nodes_cache() + if kwargs['action'] == 'post_add': + logger.debug("Asset node change signal received") + nodes = kwargs['model'].objects.filter(pk__in=kwargs['pk_set']) + system_users_assets = defaultdict(set) + system_users = SystemUser.objects.filter(nodes__in=nodes) + # 清理节点缓存 + for system_user in system_users: + system_users_assets[system_user].update({instance}) + for system_user, assets in system_users_assets.items(): + system_user.assets.add(*tuple(assets)) @receiver(m2m_changed, sender=Asset.nodes.through) def on_node_assets_changed(sender, instance=None, **kwargs): - if isinstance(instance, Node) and kwargs['action'] == 'post_add': - logger.debug("Node assets change signal received") + if isinstance(instance, Node): assets = kwargs['model'].objects.filter(pk__in=kwargs['pk_set']) - system_users = SystemUser.objects.filter(nodes=instance) - for system_user in system_users: - system_user.assets.add(*tuple(assets)) + # 清理资产节点缓存 + for asset in assets: + asset.expire_nodes_cache() + + if kwargs['action'] == 'post_add': + logger.debug("Node assets change signal received") + # 重新关联系统用户和资产的关系 + system_users = SystemUser.objects.filter(nodes=instance) + for system_user in system_users: + system_user.assets.add(*tuple(assets)) diff --git a/apps/assets/utils.py b/apps/assets/utils.py index 5fb5eae84..367c5e5f7 100644 --- a/apps/assets/utils.py +++ b/apps/assets/utils.py @@ -1,7 +1,8 @@ # ~*~ coding: utf-8 ~*~ # - +import os import paramiko +from paramiko.ssh_exception import SSHException from common.utils import get_object_or_none from .models import Asset, SystemUser, Label @@ -49,22 +50,23 @@ def test_gateway_connectability(gateway): """ client = paramiko.SSHClient() client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) - - proxy_command = [ - "ssh", "{}@{}".format(gateway.username, gateway.ip), - "-p", str(gateway.port), "-W", "127.0.0.1:{}".format(gateway.port), - ] - - if gateway.password: - proxy_command.insert(0, "sshpass -p '{}'".format(gateway.password)) - if gateway.private_key: - proxy_command.append("-i {}".format(gateway.private_key_file)) + proxy = paramiko.SSHClient() + proxy.load_host_keys(os.path.expanduser('~/.ssh/known_hosts')) + proxy.set_missing_host_key_policy(paramiko.AutoAddPolicy()) try: - sock = paramiko.ProxyCommand(" ".join(proxy_command)) - except paramiko.ProxyCommandFailure as e: + proxy.connect(gateway.ip, username=gateway.username, + password=gateway.password, + pkey=gateway.private_key_obj) + except(paramiko.AuthenticationException, + paramiko.BadAuthenticationType, + SSHException) as e: return False, str(e) + sock = proxy.get_transport().open_channel( + 'direct-tcpip', ('127.0.0.1', gateway.port), ('127.0.0.1', 0) + ) + try: client.connect("127.0.0.1", port=gateway.port, username=gateway.username, diff --git a/apps/perms/api.py b/apps/perms/api.py index e104cc1c4..bd2fb1139 100644 --- a/apps/perms/api.py +++ b/apps/perms/api.py @@ -147,13 +147,8 @@ class UserGrantedNodeAssetsApi(ListAPIView): user = get_object_or_404(User, id=user_id) else: user = self.request.user + node = get_object_or_404(Node, id=node_id) nodes = AssetPermissionUtil.get_user_nodes_with_assets(user) - node = get_object_or_none(Node, id=node_id) - - if not node: - unnode = [node for node in nodes if node.name == 'Unnode'] - node = unnode[0] if unnode else None - assets = nodes.get(node, []) for asset, system_users in assets.items(): asset.system_users_granted = system_users diff --git a/apps/perms/utils.py b/apps/perms/utils.py index adbd8cd1c..4844cda06 100644 --- a/apps/perms/utils.py +++ b/apps/perms/utils.py @@ -15,7 +15,7 @@ logger = get_logger(__file__) class Tree: def __init__(self): - self.__all_nodes = list(Node.objects.all()) + self.__all_nodes = list(Node.objects.all().prefetch_related('assets')) self.__node_asset_map = defaultdict(set) self.nodes = defaultdict(dict) self.root = Node.root() @@ -134,7 +134,7 @@ class AssetPermissionUtil: _assets = cls.get_user_group_assets(group) tree = Tree() for asset, _system_users in _assets.items(): - _nodes = asset.get_nodes() + _nodes = asset.get_nodes_or_cache() tree.add_nodes(_nodes) for node in _nodes: tree.nodes[node][asset].update(_system_users) diff --git a/jms b/jms index 0f07f0760..47eb81859 100755 --- a/jms +++ b/jms @@ -123,6 +123,7 @@ def start_gunicorn(): 'gunicorn', 'jumpserver.wsgi', '-b', bind, '-w', str(WORKERS), + '-k', 'eventlet', '--access-logformat', log_format, '-p', pid_file, ]