Merge pull request #15047 from jumpserver/pr@dev@excution

perf: Automation execution

apps/accounts/api/automations/change_secret_dashboard.py

@@ -2,6 +2,7 @@
 #
 from collections import defaultdict
 
+from django.core.cache import cache
 from django.http.response import JsonResponse
 from django.utils import timezone
 from rest_framework.views import APIView
@@ -25,6 +26,7 @@ class ChangeSecretDashboardApi(APIView):
 
     tp = AutomationTypes.change_secret
     task_name = 'accounts.tasks.automation.execute_account_automation_task'
+    ongoing_change_secret_cache_key = "ongoing_change_secret_cache_key"
 
     @lazyproperty
     def days(self):
@@ -146,24 +148,33 @@ class ChangeSecretDashboardApi(APIView):
             })
 
         if _all or query_params.get('total_count_ongoing_change_secret'):
-            execution_ids = []
-            inspect = app.control.inspect()
-            active_tasks = inspect.active()
-            if active_tasks:
-                for tasks in active_tasks.values():
-                    for task in tasks:
-                        _id = task.get('id')
-                        name = task.get('name')
-                        tp = task.get('kwargs', {}).get('tp')
-                        if name == self.task_name and tp == self.tp:
-                            execution_ids.append(_id)
+            ongoing_counts = cache.get(self.ongoing_change_secret_cache_key)
+            if ongoing_counts is None:
+                execution_ids = []
+                inspect = app.control.inspect()
+                active_tasks = inspect.active()
+                if active_tasks:
+                    for tasks in active_tasks.values():
+                        for task in tasks:
+                            _id = task.get('id')
+                            name = task.get('name')
+                            tp = task.get('kwargs', {}).get('tp')
+                            if name == self.task_name and tp == self.tp:
+                                execution_ids.append(_id)
 
-            snapshots = AutomationExecution.objects.filter(id__in=execution_ids).values_list('snapshot', flat=True)
+                snapshots = AutomationExecution.objects.filter(id__in=execution_ids).values_list('snapshot', flat=True)
 
-            asset_ids = {asset for i in snapshots for asset in i.get('assets', [])}
-            account_ids = {account for i in snapshots for account in i.get('accounts', [])}
-            data['total_count_ongoing_change_secret'] = len(execution_ids)
-            data['total_count_ongoing_change_secret_assets'] = len(asset_ids)
-            data['total_count_ongoing_change_secret_accounts'] = len(account_ids)
+                asset_ids = {asset for i in snapshots for asset in i.get('assets', [])}
+                account_ids = {account for i in snapshots for account in i.get('accounts', [])}
+
+                ongoing_counts = (len(execution_ids), len(asset_ids), len(account_ids))
+                data['total_count_ongoing_change_secret'] = ongoing_counts[0]
+                data['total_count_ongoing_change_secret_assets'] = ongoing_counts[1]
+                data['total_count_ongoing_change_secret_accounts'] = ongoing_counts[2]
+                cache.set(self.ongoing_change_secret_cache_key, ongoing_counts, 60)
+            else:
+                data['total_count_ongoing_change_secret'] = ongoing_counts[0]
+                data['total_count_ongoing_change_secret_assets'] = ongoing_counts[1]
+                data['total_count_ongoing_change_secret_accounts'] = ongoing_counts[2]
 
         return JsonResponse(data, status=200)

apps/accounts/automations/backup_account/handlers.py

@@ -36,17 +36,24 @@ class BaseAccountHandler:
             if isinstance(v, OrderedDict):
                 cls.unpack_data(v, data)
             else:
+                if isinstance(v, dict):
+                    v = v.get('label')
                 data[k] = v
         return data
 
     @classmethod
     def get_header_fields(cls, serializer: serializers.Serializer):
         try:
-            backup_fields = getattr(serializer, 'Meta').fields_backup
+            exclude_backup_fields = getattr(serializer, 'Meta').exclude_backup_fields
         except AttributeError:
-            backup_fields = serializer.fields.keys()
+            exclude_backup_fields = []
+        backup_fields = serializer.fields.keys()
+
         header_fields = {}
         for field in backup_fields:
