Ver Fonte

[jwt] Implement custom JWT Authentication and update UTs

- Decode JWT token
- Find or create user from payload data
- Update tests for existing user, new user, authentication scenarios

TODO
- Token persistance
- Secret key handling and decide algorithm
Harshg999 há 4 anos atrás
pai
commit
f0fa89f0ea

+ 18 - 19
desktop/core/src/desktop/auth/api_authentications.py

@@ -16,9 +16,11 @@
 # limitations under the License.
 
 import logging
+import jwt
 
-from rest_framework import authentication
-from rest_framework import exceptions
+from rest_framework import authentication, exceptions
+
+from desktop.auth.backend import find_or_create_user, ensure_has_a_group, rewrite_user
 
 from useradmin.models import User
 
@@ -33,28 +35,25 @@ class JwtAuthentication(authentication.BaseAuthentication):
       LOG.debug('JwtAuthentication: no authorization header')
       return None
 
-    bearer = authorization_header[len('Bearer '):]
-    if not bearer:
+    access_token = authorization_header[len('Bearer '):]
+    if not access_token:
       LOG.debug('JwtAuthentication: no Bearer value')
       return None
 
-    LOG.debug('JwtAuthentication: got token %s' % bearer)
-
-    # Decode token via jwt module
-    # check expiration, get userId, handle errors
-
-    # token = '...'
+    LOG.debug('JwtAuthentication: got access token %s' % access_token)
 
-    # cf. below similar to backend.py
-    # user = find_or_create_user(
-    #   username,
-    #   password,
-    #   is_superuser=False
-    # )
-    user = User.objects.get(username='test')
+    try:
+      payload = jwt.decode(access_token, 'secret', algorithms=["HS256"])
+    except jwt.DecodeError:
+      raise exceptions.AuthenticationFailed('JwtAuthentication: Invalid token')
+    except jwt.ExpiredSignatureError:
+      raise exceptions.AuthenticationFailed('JwtAuthentication: Token expired')
+    except Exception as e:
+      raise exceptions.AuthenticationFailed(e)
 
-    # ensure_has_a_group(user)
-    # user = rewrite_user(user)
+    user = find_or_create_user(payload['username'], is_superuser=False)
+    ensure_has_a_group(user)
+    user = rewrite_user(user)
 
     # We should persist the token (to reuse for communicating with external services as the user, e.g. Impala)
     # either via a new DB field (might just be cleaner, even if requires a DB migration) or in the json blob.

+ 45 - 11
desktop/core/src/desktop/auth/api_authentications_tests.py

@@ -17,13 +17,17 @@
 
 import sys
 
-from nose.tools import assert_equal, assert_true
+from nose.tools import assert_equal, assert_true, assert_false, assert_raises
 
 from desktop.auth.backend import rewrite_user
 from desktop.auth.api_authentications import JwtAuthentication
 from desktop.lib.django_test_util import make_logged_in_client
+
+from rest_framework import exceptions
+
 from useradmin.models import User
 
+
 if sys.version_info[0] > 2:
   from unittest.mock import patch, Mock, MagicMock
 else:
@@ -36,16 +40,46 @@ class TestJwtAuthentication():
     self.user = rewrite_user(User.objects.get(username="test"))
 
   def test_authenticate(self):
