perf: 重构扫描

jumpserver · Jan 8, 2025 · 7b38021 · 7b38021
2 parents 49f34b8 + 6036332
commit 7b38021
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 31 deletions.
diff --git a/apps/accounts/api/account/pam_dashboard.py b/apps/accounts/api/account/pam_dashboard.py
@@ -1,9 +1,11 @@
 # -*- coding: utf-8 -*-
 #
+from django.db.models import Count, F, Q
 from django.http.response import JsonResponse
 from rest_framework.views import APIView
 
 from accounts.models import Account, RiskChoice
+from assets.const import AllTypes
 from common.utils.timezone import local_monday
 
 __all__ = ['PamDashboardApi']
@@ -15,39 +17,77 @@ class PamDashboardApi(APIView):
         'GET': 'accounts.view_account',
     }
 
+    @staticmethod
+    def get_type_to_accounts():
+        result = Account.objects.annotate(type=F('asset__platform__type')). \
+            values('type').order_by('type').annotate(total=Count(1))
+        all_types_dict = dict(AllTypes.choices())
+
+        result = [
+            {
+                **i,
+                'label': all_types_dict.get(i['type'], i['type'])
+            }
+            for i in result
+        ]
+        return result
+
     def get(self, request, *args, **kwargs):
-        query_params = self.request.query_params
         data = {}
+        monday_time = local_monday()
+        query_params = self.request.query_params
+
+        account_stats = Account.objects.aggregate(
+            total_count=Count('id'),
+            privileged_count=Count('id', filter=Q(privileged=True)),
+            connectivity_ok_count=Count('id', filter=Q(connectivity='ok')),
+            secret_reset_count=Count('id', filter=Q(secret_reset=True)),
+            unavailable_count=Count('id', filter=Q(is_active=False)),
+            week_add_count=Count('id', filter=Q(date_created__gte=monday_time)),
+        )
+
+        _all = query_params.get('all')
+
+        if _all or query_params.get('total_accounts'):
+            data['total_accounts'] = account_stats['total_count']
+
+        if _all or query_params.get('total_week_add_accounts'):
+            data['total_week_add_accounts'] = account_stats['week_add_count']
+
+        if _all or query_params.get('total_privileged_accounts'):
+            data['total_privileged_accounts'] = account_stats['privileged_count']
 
-        account_count = Account.objects.count()
-        privileged_account_count = Account.objects.filter(privileged=True).count()
-
-        if query_params.get('total_accounts'):
-            data['total_accounts'] = account_count
+        if _all or query_params.get('total_connectivity_ok_accounts'):
+            data['total_connectivity_ok_accounts'] = account_stats['connectivity_ok_count']
 
-        if query_params.get('total_week_add_accounts'):
-            monday_time = local_monday()
-            data['total_week_add_accounts'] = Account.objects.filter(date_created__gte=monday_time).count()
+        if _all or query_params.get('total_secret_reset_accounts'):
+            data['total_secret_reset_accounts'] = account_stats['secret_reset_count']
 
-        if query_params.get('total_privileged_accounts'):
-            data['total_privileged_accounts'] = privileged_account_count
+        if _all or query_params.get('total_ordinary_accounts'):
+            data['total_ordinary_accounts'] = account_stats['total_count'] - account_stats['privileged_count']
 
-        if query_params.get('total_ordinary_accounts'):
-            data['total_ordinary_accounts'] = account_count - privileged_account_count
+        if _all or query_params.get('total_unavailable_accounts'):
+            data['total_unavailable_accounts'] = account_stats['unavailable_count']
 
-        if query_params.get('total_unmanaged_accounts'):
-            data['total_unmanaged_accounts'] = Account.get_risks(risk_type=RiskChoice.new_found).count()
+        if _all or query_params.get('total_unmanaged_accounts'):
+            data['total_unmanaged_accounts'] = Account.get_risks(
+                risk_type=RiskChoice.new_found).count()
 
-        if query_params.get('total_unavailable_accounts'):
-            data['total_unavailable_accounts'] = Account.objects.filter(is_active=False).count()
+        if _all or query_params.get('total_long_time_no_login_accounts'):
+            data['total_long_time_no_login_accounts'] = Account.get_risks(
+                risk_type=RiskChoice.long_time_no_login).count()
 
