perf: 合并代码

apps/assets/automations/base/manager.py

@@ -49,7 +49,7 @@ class BasePlaybookManager:
         ansible_dir = settings.ANSIBLE_DIR
         dir_name = '{}_{}'.format(self.automation.name.replace(' ', '_'), self.execution.id)
         path = os.path.join(
-            ansible_dir, 'automations', self.automation.type,
+            ansible_dir, 'automations', self.execution.snapshot['type'],
             dir_name, timezone.now().strftime('%Y%m%d_%H%M%S')
         )
         if not os.path.exists(path):
@@ -77,9 +77,9 @@ class BasePlaybookManager:
 
     def generate_inventory(self, platformed_assets, inventory_path):
         inventory = JMSInventory(
+            manager=self,
             assets=platformed_assets,
             account_policy=self.ansible_account_policy,
-            host_callback=self.host_callback
         )
         inventory.write_to_file(inventory_path)
 

apps/assets/automations/change_secret/host/linux/main.yml

@@ -13,26 +13,26 @@
         name: "{{ account.username }}"
         password: "{{ account.secret | password_hash('sha512') }}"
         update_password: always
-      when: "{{ secret_type == 'password' }}"
+      when: secret_type == "password"
 
     - name: create user If it already exists, no operation will be performed
       ansible.builtin.user:
         name: "{{ account.username }}"
-      when: "{{ secret_type == 'ssh_key' }}"
+      when: secret_type == "ssh_key"
 
     - name: remove jumpserver ssh key
       ansible.builtin.lineinfile:
         dest: "{{ kwargs.dest }}"
         regexp: "{{ kwargs.regexp }}"
         state: absent
-      when: "{{ secret_type == 'ssh_key' and kwargs.strategy == 'set_jms' }}"
+      when: secret_type == "ssh_key" and kwargs.strategy == "set_jms"
 
     - name: Change SSH key
       ansible.builtin.authorized_key:
         user: "{{ account.username }}"
         key: "{{ account.secret }}"
         exclusive: "{{ kwargs.exclusive }}"
-      when: "{{ secret_type == 'ssh_key' }}"
+      when: secret_type == "ssh_key"
 
     - name: Refresh connection
       ansible.builtin.meta: reset_connection
@@ -44,7 +44,7 @@
         ansible_user: "{{ account.username }}"
         ansible_password: "{{ account.secret }}"
         ansible_become: no
-      when: "{{ secret_type == 'password' }}"
+      when: secret_type == "password"
 
     - name: Verify SSH key
       ansible.builtin.ping:
@@ -53,4 +53,4 @@
         ansible_user: "{{ account.username }}"
         ansible_ssh_private_key_file: "{{ account.private_key_path }}"
         ansible_become: no
-      when: "{{ secret_type == 'ssh_key' }}"
+      when: secret_type == "ssh_key"

apps/assets/automations/change_secret/host/windows/main.yml

@@ -1,7 +1,7 @@
 - hosts: demo
   gather_facts: no
   tasks:
-    - name: ping
+    - name: Test privileged account
       ansible.windows.win_ping:
 
 #    - name: Print variables
@@ -23,3 +23,4 @@
       vars:
         ansible_user: "{{ account.username }}"
         ansible_password: "{{ account.secret }}"
+      when: account.secret_type == "password"

apps/assets/automations/change_secret/manager.py

@@ -20,15 +20,15 @@ class ChangeSecretManager(BasePlaybookManager):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.method_hosts_mapper = defaultdict(list)
-        self.secret_type = self.execution.plan_snapshot.get('secret_type')
-        self.secret_strategy = self.execution.plan_snapshot['secret_strategy']
+        self.secret_type = self.execution.snapshot['secret_type']
+        self.secret_strategy = self.execution.snapshot['secret_strategy']
         self._password_generated = None
         self._ssh_key_generated = None
         self.name_recorder_mapper = {}  # 做个映射，方便后面处理
 
     @classmethod
     def method_type(cls):