-    HEADERS = {
-      "Authorization":
-      "Bearer eyJhbGciOiJSUzI1NiJ9.eyJhdWQiOlsid29ya2xvYWQtYXBwIiwicmFuZ2VyIl0sImV4cCI6MTYyNjI1Njg5MywiaWF0IjoxNjI2MjU2NTkzLCJpc3MiOiJDbG91ZGVyYTEiLCJqdGkiOiJpZDEiLCJzdWIiOiJ0ZXN0LXN1YmplY3QiLCJ1c2VyIjoidGVzdF91c2VyIn0.jvyVDxbWTAik0jbdUcIc9ZANNrJZUCWH-Pg7FloRhg0ZYAETd_AO3p5v_ppoMmVcPD2xBSrngA5J3_A_zPBvQ_hdDlpb0_-mCCJfGhC5tju4bI9EE9Akdn2FrrsqrvQQ8cPyGsIlvoIxrK1De4f74MmUaxfN7Hrrcue1PTY4u4IB9cWQqV9vIcX99Od5PUaNekLIee-I8gweqvfGEEsW7qWUM63nh59_TOB3LLq-YcEuaX1h_oiTATeCssjk_ee9RrJGLNyKmC0WJ4UrEWn8a_T3bwCy8CMe0zV5PSuuvPHy0FvnTo2il5SDjGimxKcbpgNiJdfblslu6i35DlfiWg"
-    }
-    request = MagicMock(HEADERS=HEADERS)
+    with patch('desktop.auth.api_authentications.jwt.decode') as jwt_decode:
+
+      HEADERS = {
+        'HTTP_AUTHORIZATION':
+        "Bearer eyJhbGciOiJSUzI1NiJ9.eyJhdWQiOlsid29ya2xvYWQtYXBwIiwicmFuZ2VyIl0sImV4cCI6MTYyNjI1Njg5MywiaWF0IjoxNjI2MjU2NTkzLCJpc3MiOiJDbG91ZGVyYTEiLCJqdGkiOiJpZDEiLCJzdWIiOiJ0ZXN0LXN1YmplY3QiLCJ1c2VyIjoidGVzdF91c2VyIn0.jvyVDxbWTAik0jbdUcIc9ZANNrJZUCWH-Pg7FloRhg0ZYAETd_AO3p5v_ppoMmVcPD2xBSrngA5J3_A_zPBvQ_hdDlpb0_-mCCJfGhC5tju4bI9EE9Akdn2FrrsqrvQQ8cPyGsIlvoIxrK1De4f74MmUaxfN7Hrrcue1PTY4u4IB9cWQqV9vIcX99Od5PUaNekLIee-I8gweqvfGEEsW7qWUM63nh59_TOB3LLq-YcEuaX1h_oiTATeCssjk_ee9RrJGLNyKmC0WJ4UrEWn8a_T3bwCy8CMe0zV5PSuuvPHy0FvnTo2il5SDjGimxKcbpgNiJdfblslu6i35DlfiWg"
+      }
+      
+      request = MagicMock(HEADERS=HEADERS)
+
+      # Existing user
+      jwt_decode.return_value = {
+        "username": "test"
+      }
+      assert_equal(1, User.objects.count())
+
+      user, token = JwtAuthentication().authenticate(request=request)
+
+      assert_equal(user, self.user)
+      assert_true(user.is_authenticated)
+      assert_equal(1, User.objects.count())
+
+      # New user
+      jwt_decode.return_value = {
+        "username": "test_new_user"
+      }
+      assert_equal(1, User.objects.count())
+
+      user, token = JwtAuthentication().authenticate(request=request)
 
-    user, token = JwtAuthentication().authenticate(request=request)
+      assert_equal(User.objects.get(username="test_new_user"), user)
+      assert_true(user.is_authenticated)
+      assert_equal(2, User.objects.count())
 
-    assert_equal(user, self.user)
-    assert_true(user.is_authenticated)
+      # Invalid token
+      jwt_decode.side_effect = exceptions.AuthenticationFailed('JwtAuthentication: Invalid token')
+      assert_raises(exceptions.AuthenticationFailed, JwtAuthentication().authenticate, request)
 
-  # TODO:
-  # check user.token, test new user, existing user, failure to authenticate
+      # Expired token
+      jwt_decode.side_effect = exceptions.AuthenticationFailed('JwtAuthentication: Token expired')
+      assert_raises(exceptions.AuthenticationFailed, JwtAuthentication().authenticate, request)
+    
+      # TODO:
+      # check user.token