-        if query_params.get('total_long_time_no_login_accounts'):
-            data['total_long_time_no_login_accounts'] = Account.get_risks(risk_type=RiskChoice.long_time_no_login).count()
+        if _all or query_params.get('total_weak_password_accounts'):
+            data['total_weak_password_accounts'] = Account.get_risks(
+                risk_type=RiskChoice.weak_password).count()
 
-        if query_params.get('total_weak_password_accounts'):
-            data['total_weak_password_accounts'] = Account.get_risks(risk_type=RiskChoice.weak_password).count()
+        if _all or query_params.get('total_long_time_change_password_accounts'):
+            data['total_long_time_change_password_accounts'] = Account.get_risks(
+                risk_type=RiskChoice.long_time_password).count()
 
-        if query_params.get('total_long_time_change_password_accounts'):
-            data['total_long_time_change_password_accounts'] = Account.get_risks(risk_type=RiskChoice.long_time_password).count()
+        if _all or query_params.get('total_count_type_to_accounts_amount'):
+            data.update({
+                'total_count_type_to_accounts_amount': self.get_type_to_accounts(),
+            })
 
         return JsonResponse(data, status=200)
diff --git a/apps/accounts/models/account.py b/apps/accounts/models/account.py
@@ -76,8 +76,9 @@ class Account(AbsConnectivity, LabeledMixin, BaseAccount):
     date_last_login = models.DateTimeField(null=True, blank=True, verbose_name=_('Date last access'))
     login_by = models.CharField(max_length=128, null=True, blank=True, verbose_name=_('Access by'))
     date_change_secret = models.DateTimeField(null=True, blank=True, verbose_name=_('Date change secret'))
-    change_secret_status = models.CharField(max_length=16, null=True, blank=True,
-                                            verbose_name=_('Change secret status'))
+    change_secret_status = models.CharField(
+        max_length=16, null=True, blank=True, verbose_name=_('Change secret status')
+    )
 
     class Meta:
         verbose_name = _('Account')

diff --git a/apps/accounts/risk_handlers.py b/apps/accounts/risk_handlers.py
@@ -1,14 +1,15 @@
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 
-from accounts.const import AutomationTypes
+from accounts.const import AutomationTypes, Source
 from accounts.models import (
     GatheredAccount,
     AccountRisk,
     SecretType,
-    AutomationExecution, RiskChoice,
+    AutomationExecution, RiskChoice, Account
 )
 from common.const import ConfirmOrIgnore
+from common.utils import random_string
 
 TYPE_CHOICES = [
     ("ignore", _("Ignore")),
@@ -17,7 +18,7 @@
     ("delete_both", _("Delete remote")),
     ("add_account", _("Add account")),
     ("change_password_add", _("Change password and Add")),
-    ("change_password", _("Change password")),
+    ("change_password", _("Change password"))
 ]
 
 
@@ -107,9 +108,6 @@ def _handle_delete(self, delete="both"):
     def handle_delete_both(self):
         self._handle_delete(delete="both")
 
-    def handle_change_password_add(self):
-        pass
-
     def handle_change_password(self):
         asset = self.asset
         execution = AutomationExecution()
@@ -124,3 +122,34 @@ def handle_change_password(self):
         execution.save()
         execution.start()
         return execution.summary
+
+    def handle_change_password_add(self):
+        asset = self.asset
+        secret_type = SecretType.PASSWORD
+        secret = random_string(30)
+        account_data = {
+            "username": self.username,
+            "name": f'{self.username}-{secret_type}',
+            "secret_type": SecretType.PASSWORD,
+            "source": Source.DISCOVERY,
+            "asset": asset,
+            "secret": secret
+        }
+        account, _ = self.asset.accounts.get_or_create(defaults=account_data, username=self.username)
+        execution = AutomationExecution()
+        execution.snapshot = {
+            "assets": [str(asset.id)],
+            "accounts": [str(account.id)],
+            "type": AutomationTypes.push_account,
+            "secret_type": secret_type,
+            'nodes': [],
+            'org_id': self.asset.org_id,
+            "secret_strategy": "random",
+            "secret": secret,
+            'ssh_key_change_strategy': 'set_jms',
+            'check_conn_after_change': True,
+            "name": "Push account password: {}@{}".format(self.username, asset.name),
+        }
+        execution.save()
+        execution.start()
+        return execution.summary