You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
jumpserver/apps/accounts/api/account/task.py

64 lines
2.1 KiB

from django.db.models import Q
from rest_framework.generics import CreateAPIView
from accounts import serializers
from accounts.models import Account
from accounts.permissions import AccountTaskActionPermission
from accounts.tasks import (
remove_accounts_task, verify_accounts_connectivity_task, push_accounts_to_assets_task
)
from authentication.permissions import UserConfirmation, ConfirmType
__all__ = [
'AccountsTaskCreateAPI',
]
class AccountsTaskCreateAPI(CreateAPIView):
serializer_class = serializers.AccountTaskSerializer
permission_classes = (AccountTaskActionPermission,)
def get_permissions(self):
act = self.request.data.get('action')
if act == 'remove':
self.permission_classes = [
AccountTaskActionPermission,
UserConfirmation.require(ConfirmType.PASSWORD)
]
return super().get_permissions()
@staticmethod
def get_account_ids(data, action):
account_type = 'gather_accounts' if action == 'remove' else 'accounts'
accounts = data.get(account_type, [])
account_ids = [str(a.id) for a in accounts]
if action == 'remove':
return account_ids
assets = data.get('assets', [])
asset_ids = [str(a.id) for a in assets]
ids = Account.objects.filter(
Q(id__in=account_ids) | Q(asset_id__in=asset_ids)
).distinct().values_list('id', flat=True)
return [str(_id) for _id in ids]
def perform_create(self, serializer):
data = serializer.validated_data
action = data['action']
ids = self.get_account_ids(data, action)
if action == 'push':
task = push_accounts_to_assets_task.delay(ids, data.get('params'))
elif action == 'remove':
task = remove_accounts_task.delay(ids)
elif action == 'verify':
task = verify_accounts_connectivity_task.delay(ids)
else:
raise ValueError(f"Invalid action: {action}")
data = getattr(serializer, '_data', {})
data["task"] = task.id
setattr(serializer, '_data', data)
return task