From aeff0ab5f338a7d5c842f4657ef584cb97dc47b0 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Mon, 11 Nov 2019 12:43:00 +0800
Subject: [PATCH] =?UTF-8?q?[Update]=20=E4=BF=AE=E6=94=B9public=20key=20?=
 =?UTF-8?q?=E6=A0=A1=E9=AA=8C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/authentication/backends/pubkey.py | 39 ++++++++++++++++++++++++++
 apps/authentication/utils.py           | 14 ++-------
 apps/jumpserver/settings.py            |  1 +
 apps/users/models/user.py              | 17 +++++++++++
 4 files changed, 60 insertions(+), 11 deletions(-)
 create mode 100644 apps/authentication/backends/pubkey.py

diff --git a/apps/authentication/backends/pubkey.py b/apps/authentication/backends/pubkey.py
new file mode 100644
index 000000000..db0ace648
--- /dev/null
+++ b/apps/authentication/backends/pubkey.py
@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+#
+from django.contrib.auth import get_user_model
+
+UserModel = get_user_model()
+
+__all__ = ['PublicKeyAuthBackend']
+
+
+class PublicKeyAuthBackend:
+    def authenticate(self, request, username=None, public_key=None, **kwargs):
+        if not public_key:
+            return None
+        if username is None:
+            username = kwargs.get(UserModel.USERNAME_FIELD)
+        try:
+            user = UserModel._default_manager.get_by_natural_key(username)
+        except UserModel.DoesNotExist:
+            return None
+        else:
+            if user.check_public_key(public_key) and \
+                    self.user_can_authenticate(user):
+                return user
+
+    @staticmethod
+    def user_can_authenticate(user):
+        """
+        Reject users with is_active=False. Custom user models that don't have
+        that attribute are allowed.
+        """
+        is_active = getattr(user, 'is_active', None)
+        return is_active or is_active is None
+
+    def get_user(self, user_id):
+        try:
+            user = UserModel._default_manager.get(pk=user_id)
+        except UserModel.DoesNotExist:
+            return None
+        return user if self.user_can_authenticate(user) else None
diff --git a/apps/authentication/utils.py b/apps/authentication/utils.py
index 1f190bef5..eb1649885 100644
--- a/apps/authentication/utils.py
+++ b/apps/authentication/utils.py
@@ -33,17 +33,9 @@ def check_user_valid(**kwargs):
     elif user.password_has_expired:
         return None, errors.reason_password_expired
 
-    if password:
-        user = authenticate(request, username=username, password=password)
+    if password or public_key:
+        user = authenticate(request, username=username,
+                            password=password, public_key=public_key)
         if user:
             return user, ''
-
-    if public_key and user.public_key:
-        public_key_saved = user.public_key.split()
-        if len(public_key_saved) == 1:
-            public_key_saved = public_key_saved[0]
-        else:
-            public_key_saved = public_key_saved[1]
-        if public_key == public_key_saved:
-            return user, ''
     return None, errors.reason_password_failed
diff --git a/apps/jumpserver/settings.py b/apps/jumpserver/settings.py
index 3eea3c78a..db2785d58 100644
--- a/apps/jumpserver/settings.py
+++ b/apps/jumpserver/settings.py
@@ -411,6 +411,7 @@ REST_FRAMEWORK = {
 
 AUTHENTICATION_BACKENDS = [
     'django.contrib.auth.backends.ModelBackend',
+    'authentication.backends.pubkey.PublicKeyAuthBackend',
 ]
 
 # Custom User Auth model
diff --git a/apps/users/models/user.py b/apps/users/models/user.py
index c987fbd91..bf8186174 100644
--- a/apps/users/models/user.py
+++ b/apps/users/models/user.py
@@ -120,6 +120,23 @@ class AuthMixin:
                 return s
         return False
 
+    @staticmethod
+    def get_public_key_body(key):
+        for i in key.split():
+            if len(i) > 256:
+                return i
+        return key
+
+    def check_public_key(self, key):
+        if not self.public_key:
+            return False
+        key = self.get_public_key_body(key)
+        key_saved = self.get_public_key_body(self.public_key)
+        if key == key_saved:
+            return True
+        else:
+            return False
+
 
 class RoleMixin:
     ROLE_ADMIN = 'Admin'