+            if field in exclude_backup_fields:
+                continue
+
             v = serializer.fields[field]
             if isinstance(v, serializers.Serializer):
                 _fields = cls.get_header_fields(v)
@@ -189,8 +196,8 @@ class AccountBackupHandler:
                 attachment_list = [attachment]
             AccountBackupExecutionTaskMsg(name, user).publish(attachment_list)
 
-        for file in files:
-            os.remove(file)
+        # for file in files:
+        #     os.remove(file)
 
     def send_backup_obj_storage(self, files, recipients, password):
         if not files:
@@ -275,7 +282,8 @@ class AccountBackupHandler:
         else:
             recipients = recipients_part_one or recipients_part_two
             files = self.create_excel()
-            self.send_backup_mail(files, recipients)
+            print(files)
+            # self.send_backup_mail(files, recipients)
 
     def run(self):
         print('{}: {}'.format(_('Plan start'), local_now_display()))

apps/accounts/models/automations/backup_account.py

@@ -7,8 +7,8 @@ from django.utils.translation import gettext_lazy as _
 
 from accounts.const import AccountBackupType, AutomationTypes
 from common.db import fields
-from orgs.mixins.models import OrgManager
 from common.utils import get_logger
+from orgs.mixins.models import OrgManager
 from .base import AccountBaseAutomation
 
 __all__ = ['BackupAccountAutomation']
@@ -59,6 +59,7 @@ class BackupAccountAutomation(AccountBaseAutomation):
     def to_attr_json(self):
         attr_json = super().to_attr_json()
         attr_json.update({
+            'id': str(self.id),
             'types': self.types,
             'backup_type': self.backup_type,
             'is_password_divided_by_email': self.is_password_divided_by_email,

apps/accounts/serializers/account/account.py

@@ -460,6 +460,7 @@ class AccountSecretSerializer(SecretReadableMixin, AccountSerializer):
             'secret': {'write_only': False},
             'spec_info': {'label': _('Spec info')},
         }
+        exclude_backup_fields = ['passphrase', 'push_now', 'params']
 
 
 class AccountHistorySerializer(serializers.ModelSerializer):

apps/accounts/tasks/common.py

@@ -13,6 +13,6 @@ def quickstart_automation_by_snapshot(task_name, tp, task_snapshot=None):
         data['id'] = str(uuid.uuid4())
 
     execution = AutomationExecution.objects.create(
-        trigger=Trigger.manual, **data
+        type=tp, trigger=Trigger.manual, **data
     )
     execution.start()

apps/assets/models/automations/base.py

@@ -109,6 +109,7 @@ class BaseAutomation(PeriodTaskModelMixin, JMSOrgBaseModel):
 
         execution = self.execution_model.objects.create(
             id=eid,
+            type=self.type,
             trigger=trigger,
             automation=self,
             snapshot=self.to_attr_json(),

apps/assets/tasks/common.py

@@ -42,7 +42,7 @@ def quickstart_automation(task_name, tp, task_snapshot=None):
             break
 
     execution = AutomationExecution.objects.create(
-        trigger=Trigger.manual, **data
+        type=tp, trigger=Trigger.manual, **data
     )
     execution.start()
     return execution

apps/i18n/lina/en.json

@@ -1521,5 +1521,7 @@
     "IpGroupHelpText": "* indicates match all. for example: 192.168.10.1, 192.168.1.0/24, 10.1.1.1-10.1.1.20, 2001:db8:2de::e13, 2001:db8:1a:1110::/64",
     "SelectRisk": "Select risk",
     "IpGroup": "IP group",
+    "PublicIP": "Public IP",
+    "PrivateIP": "Private IP",
     "ExecuteAfterSaving": "Execute after saving"
 }

apps/i18n/lina/zh.json

@@ -1523,5 +1523,7 @@
     "Review": "审查",
     "SelectRisk": "选择风险",
     "IpGroup": "IP 组",
+    "PublicIP": "公有 IP",
+    "PrivateIP": "私有 IP",
     "ExecuteAfterSaving": "保存后执行"
 }