-        return AutomationTypes.method_id_meta_mapper
+        return AutomationTypes.change_secret
 
     @lazyproperty
     def related_accounts(self):
@@ -53,7 +53,7 @@ class ChangeSecretManager(BasePlaybookManager):
         return key_path
 
     def generate_password(self):
-        kwargs = self.automation.plan_snapshot['password_rules'] or {}
+        kwargs = self.execution.snapshot['password_rules'] or {}
         length = int(kwargs.get('length', DEFAULT_PASSWORD_RULES['length']))
         symbol_set = kwargs.get('symbol_set')
         if symbol_set is None:
@@ -69,7 +69,7 @@ class ChangeSecretManager(BasePlaybookManager):
 
     def get_ssh_key(self):
         if self.secret_strategy == SecretStrategy.custom:
-            ssh_key = self.automation.plan_snapshot['ssh_key']
+            ssh_key = self.execution.snapshot['ssh_key']
             if not ssh_key:
                 raise ValueError("Automation SSH key must be set")
             return ssh_key
@@ -82,7 +82,7 @@ class ChangeSecretManager(BasePlaybookManager):
 
     def get_password(self):
         if self.secret_strategy == SecretStrategy.custom:
-            password = self.automation.plan_snapshot['password']
+            password = self.execution.snapshot['secret']
             if not password:
                 raise ValueError("Automation Password must be set")
             return password
@@ -106,7 +106,7 @@ class ChangeSecretManager(BasePlaybookManager):
         kwargs = {}
         if self.secret_type != SecretType.ssh_key:
             return kwargs
-        kwargs['strategy'] = self.automation.plan_snapshot['ssh_key_change_strategy']
+        kwargs['strategy'] = self.execution.snapshot['ssh_key_change_strategy']
         kwargs['exclusive'] = 'yes' if kwargs['strategy'] == SSHKeyStrategy.set else 'no'
 
         if kwargs['strategy'] == SSHKeyStrategy.set_jms:
@@ -123,11 +123,11 @@ class ChangeSecretManager(BasePlaybookManager):
         accounts = asset.accounts.all()
         if account:
             accounts = accounts.exclude(id=account.id)
-        if '*' not in self.automation.accounts:
-            accounts = accounts.filter(
-                username__in=self.automation.accounts, secret_type=self.secret_type
-            )
 
+        if '*' not in self.execution.snapshot['accounts']:
+            accounts = accounts.filter(username__in=self.execution.snapshot['accounts'])
+
+        accounts = accounts.filter(secret_type=self.secret_type)
         method_attr = getattr(automation, self.method_type() + '_method')
         method_hosts = self.method_hosts_mapper[method_attr]
         method_hosts = [h for h in method_hosts if h != host['name']]
@@ -153,7 +153,6 @@ class ChangeSecretManager(BasePlaybookManager):
                 new_secret = self.generate_public_key(new_secret)
 
             h['kwargs'] = self.get_kwargs(account, new_secret)
