feat: support configuring multiple SSH keys for users

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 *

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()

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)
+    ]

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 *
-

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')

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 *

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

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('<str:backend>/qr/unbind/', api.QRUnBindForUserApi.as_view(), name='qr-unbind'),

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

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()

apps/users/models/user/__init__.py

@@ -28,7 +28,8 @@ logger = get_logger(__file__)
 __all__ = [
     "User",
     "UserPasswordHistory",
-    "MFAMixin"
+    "MFAMixin",
+    "AuthMixin"
 ]
 
 

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
-

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)

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/<uuid:pk>/mfa/reset/', api.UserResetMFAApi.as_view(), name='user-reset-mfa'),
@@ -32,5 +30,3 @@ urlpatterns = [
     path('users/<uuid:pk>/unblock/', api.UserUnblockPKApi.as_view(), name='user-unblock'),
 ]
 urlpatterns += router.urls
-
-

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 *

apps/users/views/profile/pubkey.py

@@ -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