perf(authentication): 优化connection token的使用

pull/5541/head
ibuler 2021-01-21 13:50:29 +08:00 committed by Jiangjie.Bai
parent dd5b2b9101
commit 23afe81ff5
4 changed files with 66 additions and 27 deletions

View File

@ -2,54 +2,59 @@
# #
import uuid import uuid
from django.conf import settings
from django.core.cache import cache from django.core.cache import cache
from django.shortcuts import get_object_or_404
from rest_framework.response import Response from rest_framework.response import Response
from rest_framework.views import APIView from rest_framework.generics import CreateAPIView
from common.utils import get_logger from common.utils import get_logger
from common.permissions import IsOrgAdminOrAppUser from common.permissions import IsSuperUserOrAppUser
from orgs.mixins.api import RootOrgViewMixin from orgs.mixins.api import RootOrgViewMixin
from users.models import User
from assets.models import Asset, SystemUser from ..serializers import ConnectionTokenSerializer
logger = get_logger(__name__) logger = get_logger(__name__)
__all__ = [ __all__ = ['UserConnectionTokenApi']
'UserConnectionTokenApi',
]
class UserConnectionTokenApi(RootOrgViewMixin, APIView): class UserConnectionTokenApi(RootOrgViewMixin, CreateAPIView):
permission_classes = (IsOrgAdminOrAppUser,) permission_classes = (IsSuperUserOrAppUser,)
serializer_class = ConnectionTokenSerializer
def post(self, request): def perform_create(self, serializer):
user_id = request.data.get('user', '') user = serializer.validated_data['user']
asset_id = request.data.get('asset', '') asset = serializer.validated_data['asset']
system_user_id = request.data.get('system_user', '') system_user = serializer.validated_data['system_user']
token = str(uuid.uuid4()) token = str(uuid.uuid4())
user = get_object_or_404(User, id=user_id)
asset = get_object_or_404(Asset, id=asset_id)
system_user = get_object_or_404(SystemUser, id=system_user_id)
value = { value = {
'user': user_id, 'user': str(user.id),
'username': user.username, 'username': user.username,
'asset': asset_id, 'asset': str(asset.id),
'hostname': asset.hostname, 'hostname': asset.hostname,
'system_user': system_user_id, 'system_user': str(system_user.id),
'system_user_name': system_user.name 'system_user_name': system_user.name
} }
cache.set(token, value, timeout=20) cache.set(token, value, timeout=20)
return token
def create(self, request, *args, **kwargs):
if not settings.CONNECTION_TOKEN_ENABLED:
data = {'error': 'Connection token disabled'}
return Response(data, status=400)
if not request.user.is_superuser:
data = {'error': 'Only super user can create token'}
return Response(data, status=403)
serializer = self.get_serializer(data=request.data)
serializer.is_valid(raise_exception=True)
token = self.perform_create(serializer)
return Response({"token": token}, status=201) return Response({"token": token}, status=201)
def get(self, request): def get(self, request):
token = request.query_params.get('token') token = request.query_params.get('token')
user_only = request.query_params.get('user-only', None)
value = cache.get(token, None) value = cache.get(token, None)
if not value: if not value:
return Response('', status=404) return Response('', status=404)
return Response(value)
if not user_only:
return Response(value)
else:
return Response({'user': value['user']})

View File

@ -11,6 +11,7 @@ from .models import AccessKey, LoginConfirmSetting, SSOToken
__all__ = [ __all__ = [
'AccessKeySerializer', 'OtpVerifySerializer', 'BearerTokenSerializer', 'AccessKeySerializer', 'OtpVerifySerializer', 'BearerTokenSerializer',
'MFAChallengeSerializer', 'LoginConfirmSettingSerializer', 'SSOTokenSerializer', 'MFAChallengeSerializer', 'LoginConfirmSettingSerializer', 'SSOTokenSerializer',
'ConnectionTokenSerializer',
] ]
@ -82,3 +83,33 @@ class SSOTokenSerializer(serializers.Serializer):
username = serializers.CharField(write_only=True) username = serializers.CharField(write_only=True)
login_url = serializers.CharField(read_only=True) login_url = serializers.CharField(read_only=True)
next = serializers.CharField(write_only=True, allow_blank=True, required=False, allow_null=True) next = serializers.CharField(write_only=True, allow_blank=True, required=False, allow_null=True)
class ConnectionTokenSerializer(serializers.Serializer):
user = serializers.CharField(max_length=128, required=True)
system_user = serializers.CharField(max_length=128, required=True)
asset = serializers.CharField(max_length=128, required=True)
@staticmethod
def validate_user(user_id):
from users.models import User
user = User.objects.filter(id=user_id).first()
if user is None:
raise serializers.ValidationError('user id not exist')
return user
@staticmethod
def validate_system_user(system_user_id):
from assets.models import SystemUser
system_user = SystemUser.objects.filter(id=system_user_id).first()
if system_user is None:
raise serializers.ValidationError('system_user id not exist')
return system_user
@staticmethod
def validate_asset(asset_id):
from assets.models import Asset
asset = Asset.objects.filter(id=asset_id).first()
if asset is None:
raise serializers.ValidationError('asset id not exist')
return asset

View File

@ -280,7 +280,8 @@ class Config(dict):
'SESSION_COOKIE_SECURE': False, 'SESSION_COOKIE_SECURE': False,
'CSRF_COOKIE_SECURE': False, 'CSRF_COOKIE_SECURE': False,
'REFERER_CHECK_ENABLED': False, 'REFERER_CHECK_ENABLED': False,
'SERVER_REPLAY_STORAGE': {} 'SERVER_REPLAY_STORAGE': {},
'CONNECTION_TOKEN_ENABLED': False,
} }
def compatible_auth_openid_of_key(self): def compatible_auth_openid_of_key(self):

View File

@ -115,3 +115,5 @@ DATETIME_DISPLAY_FORMAT = '%Y-%m-%d %H:%M:%S'
TICKETS_ENABLED = CONFIG.TICKETS_ENABLED TICKETS_ENABLED = CONFIG.TICKETS_ENABLED
REFERER_CHECK_ENABLED = CONFIG.REFERER_CHECK_ENABLED REFERER_CHECK_ENABLED = CONFIG.REFERER_CHECK_ENABLED
CONNECTION_TOKEN_ENABLED = CONFIG.CONNECTION_TOKEN_ENABLED