jumpserver/apps/users/api/user.py

# ~*~ coding: utf-8 ~*~
from collections import defaultdict

from django.utils.translation import ugettext as _
from rest_framework.decorators import action
from rest_framework import generics
from rest_framework.response import Response
from rest_framework_bulk import BulkModelViewSet
from django.db.models import Prefetch

from users.notifications import ResetMFAMsg
from common.permissions import (
    IsOrgAdmin, IsOrgAdminOrAppUser,
    CanUpdateDeleteUser, IsSuperUser
)
from common.mixins import CommonApiMixin
from common.utils import get_logger
from orgs.utils import current_org
from orgs.models import ROLE as ORG_ROLE, OrganizationMember
from users.utils import LoginBlockUtil, MFABlockUtils
from .. import serializers
from ..serializers import UserSerializer, UserRetrieveSerializer, MiniUserSerializer, InviteSerializer
from .mixins import UserQuerysetMixin
from ..models import User
from ..signals import post_user_create
from ..filters import OrgRoleUserFilterBackend, UserFilter

logger = get_logger(__name__)
__all__ = [
    'UserViewSet', 'UserChangePasswordApi',
    'UserUnblockPKApi', 'UserResetOTPApi',
]


class UserViewSet(CommonApiMixin, UserQuerysetMixin, BulkModelViewSet):
    filterset_class = UserFilter
    search_fields = ('username', 'email', 'name', 'id', 'source', 'role')
    permission_classes = (IsOrgAdmin, CanUpdateDeleteUser)
    serializer_classes = {
        'default': UserSerializer,
        'retrieve': UserRetrieveSerializer,
        'suggestion': MiniUserSerializer,
        'invite': InviteSerializer,
    }
    extra_filter_backends = [OrgRoleUserFilterBackend]

    def get_queryset(self):
        queryset = super().get_queryset().prefetch_related(
            'groups'
        )
        if not current_org.is_root():
            # 为在列表中计算用户在真实组织里的角色
            queryset = queryset.prefetch_related(
                Prefetch(
                    'm2m_org_members',
                    queryset=OrganizationMember.objects.filter(org__id=current_org.id)
                )
            )
        return queryset

    def send_created_signal(self, users):
        if not isinstance(users, list):
            users = [users]
        for user in users:
            post_user_create.send(self.__class__, user=user)

    @staticmethod
    def set_users_to_org(users, org_roles, update=False):
        # 只有真实存在的组织才真正关联用户
        if not current_org or current_org.is_root():
            return
        for user, roles in zip(users, org_roles):
            if update and roles is None:
                continue
            if not roles:
                # 当前组织创建的用户，至少是该组织的`User`
                roles = [ORG_ROLE.USER]
            OrganizationMember.objects.set_user_roles(current_org, user, roles)

    def perform_create(self, serializer):
        org_roles = self.get_serializer_org_roles(serializer)
        # 创建用户
        users = serializer.save()
        if isinstance(users, User):
            users = [users]
        self.set_users_to_org(users, org_roles)
        self.send_created_signal(users)

    def get_permissions(self):
        if self.action in ["retrieve", "list"]:
            if self.request.query_params.get('all'):
                self.permission_classes = (IsSuperUser,)
            else:
                self.permission_classes = (IsOrgAdminOrAppUser,)
        elif self.action in ['destroy']:
            self.permission_classes = (IsSuperUser,)
        return super().get_permissions()

    def perform_bulk_destroy(self, objects):
        for obj in objects:
            self.check_object_permissions(self.request, obj)
            self.perform_destroy(obj)

    @staticmethod
    def get_serializer_org_roles(serializer):
        validated_data = serializer.validated_data
        # `org_roles` 先 `pop`
        if isinstance(validated_data, list):
            org_roles = [item.pop('org_roles', None) for item in validated_data]
        else:
            org_roles = [validated_data.pop('org_roles', None)]
        return org_roles

    def perform_update(self, serializer):
        org_roles = self.get_serializer_org_roles(serializer)
        users = serializer.save()
        if isinstance(users, User):
            users = [users]
        self.set_users_to_org(users, org_roles, update=True)

    def perform_bulk_update(self, serializer):
        # TODO: 需要测试
        user_ids = [
            d.get("id") or d.get("pk") for d in serializer.validated_data
        ]
        users = current_org.get_members().filter(id__in=user_ids)
        for user in users:
            self.check_object_permissions(self.request, user)
        return super().perform_bulk_update(serializer)

    @action(methods=['get'], detail=False, permission_classes=(IsOrgAdmin,))
    def suggestion(self, request):
        queryset = User.objects.exclude(role=User.ROLE.APP)
        queryset = self.filter_queryset(queryset)[:3]
        serializer = self.get_serializer(queryset, many=True)
        return Response(serializer.data)

    @action(methods=['post'], detail=False, permission_classes=(IsOrgAdmin,))
    def invite(self, request):
        data = request.data
        if not isinstance(data, list):
            data = [request.data]
        if not current_org or current_org.is_root():
            error = {"error": "Not a valid org"}
            return Response(error, status=400)

        serializer_cls = self.get_serializer_class()
        serializer = serializer_cls(data=data, many=True)
        serializer.is_valid(raise_exception=True)
        validated_data = serializer.validated_data

        users_by_role = defaultdict(list)
        for i in validated_data:
            users_by_role[i['role']].append(i['user'])

        OrganizationMember.objects.add_users_by_role(
            current_org,
            users=users_by_role[ORG_ROLE.USER],
            admins=users_by_role[ORG_ROLE.ADMIN],
            auditors=users_by_role[ORG_ROLE.AUDITOR]
        )
        return Response(serializer.data, status=201)

    @action(methods=['post'], detail=True, permission_classes=(IsOrgAdmin,))
    def remove(self, request, *args, **kwargs):
        instance = self.get_object()
        instance.remove()
        return Response(status=204)

    @action(methods=['post'], detail=False, permission_classes=(IsOrgAdmin,), url_path='remove')
    def bulk_remove(self, request, *args, **kwargs):
        qs = self.get_queryset()
        filtered = self.filter_queryset(qs)

        for instance in filtered:
            instance.remove()
        return Response(status=204)


class UserChangePasswordApi(UserQuerysetMixin, generics.UpdateAPIView):
    permission_classes = (IsOrgAdmin,)
    serializer_class = serializers.ChangeUserPasswordSerializer

    def perform_update(self, serializer):
        user = self.get_object()
        user.password_raw = serializer.validated_data["password"]
        user.save()


class UserUnblockPKApi(UserQuerysetMixin, generics.UpdateAPIView):
    permission_classes = (IsOrgAdmin,)
    serializer_class = serializers.UserSerializer

    def perform_update(self, serializer):
        user = self.get_object()
        username = user.username if user else ''
        LoginBlockUtil.unblock_user(username)
        MFABlockUtils.unblock_user(username)


class UserResetOTPApi(UserQuerysetMixin, generics.RetrieveAPIView):
    permission_classes = (IsOrgAdmin,)
    serializer_class = serializers.ResetOTPSerializer

    def retrieve(self, request, *args, **kwargs):
        user = self.get_object() if kwargs.get('pk') else request.user
        if user == request.user:
            msg = _("Could not reset self otp, use profile reset instead")
            return Response({"error": msg}, status=401)
        if user.mfa_enabled:
            user.reset_mfa()
            user.save()

            ResetMFAMsg(user).publish_async()
        return Response({"msg": "success"})