diff --git a/apps/assets/tasks/push_system_user.py b/apps/assets/tasks/push_system_user.py index 5b61c656e..7506d1fe2 100644 --- a/apps/assets/tasks/push_system_user.py +++ b/apps/assets/tasks/push_system_user.py @@ -7,7 +7,7 @@ from django.utils.translation import ugettext as _ from django.db.models import Empty from common.utils import encrypt_password, get_logger -from assets.models import SystemUser, Asset +from assets.models import SystemUser, Asset, AuthBook from orgs.utils import org_aware_func, tmp_to_root_org from . import const from .utils import clean_ansible_task_hosts, group_asset_by_platform @@ -190,15 +190,12 @@ def get_push_system_user_tasks(system_user, platform="unixlike", username=None): @org_aware_func("system_user") def push_system_user_util(system_user, assets, task_name, username=None): from ops.utils import update_or_create_ansible_task - hosts = clean_ansible_task_hosts(assets, system_user=system_user) - if not hosts: + assets = clean_ansible_task_hosts(assets, system_user=system_user) + if not assets: return {} - platform_hosts_map = {} - hosts_sorted = sorted(hosts, key=group_asset_by_platform) - platform_hosts = groupby(hosts_sorted, key=group_asset_by_platform) - for i in platform_hosts: - platform_hosts_map[i[0]] = list(i[1]) + assets_sorted = sorted(assets, key=group_asset_by_platform) + platform_hosts = groupby(assets_sorted, key=group_asset_by_platform) def run_task(_tasks, _hosts): if not _tasks: @@ -209,23 +206,51 @@ def push_system_user_util(system_user, assets, task_name, username=None): ) task.run() - for platform, _hosts in platform_hosts_map.items(): - if not _hosts: + if system_user.username_same_with_user: + if username is None: + # 动态系统用户,但是没有指定 username + usernames = list(system_user.users.all().values_list('username', flat=True).distinct()) + else: + usernames = [username] + else: + # 非动态系统用户指定 username 无效 + assert username is None, 'Only Dynamic user can assign `username`' + usernames = [system_user.username] + + for platform, _assets in platform_hosts: + _assets = list(_assets) + if not _assets: continue print(_("Start push system user for platform: [{}]").format(platform)) - print(_("Hosts count: {}").format(len(_hosts))) + print(_("Hosts count: {}").format(len(_assets))) - # 如果没有特殊密码设置,就不需要单独推送某台机器了 - if not system_user.has_special_auth(username=username): - logger.debug("System user not has special auth") - tasks = get_push_system_user_tasks(system_user, platform, username=username) - run_task(tasks, _hosts) - continue + id_asset_map = {_asset.id: _asset for _asset in _assets} + assets_id = id_asset_map.keys() + no_special_auth = [] + special_auth_set = set() - for _host in _hosts: - system_user.load_asset_special_auth(_host, username=username) - tasks = get_push_system_user_tasks(system_user, platform, username=username) - run_task(tasks, [_host]) + auth_books = AuthBook.objects.filter(username__in=usernames, asset_id__in=assets_id) + + for auth_book in auth_books: + special_auth_set.add((auth_book.username, auth_book.asset_id)) + + for _username in usernames: + no_special_assets = [] + for asset_id in assets_id: + if (_username, asset_id) not in special_auth_set: + no_special_assets.append(id_asset_map[asset_id]) + if no_special_assets: + no_special_auth.append((_username, no_special_assets)) + + for _username, no_special_assets in no_special_auth: + tasks = get_push_system_user_tasks(system_user, platform, username=_username) + run_task(tasks, no_special_assets) + + for auth_book in auth_books: + system_user._merge_auth(auth_book) + tasks = get_push_system_user_tasks(system_user, platform, username=auth_book.username) + asset = id_asset_map[auth_book.asset_id] + run_task(tasks, [asset]) @shared_task(queue="ansible") @@ -264,16 +289,6 @@ def push_system_user_to_assets(system_user_id, assets_id, username=None): assets = get_objects(Asset, assets_id) task_name = _("Push system users to assets: {}").format(system_user.name) - if username is None and system_user.username_same_with_user: - # 动态系统用户,把与系统用户关联的所有用户推送到新关联的资产上 - usernames = system_user.users.all().values_list('username', flat=True).distinct() - ret = [] - for username in usernames: - ret.append( - push_system_user_util(system_user, assets, task_name, username=username) - ) - return ret - return push_system_user_util(system_user, assets, task_name, username=username) # @shared_task