From 2a5c41dfaf0889f4bff0a6787a8ba79beb75fb90 Mon Sep 17 00:00:00 2001 From: wangruidong <940853815@qq.com> Date: Mon, 29 Jul 2024 19:37:50 +0800 Subject: [PATCH] feat: support configuring multiple SSH keys for users --- apps/authentication/api/__init__.py | 1 + apps/authentication/api/ssh_key.py | 21 ++++++++ apps/authentication/migrations/0003_sshkey.py | 51 +++++++++++++++++++ apps/authentication/models/__init__.py | 2 +- apps/authentication/models/ssh_key.py | 27 ++++++++++ apps/authentication/serializers/__init__.py | 1 + apps/authentication/serializers/ssh_key.py | 44 ++++++++++++++++ apps/authentication/urls/api_urls.py | 1 + apps/authentication/urls/view_urls.py | 1 - apps/users/api/profile.py | 15 +----- apps/users/models/user/__init__.py | 3 +- apps/users/models/user/_auth.py | 25 +++++---- apps/users/serializers/profile.py | 27 ---------- apps/users/urls/api_urls.py | 4 -- apps/users/views/profile/__init__.py | 1 - apps/users/views/profile/pubkey.py | 25 --------- 16 files changed, 162 insertions(+), 87 deletions(-) create mode 100644 apps/authentication/api/ssh_key.py create mode 100644 apps/authentication/migrations/0003_sshkey.py create mode 100644 apps/authentication/models/ssh_key.py create mode 100644 apps/authentication/serializers/ssh_key.py delete mode 100644 apps/users/views/profile/pubkey.py diff --git a/apps/authentication/api/__init__.py b/apps/authentication/api/__init__.py index 4cfc6cc4a..53d59c9a6 100644 --- a/apps/authentication/api/__init__.py +++ b/apps/authentication/api/__init__.py @@ -11,6 +11,7 @@ from .login_confirm import * from .mfa import * from .password import * from .session import * +from .ssh_key import * from .sso import * from .temp_token import * from .token import * diff --git a/apps/authentication/api/ssh_key.py b/apps/authentication/api/ssh_key.py new file mode 100644 index 000000000..aa4b34ad4 --- /dev/null +++ b/apps/authentication/api/ssh_key.py @@ -0,0 +1,21 @@ +from django.utils import timezone +from rest_framework.response import Response +from rest_framework.decorators import action + +from rbac.permissions import RBACPermission +from common.api import JMSModelViewSet +from common.permissions import IsValidUser +from ..serializers import SSHKeySerializer +from users.notifications import ResetPublicKeySuccessMsg + + +class SSHkeyViewSet(JMSModelViewSet): + serializer_class = SSHKeySerializer + permission_classes = [IsValidUser] + + def get_queryset(self): + return self.request.user.ssh_keys.all() + + def perform_update(self, serializer): + super().perform_update(serializer) + ResetPublicKeySuccessMsg(self.request.user, self.request).publish_async() diff --git a/apps/authentication/migrations/0003_sshkey.py b/apps/authentication/migrations/0003_sshkey.py new file mode 100644 index 000000000..22fbc8d34 --- /dev/null +++ b/apps/authentication/migrations/0003_sshkey.py @@ -0,0 +1,51 @@ +# Generated by Django 4.1.13 on 2024-07-29 02:25 + +import common.db.fields +import common.db.models +from django.conf import settings +from django.db import migrations, models +import uuid + + +def migrate_user_public_and_private_key(apps, schema_editor): + user_model = apps.get_model('users', 'User') + users = user_model.objects.all() + ssh_key_model = apps.get_model('authentication', 'SSHKey') + db_alias = schema_editor.connection.alias + for user in users: + if user.public_key: + ssh_key_model.objects.using(db_alias).create( + public_key=user.public_key, private_key=user.private_key, user=user + ) + + +class Migration(migrations.Migration): + dependencies = [ + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ('authentication', '0002_auto_20190729_1423'), + ] + + operations = [ + migrations.CreateModel( + name='SSHKey', + fields=[ + ('created_by', models.CharField(blank=True, max_length=128, null=True, verbose_name='Created by')), + ('updated_by', models.CharField(blank=True, max_length=128, null=True, verbose_name='Updated by')), + ('date_created', models.DateTimeField(auto_now_add=True, null=True, verbose_name='Date created')), + ('date_updated', models.DateTimeField(auto_now=True, verbose_name='Date updated')), + ('comment', models.TextField(blank=True, default='', verbose_name='Comment')), + ('id', models.UUIDField(default=uuid.uuid4, primary_key=True, serialize=False)), + ('name', models.CharField(max_length=128, verbose_name='Name')), + ('is_active', models.BooleanField(default=True, verbose_name='Active')), + ('private_key', common.db.fields.EncryptTextField(blank=True, null=True, verbose_name='Private key')), + ('public_key', common.db.fields.EncryptTextField(blank=True, null=True, verbose_name='Public key')), + ('date_last_used', models.DateTimeField(blank=True, null=True, verbose_name='Date last used')), + ('user', models.ForeignKey(db_constraint=False, on_delete=common.db.models.CASCADE_SIGNAL_SKIP, + related_name='ssh_keys', to=settings.AUTH_USER_MODEL, verbose_name='User')), + ], + options={ + 'verbose_name': 'SSH key', + }, + ), + migrations.RunPython(migrate_user_public_and_private_key) + ] diff --git a/apps/authentication/models/__init__.py b/apps/authentication/models/__init__.py index 23f8838ec..c11b5718f 100644 --- a/apps/authentication/models/__init__.py +++ b/apps/authentication/models/__init__.py @@ -1,7 +1,7 @@ from .access_key import * from .connection_token import * from .private_token import * +from .ssh_key import * from .sso_token import * from .temp_token import * from ..backends.passkey.models import * - diff --git a/apps/authentication/models/ssh_key.py b/apps/authentication/models/ssh_key.py new file mode 100644 index 000000000..bf2168ffc --- /dev/null +++ b/apps/authentication/models/ssh_key.py @@ -0,0 +1,27 @@ +import sshpubkeys + +from django.db import models +from django.utils.translation import gettext_lazy as _ + +from common.db.models import JMSBaseModel, CASCADE_SIGNAL_SKIP +from users.models import AuthMixin +from common.db import fields + + +class SSHKey(JMSBaseModel, AuthMixin): + name = models.CharField(max_length=128, verbose_name=_("Name")) + is_active = models.BooleanField(default=True, verbose_name=_('Active')) + private_key = fields.EncryptTextField( + blank=True, null=True, verbose_name=_("Private key") + ) + public_key = fields.EncryptTextField( + blank=True, null=True, verbose_name=_("Public key") + ) + date_last_used = models.DateTimeField(null=True, blank=True, verbose_name=_('Date last used')) + user = models.ForeignKey( + 'users.User', on_delete=CASCADE_SIGNAL_SKIP, verbose_name=_('User'), db_constraint=False, + related_name='ssh_keys' + ) + + class Meta: + verbose_name = _('SSH key') diff --git a/apps/authentication/serializers/__init__.py b/apps/authentication/serializers/__init__.py index d6e1671cf..9614b3efb 100644 --- a/apps/authentication/serializers/__init__.py +++ b/apps/authentication/serializers/__init__.py @@ -2,4 +2,5 @@ from .confirm import * from .connect_token_secret import * from .connection_token import * from .password_mfa import * +from .ssh_key import * from .token import * diff --git a/apps/authentication/serializers/ssh_key.py b/apps/authentication/serializers/ssh_key.py new file mode 100644 index 000000000..4d06a47c3 --- /dev/null +++ b/apps/authentication/serializers/ssh_key.py @@ -0,0 +1,44 @@ +# -*- coding: utf-8 -*- +# +from django.utils import timezone +from django.utils.translation import gettext_lazy as _ +from rest_framework import serializers + +from common.serializers.fields import ReadableHiddenField +from ..models import SSHKey +from common.utils import validate_ssh_public_key + +__all__ = ['SSHKeySerializer'] + + +class SSHKeySerializer(serializers.ModelSerializer): + user = ReadableHiddenField(default=serializers.CurrentUserDefault()) + public_key_comment = serializers.CharField( + source='get_public_key_comment', required=False, read_only=True, max_length=128 + ) + public_key_hash_md5 = serializers.CharField( + source='get_public_key_hash_md5', required=False, read_only=True, max_length=128 + ) + + class Meta: + model = SSHKey + fields_mini = ['name'] + fields_small = fields_mini + [ + 'public_key', 'is_active', + ] + read_only_fields = [ + 'id', 'user', 'public_key_comment', 'public_key_hash_md5', + 'date_last_used', 'date_created', 'date_updated' + ] + fields = fields_small + read_only_fields + + def to_representation(self, instance): + data = super().to_representation(instance) + data.pop('public_key', None) + return data + + @staticmethod + def validate_public_key(value): + if not validate_ssh_public_key(value): + raise serializers.ValidationError(_('Not a valid ssh public key')) + return value diff --git a/apps/authentication/urls/api_urls.py b/apps/authentication/urls/api_urls.py index 62f52ac9f..8f74f2907 100644 --- a/apps/authentication/urls/api_urls.py +++ b/apps/authentication/urls/api_urls.py @@ -14,6 +14,7 @@ router.register('temp-tokens', api.TempTokenViewSet, 'temp-token') router.register('connection-token', api.ConnectionTokenViewSet, 'connection-token') router.register('super-connection-token', api.SuperConnectionTokenViewSet, 'super-connection-token') router.register('confirm', api.UserConfirmationViewSet, 'confirm') +router.register('ssh-key', api.SSHkeyViewSet, 'ssh-key') urlpatterns = [ path('/qr/unbind/', api.QRUnBindForUserApi.as_view(), name='qr-unbind'), diff --git a/apps/authentication/urls/view_urls.py b/apps/authentication/urls/view_urls.py index 94c15047f..3df50260e 100644 --- a/apps/authentication/urls/view_urls.py +++ b/apps/authentication/urls/view_urls.py @@ -62,7 +62,6 @@ urlpatterns = [ path('slack/qr/login/callback/', views.SlackQRLoginCallbackView.as_view(), name='slack-qr-login-callback'), # Profile - path('profile/pubkey/generate/', users_view.UserPublicKeyGenerateView.as_view(), name='user-pubkey-generate'), path('profile/mfa/', users_view.MFASettingView.as_view(), name='user-mfa-setting'), # OTP Setting diff --git a/apps/users/api/profile.py b/apps/users/api/profile.py index 974fc136c..7e37da52c 100644 --- a/apps/users/api/profile.py +++ b/apps/users/api/profile.py @@ -18,8 +18,7 @@ from ..models import User __all__ = [ 'UserResetPasswordApi', 'UserResetPKApi', - 'UserProfileApi', 'UserPasswordApi', - 'UserPublicKeyApi' + 'UserProfileApi', 'UserPasswordApi' ] @@ -79,15 +78,3 @@ class UserPasswordApi(generics.RetrieveUpdateAPIView): resp = super().update(request, *args, **kwargs) ResetPasswordSuccessMsg(self.request.user, request).publish_async() return resp - - -class UserPublicKeyApi(generics.RetrieveUpdateAPIView): - permission_classes = (IsAuthenticated,) - serializer_class = serializers.UserUpdatePublicKeySerializer - - def get_object(self): - return self.request.user - - def perform_update(self, serializer): - super().perform_update(serializer) - ResetPublicKeySuccessMsg(self.get_object(), self.request).publish_async() diff --git a/apps/users/models/user/__init__.py b/apps/users/models/user/__init__.py index 7ab2ea271..725377ba1 100644 --- a/apps/users/models/user/__init__.py +++ b/apps/users/models/user/__init__.py @@ -28,7 +28,8 @@ logger = get_logger(__file__) __all__ = [ "User", "UserPasswordHistory", - "MFAMixin" + "MFAMixin", + "AuthMixin" ] diff --git a/apps/users/models/user/_auth.py b/apps/users/models/user/_auth.py index 7ea67d1d4..4b014b0e5 100644 --- a/apps/users/models/user/_auth.py +++ b/apps/users/models/user/_auth.py @@ -21,7 +21,6 @@ from users.signals import post_user_change_password logger = get_logger(__file__) - __all__ = ['MFAMixin', 'AuthMixin'] @@ -134,12 +133,6 @@ class AuthMixin: post_user_change_password.send(self.__class__, user=self) super().set_password(raw_password) # noqa - def set_public_key(self, public_key): - if self.can_update_ssh_key(): - self.public_key = public_key - self.save() - post_user_change_password.send(self.__class__, user=self) - def can_update_password(self): return self.is_local @@ -167,7 +160,7 @@ class AuthMixin: Check if the user's ssh public key is valid. This function is used in base.html. """ - if self.public_key: + if self.user_ssh_keys: return True return False @@ -233,14 +226,21 @@ class AuthMixin: except Exception as e: return "" + @property + def user_ssh_keys(self): + return self.ssh_keys.filter(is_active=True).all() + def check_public_key(self, key): - if not self.public_key: - return False key_md5 = self.get_public_key_md5(key) if not key_md5: return False - self_key_md5 = self.get_public_key_md5(self.public_key) - return key_md5 == self_key_md5 + for ssh_key in self.user_ssh_keys: + self_key_md5 = self.get_public_key_md5(ssh_key.public_key) + if key_md5 == self_key_md5: + ssh_key.date_last_used = timezone.now() + ssh_key.save(update_fields=['date_last_used']) + return True + return False def cache_login_password_if_need(self, password): from common.utils import signer @@ -270,4 +270,3 @@ class AuthMixin: return "" password = signer.unsign(secret) return password - diff --git a/apps/users/serializers/profile.py b/apps/users/serializers/profile.py index 852dd927c..5bbf5598a 100644 --- a/apps/users/serializers/profile.py +++ b/apps/users/serializers/profile.py @@ -55,33 +55,6 @@ class UserUpdatePasswordSerializer(serializers.ModelSerializer): return instance -class UserUpdatePublicKeySerializer(serializers.ModelSerializer): - public_key_comment = serializers.CharField( - source='get_public_key_comment', required=False, read_only=True, max_length=128 - ) - public_key_hash_md5 = serializers.CharField( - source='get_public_key_hash_md5', required=False, read_only=True, max_length=128 - ) - - class Meta: - model = User - fields = ['public_key_comment', 'public_key_hash_md5', 'public_key'] - extra_kwargs = { - 'public_key': {'required': True, 'write_only': True, 'max_length': 2048} - } - - @staticmethod - def validate_public_key(value): - if not validate_ssh_public_key(value): - raise serializers.ValidationError(_('Not a valid ssh public key')) - return value - - def update(self, instance, validated_data): - new_public_key = self.validated_data.get('public_key') - instance.set_public_key(new_public_key) - return instance - - class UserRoleSerializer(serializers.Serializer): name = serializers.CharField(max_length=24) display = serializers.CharField(max_length=64) diff --git a/apps/users/urls/api_urls.py b/apps/users/urls/api_urls.py index fa44e67e4..1c40f234b 100644 --- a/apps/users/urls/api_urls.py +++ b/apps/users/urls/api_urls.py @@ -18,11 +18,9 @@ router.register(r'users-groups-relations', api.UserUserGroupRelationViewSet, 'us router.register(r'service-account-registrations', api.ServiceAccountRegistrationViewSet, 'service-account-registration') router.register(r'connection-token', auth_api.ConnectionTokenViewSet, 'connection-token') - urlpatterns = [ path('profile/', api.UserProfileApi.as_view(), name='user-profile'), path('profile/password/', api.UserPasswordApi.as_view(), name='user-password'), - path('profile/public-key/', api.UserPublicKeyApi.as_view(), name='user-public-key'), path('profile/mfa/reset/', api.UserResetMFAApi.as_view(), name='my-mfa-reset'), path('preference/', api.PreferenceApi.as_view(), name='preference'), path('users//mfa/reset/', api.UserResetMFAApi.as_view(), name='user-reset-mfa'), @@ -32,5 +30,3 @@ urlpatterns = [ path('users//unblock/', api.UserUnblockPKApi.as_view(), name='user-unblock'), ] urlpatterns += router.urls - - diff --git a/apps/users/views/profile/__init__.py b/apps/users/views/profile/__init__.py index 5abff8d9f..35e2c7948 100644 --- a/apps/users/views/profile/__init__.py +++ b/apps/users/views/profile/__init__.py @@ -4,4 +4,3 @@ from .password import * from .mfa import * from .otp import * from .reset import * -from .pubkey import * diff --git a/apps/users/views/profile/pubkey.py b/apps/users/views/profile/pubkey.py deleted file mode 100644 index f8a038e50..000000000 --- a/apps/users/views/profile/pubkey.py +++ /dev/null @@ -1,25 +0,0 @@ -# ~*~ coding: utf-8 ~*~ - -from django.http import HttpResponse -from django.views import View - -from common.utils import get_logger, ssh_key_gen -from common.permissions import IsValidUser -from common.views.mixins import PermissionsMixin - -__all__ = ['UserPublicKeyGenerateView'] - -logger = get_logger(__name__) - - -class UserPublicKeyGenerateView(PermissionsMixin, View): - permission_classes = [IsValidUser] - - def get(self, request, *args, **kwargs): - username = request.user.username - private, public = ssh_key_gen(username=username, hostname='jumpserver') - request.user.set_public_key(public) - response = HttpResponse(private, content_type='text/plain') - filename = "{0}-jumpserver.pem".format(username) - response['Content-Disposition'] = 'attachment; filename={}'.format(filename) - return response