jumpserver/apps/accounts/api/account/task.py

64 lines
2.1 KiB
Python
Raw Normal View History

2024-01-29 06:40:09 +00:00
from django.db.models import Q
2023-02-21 05:00:04 +00:00
from rest_framework.generics import CreateAPIView
from accounts import serializers
2024-01-29 06:40:09 +00:00
from accounts.models import Account
2023-12-06 10:48:35 +00:00
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
2023-02-21 05:00:04 +00:00
__all__ = [
'AccountsTaskCreateAPI',
]
class AccountsTaskCreateAPI(CreateAPIView):
serializer_class = serializers.AccountTaskSerializer
2023-12-06 10:48:35 +00:00
permission_classes = (AccountTaskActionPermission,)
2023-02-21 05:00:04 +00:00
2023-12-06 10:48:35 +00:00
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()
2023-02-21 05:00:04 +00:00
2024-01-29 06:40:09 +00:00
@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]
2023-02-21 05:00:04 +00:00
2024-01-29 06:40:09 +00:00
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)
2023-02-21 05:00:04 +00:00
else:
2024-01-29 06:40:09 +00:00
raise ValueError(f"Invalid action: {action}")
2023-02-21 05:00:04 +00:00
data = getattr(serializer, '_data', {})
data["task"] = task.id
setattr(serializer, '_data', data)
return task