From a99badd41ee10e5856199478b5d40c7e0701ab7e Mon Sep 17 00:00:00 2001 From: ibuler Date: Fri, 21 Mar 2025 19:08:21 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=20token=20=E7=9B=B4?= =?UTF-8?q?=E8=BF=9E=E7=9A=84=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/authentication/api/connection_token.py | 20 +++++++---- apps/authentication/const.py | 6 ++-- .../authentication/models/connection_token.py | 34 ++++++++++++++++--- .../serializers/connection_token.py | 8 +++-- 4 files changed, 52 insertions(+), 16 deletions(-) diff --git a/apps/authentication/api/connection_token.py b/apps/authentication/api/connection_token.py index 6203180f3..e0173cd3c 100644 --- a/apps/authentication/api/connection_token.py +++ b/apps/authentication/api/connection_token.py @@ -35,7 +35,8 @@ from ..models import ConnectionToken, AdminConnectionToken, date_expired_default from ..serializers import ( ConnectionTokenSerializer, ConnectionTokenSecretSerializer, SuperConnectionTokenSerializer, ConnectTokenAppletOptionSerializer, - ConnectionTokenReusableSerializer, ConnectTokenVirtualAppOptionSerializer + ConnectionTokenReusableSerializer, ConnectTokenVirtualAppOptionSerializer, + AdminConnectionTokenSerializer, ) __all__ = ['ConnectionTokenViewSet', 'SuperConnectionTokenViewSet', 'AdminConnectionTokenViewSet'] @@ -436,8 +437,7 @@ class ConnectionTokenViewSet(AuthFaceMixin, ExtraActionApiMixin, RootOrgViewMixi if ticket or self.need_face_verify: data['is_active'] = False if self.face_monitor_token: - FaceMonitorContext.get_or_create_context(self.face_monitor_token, - self.request.user.id) + FaceMonitorContext.get_or_create_context(self.face_monitor_token, self.request.user.id) data['face_monitor_token'] = self.face_monitor_token return data @@ -617,7 +617,7 @@ class SuperConnectionTokenViewSet(ConnectionTokenViewSet): raise PermissionDenied('Not allow to view secret') token_id = request.data.get('id') or '' - token = get_object_or_404(ConnectionToken, pk=token_id) + token = ConnectionToken.get_typed_connection_token(token_id) token.is_valid() serializer = self.get_serializer(instance=token) @@ -670,6 +670,9 @@ class SuperConnectionTokenViewSet(ConnectionTokenViewSet): class AdminConnectionTokenViewSet(ConnectionTokenViewSet): + serializer_classes = { + 'default': AdminConnectionTokenSerializer, + } def check_permissions(self, request): user = request.user @@ -677,11 +680,16 @@ class AdminConnectionTokenViewSet(ConnectionTokenViewSet): self.permission_denied(request) def get_queryset(self): - return AdminConnectionToken.objects.all() + return AdminConnectionToken.objects.all().filter(user=self.request.user) def get_permed_account(self, user, asset, account_name, protocol): + """ + 管理员 token 可以访问所有资产的账号 + """ with tmp_to_org(asset.org): - account = asset.accounts.all().active().get(name=account_name) + account = asset.accounts.all().active().filter(name=account_name).first() + if not account: + return None account.actions = ActionChoices.all() account.date_expired = timezone.now() + timezone.timedelta(days=365) return account diff --git a/apps/authentication/const.py b/apps/authentication/const.py index 22104cfdc..f71ab8ce2 100644 --- a/apps/authentication/const.py +++ b/apps/authentication/const.py @@ -48,6 +48,6 @@ class FaceMonitorActionChoices(TextChoices): class ConnectionTokenType(TextChoices): - ADMIN = 'admin', 'Admin' - SUPER = 'super', 'Super' - USER = 'user', 'User' + ADMIN = 'admin', 'Admin' # 管理员 Token, 可以访问所有资源 + SUPER = 'super', 'Super' # 超级 Token, 可以为不同用户生成的 Token, 遵守用户 Token 的限制 + USER = 'user', 'User' # 用户 Token, 只能访问指定资源 diff --git a/apps/authentication/models/connection_token.py b/apps/authentication/models/connection_token.py index bbf4cf990..e25476cba 100644 --- a/apps/authentication/models/connection_token.py +++ b/apps/authentication/models/connection_token.py @@ -5,8 +5,10 @@ from datetime import timedelta from django.conf import settings from django.core.cache import cache from django.db import models +from django.shortcuts import get_object_or_404 from django.utils import timezone from django.utils.translation import gettext_lazy as _ +from django.forms.models import model_to_dict from rest_framework.exceptions import PermissionDenied from accounts.models import VirtualAccount @@ -71,9 +73,19 @@ class ConnectionToken(JMSOrgBaseModel): verbose_name = _('Connection token') def save(self, *args, **kwargs): - self.type = self._meta.model._type + self.type = self._type return super().save(*args, **kwargs) + @classmethod + def get_typed_connection_token(cls, token_id): + token = get_object_or_404(cls, id=token_id) + + if token.type == ConnectionTokenType.ADMIN.value: + token = AdminConnectionToken.objects.get(id=token_id) + else: + token = ConnectionToken.objects.get(id=token_id) + return token + @property def is_expired(self): return self.date_expired < timezone.now() @@ -112,13 +124,16 @@ class ConnectionToken(JMSOrgBaseModel): self.date_expired = date_expired_default() self.save() - @lazyproperty - def permed_account(self): + def get_permed_account(self): from perms.utils import PermAssetDetailUtil permed_account = PermAssetDetailUtil(self.user, self.asset) \ .validate_permission(self.account, self.protocol) return permed_account + @lazyproperty + def permed_account(self): + return self.get_permed_account() + @lazyproperty def actions(self): return self.permed_account.actions @@ -148,7 +163,8 @@ class ConnectionToken(JMSOrgBaseModel): if timezone.now() - self.date_created < timedelta(seconds=60): return True, None - if not self.permed_account or not self.permed_account.actions: + permed_account = self.get_permed_account() + if not permed_account or not permed_account.actions: msg = 'user `{}` not has asset `{}` permission for login `{}`'.format( self.user, self.asset, self.account ) @@ -317,4 +333,12 @@ class AdminConnectionToken(ConnectionToken): return (timezone.now() + timezone.timedelta(days=365)).timestamp() def is_valid(self): - return True + return super().is_valid() + + def get_permed_account(self): + account = self.asset.accounts.filter(name=self.account).first() + if not account: + return None + account.actions = ActionChoices.all() + account.date_expired = timezone.now() + timezone.timedelta(days=5) + return account diff --git a/apps/authentication/serializers/connection_token.py b/apps/authentication/serializers/connection_token.py index f519f739e..cb8759fef 100644 --- a/apps/authentication/serializers/connection_token.py +++ b/apps/authentication/serializers/connection_token.py @@ -4,11 +4,11 @@ from rest_framework import serializers from common.serializers import CommonModelSerializer from common.serializers.fields import EncryptedField from perms.serializers.permission import ActionChoicesField -from ..models import ConnectionToken +from ..models import ConnectionToken, AdminConnectionToken __all__ = [ 'ConnectionTokenSerializer', 'SuperConnectionTokenSerializer', - 'ConnectionTokenReusableSerializer', + 'ConnectionTokenReusableSerializer', 'AdminConnectionTokenSerializer', ] @@ -74,3 +74,7 @@ class SuperConnectionTokenSerializer(ConnectionTokenSerializer): def get_user(self, attrs): return attrs.get('user') + +class AdminConnectionTokenSerializer(ConnectionTokenSerializer): + class Meta(ConnectionTokenSerializer.Meta): + model = AdminConnectionToken