-
             h['account'] = {
                 'name': account.name,
                 'username': account.username,

apps/assets/models/automations/base.py

@@ -40,15 +40,18 @@ class BaseAutomation(CommonModelMixin, PeriodTaskModelMixin, OrgModelMixin):
     def get_register_task(self):
         raise NotImplementedError
 
+    def get_many_to_many_ids(self, field: str):
+        return [str(i) for i in getattr(self, field).all().values_list('id', flat=True)]
+
     def to_attr_json(self):
         return {
             'name': self.name,
             'type': self.type,
-            'org_id': self.org_id,
+            'org_id': str(self.org_id),
             'comment': self.comment,
             'accounts': self.accounts,
-            'nodes': list(self.nodes.all().values_list('id', flat=True)),
-            'assets': list(self.assets.all().values_list('id', flat=True)),
+            'nodes': self.get_many_to_many_ids('nodes'),
+            'assets': self.get_many_to_many_ids('assets'),
         }
 
     def execute(self, trigger=Trigger.manual):
@@ -59,7 +62,7 @@ class BaseAutomation(CommonModelMixin, PeriodTaskModelMixin, OrgModelMixin):
 
         execution = self.executions.model.objects.create(
             id=eid, trigger=trigger, automation=self,
-            plan_snapshot=self.to_attr_json(),
+            snapshot=self.to_attr_json(),
         )
         return execution.start()
 

apps/assets/models/automations/change_secret.py

@@ -18,7 +18,7 @@ class ChangeSecretAutomation(BaseAutomation):
     )
     secret_strategy = models.CharField(
         choices=SecretStrategy.choices, max_length=16,
-        default=SecretStrategy.random_one, verbose_name=_('Secret strategy')
+        default=SecretStrategy.custom, verbose_name=_('Secret strategy')
     )
     secret = fields.EncryptTextField(blank=True, null=True, verbose_name=_('Secret'))
     password_rules = models.JSONField(default=dict, verbose_name=_('Password rules'))

apps/ops/ansible/inventory.py

@@ -5,28 +5,26 @@ from collections import defaultdict
 
 from django.utils.translation import gettext as _
 
-
 __all__ = ['JMSInventory']
 
 
 class JMSInventory:
-    def __init__(self, assets, account_policy='smart', account_prefer='root,administrator', host_callback=None):
+    def __init__(self, manager, assets=None, account_policy='smart', account_prefer='root,administrator'):
         """
         :param assets:
         :param account_prefer: account username name if not set use account_policy
-        :param account_policy:
-        :param host_callback: after generate host, call this callback to modify host
+        :param account_policy: smart, privileged_must, privileged_first
         """
+        self.manager = manager
         self.assets = self.clean_assets(assets)
         self.account_prefer = account_prefer
         self.account_policy = account_policy
-        self.host_callback = host_callback
 
     @staticmethod
     def clean_assets(assets):
         from assets.models import Asset
         asset_ids = [asset.id for asset in assets]
-        assets = Asset.objects.filter(id__in=asset_ids, is_active=True)\
+        assets = Asset.objects.filter(id__in=asset_ids, is_active=True) \
             .prefetch_related('platform', 'domain', 'accounts')
         return assets
 
@@ -63,6 +61,7 @@ class JMSInventory:
         var = {
             'ansible_user': account.username,
         }
+
         if not account.secret:
             return var
         if account.secret_type == 'password':
@@ -79,7 +78,10 @@ class JMSInventory:
         ssh_protocol_matched = list(filter(lambda x: x.name == 'ssh', protocols))
         ssh_protocol = ssh_protocol_matched[0] if ssh_protocol_matched else None
         host['ansible_host'] = asset.address
-        host['ansible_port'] = ssh_protocol.port if ssh_protocol else 22
+        if asset.port == 0:
+            host['ansible_port'] = ssh_protocol.port if ssh_protocol else 22
+        else:
+            host['ansible_port'] = asset.port
 
         su_from = account.su_from
         if platform.su_enabled and su_from:
@@ -107,7 +109,7 @@ class JMSInventory:
                 'protocol': asset.protocol, 'port': asset.port,
                 'protocols': [{'name': p.name, 'port': p.port} for p in protocols],
             },
-            'jms_account':  {
+            'jms_account': {
                 'id': str(account.id), 'username': account.username,
                 'secret': account.secret, 'secret_type': account.secret_type
             } if account else None
@@ -170,8 +172,8 @@ class JMSInventory:
                 if not automation.ansible_enabled:
                     host['error'] = _('Ansible disabled')
 
-                if self.host_callback is not None:
-                    host = self.host_callback(
+                if self.manager.host_callback is not None:
+                    host = self.manager.host_callback(
                         host, asset=asset, account=account,
                         platform=platform, automation=automation,
                         path_dir=path_dir