From 1d26ccc7836909c25eaedc863defc78277303dcc Mon Sep 17 00:00:00 2001
From: lijianguo <lijianguo@zhumengyuan.com>
Date: Tue, 16 Mar 2021 16:08:39 +0800
Subject: [PATCH] feat(oauth2): add oauth2 support

---
 rootfs/api/authentication.py      | 48 ++++++++++++++++++++++++++++-
 rootfs/api/backend.py             | 31 +++++++++++++++++++
 rootfs/api/oauth.py               | 50 +++++++++++++++++++++++++++++++
 rootfs/api/serializers.py         | 23 ++++++++++++++
 rootfs/api/settings/production.py | 13 ++++++--
 rootfs/api/views.py               |  8 +++--
 rootfs/requirements.txt           |  1 +
 7 files changed, 168 insertions(+), 6 deletions(-)
 create mode 100644 rootfs/api/backend.py
 create mode 100644 rootfs/api/oauth.py

diff --git a/rootfs/api/authentication.py b/rootfs/api/authentication.py
index d1c72cf2d..36cb9a00c 100644
--- a/rootfs/api/authentication.py
+++ b/rootfs/api/authentication.py
@@ -1,8 +1,14 @@
 import logging
+from django.conf import settings
 from django.contrib.auth.models import AnonymousUser
+from django.core.cache import cache
+from django.utils.translation import gettext_lazy as _
 from rest_framework import authentication
-from rest_framework.authentication import TokenAuthentication
+from rest_framework.authentication import TokenAuthentication, \
+    get_authorization_header
+from rest_framework import exceptions
 
+from api.oauth import OAuthManager
 
 logger = logging.getLogger(__name__)
 
@@ -27,3 +33,43 @@ def authenticate(self, request):
         except Exception as e:
             logger.debug(e)
             return AnonymousUser(), None
+
+
+class DryccTokenAuthentication(TokenAuthentication):
+    def authenticate(self, request):
+        if 'manager' in request.META.get('HTTP_USER_AGENT', ''):
+            auth = get_authorization_header(request).split()
+
+            if not auth or auth[0].lower() != self.keyword.lower().encode():
+                return None
+
+            if len(auth) == 1:
+                msg = _('Invalid token header. No credentials provided.')
+                raise exceptions.AuthenticationFailed(msg)
+            elif len(auth) > 2:
+                msg = _('Invalid token header. Token string should not contain spaces.')  # noqa
+                raise exceptions.AuthenticationFailed(msg)
+
+            try:
+                token = auth[1].decode()
+            except UnicodeError:
+                msg = _('Invalid token header. Token string should not contain invalid characters.')  # noqa
+                raise exceptions.AuthenticationFailed(msg)
+            return self._check_oauth_token(token)
+        return super(DryccTokenAuthentication, self).authenticate(request)  # noqa
+
+    def _check_oauth_token(self, key):
+        user = cache.get(key)
+        if user:
+            return user, None
+        try:
+            user_info = OAuthManager().get_user_by_token(key)
+            if not user_info.get('email'):
+                user_info['email'] = OAuthManager().get_email_by_token(key)
+        except Exception as e:
+            logger.info(e)
+            raise exceptions.AuthenticationFailed(_('Verify token fail.'))
+        from api import serializers
+        user = serializers.UserSerializer.update_or_create(user_info)
+        cache.set(key, user, settings.OAUTH_USER_CACHE_TIME)
+        return user, None
diff --git a/rootfs/api/backend.py b/rootfs/api/backend.py
new file mode 100644
index 000000000..9490c5e97
--- /dev/null
+++ b/rootfs/api/backend.py
@@ -0,0 +1,31 @@
+from django.contrib.auth import get_user_model
+from django.core.exceptions import ObjectDoesNotExist
+
+from api import serializers
+from api.oauth import OAuthManager
+
+
+class DryccOauthBackend(object):
+
+    # The Django auth backend API
+    def authenticate(self, request, username=None, password=None, **kwargs):
+        if username is None:
+            return None
+
+        with OAuthManager() as client:
+            client.fetch_token(username, password)
+            user_info = client.get_user()
+            user_info['username'] = username
+            user_info['password'] = password
+            if not user_info.get('email'):
+                user_info['email'] = client.get_email()
+            user = serializers.UserSerializer.update_or_create(user_info)
+        return user
+
+    def get_user(self, user_id):
+        user = None
+        try:
+            user = get_user_model().objects.get(pk=user_id)
+        except ObjectDoesNotExist:
+            pass
+        return user
diff --git a/rootfs/api/oauth.py b/rootfs/api/oauth.py
new file mode 100644
index 000000000..34c2061ff
--- /dev/null
+++ b/rootfs/api/oauth.py
@@ -0,0 +1,50 @@
+from typing import Dict
+
+import requests
+from django.conf import settings
+from authlib.integrations.requests_client import OAuth2Session
+
+
+class OAuthManager(object):
+    def __enter__(self):
+        return self
+
+    def __init__(self):
+        self.client_id = settings.OAUTH_CLIENT_ID
+        self.client_secret = settings.OAUTH_CLIENT_SECRET
+        self.token_url = settings.OAUTH_ACCESS_TOKEN_URL
+        self.api_url = settings.OAUTH_ACCESS_API_URL
+        self.client = OAuth2Session(self.client_id, self.client_secret)
+
+    def fetch_token(self, username: str, password: str) -> Dict:
+        response = self.client.fetch_token(self.token_url,
+                                           username=username,
+                                           password=password)
+        return response
+
+    def get_user(self) -> Dict:
+        response = self.client.get(f'{self.api_url}/users')
+        result = response.json()
+        return result
+
+    def get_email(self) -> str:
+        response = self.client.get(f'{self.api_url}/users/emails')
+        result = response.json()
+        return result['email']
+
+    def get_user_by_token(self, token: str) -> Dict:
+        response = requests.get(f'{self.api_url}/users', headers={
+            'Authorization': f'Bearer {token}'
+        })
+        result = response.json()
+        return result
+
+    def get_email_by_token(self, token: str) -> Dict:
+        response = requests.get(f'{self.api_url}/users/emails', headers={
+            'Authorization': f'Bearer {token}'
+        })
+        result = response.json()
+        return result
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.client.close()
diff --git a/rootfs/api/serializers.py b/rootfs/api/serializers.py
index 710f23436..3c5877895 100644
--- a/rootfs/api/serializers.py
+++ b/rootfs/api/serializers.py
@@ -166,6 +166,29 @@ def create(self, validated_data):
         user.save()
         return user
 
+    @staticmethod
+    def update_or_create(data):
+        now = timezone.now()
+        user, created = User.objects.update_or_create(
+            username=data['username'],
+            defaults={
+                'email': data['email'],
+                'first_name': data.get('first_name', ''),
+                'last_name': data.get('first_name', ''),
+                'last_login': now})
+        if created:
+            user.date_joined = now
+            user.is_active = True
+        if data.get('password'):
+            user.set_password(data['password'])
+        elif created and not data.get('password'):
+            user.set_unusable_password()
+        # Make the first signup an admin / superuser
+        if not User.objects.filter(is_superuser=True).exists():
+            user.is_superuser = user.is_staff = True
+        user.save()
+        return user
+
 
 class AdminUserSerializer(serializers.ModelSerializer):
     """Serialize admin status for a User model."""
diff --git a/rootfs/api/settings/production.py b/rootfs/api/settings/production.py
index 83e33afdd..03e038f1f 100644
--- a/rootfs/api/settings/production.py
+++ b/rootfs/api/settings/production.py
@@ -155,7 +155,7 @@
         'rest_framework.permissions.IsAuthenticated',
     ),
     'DEFAULT_AUTHENTICATION_CLASSES': (
-        'rest_framework.authentication.TokenAuthentication',
+        'api.authentication.DryccTokenAuthentication',
     ),
     'DEFAULT_RENDERER_CLASSES': (
         'rest_framework.renderers.JSONRenderer',
@@ -434,7 +434,6 @@
 # check if we can register users with `drycc register`
 REGISTRATION_MODE = os.environ.get('REGISTRATION_MODE', 'enabled')
 
-
 DRYCC_DATABASE_URL = os.environ.get('DRYCC_DATABASE_URL', 'postgres://:@:5432/drycc')
 DATABASES = {
     'default': dj_database_url.config(default=DRYCC_DATABASE_URL, conn_max_age=600)
@@ -442,6 +441,15 @@
 
 APP_URL_REGEX = '[a-z0-9-]+'
 
+# Oauth settings
+OAUTH_ACCESS_TOKEN_URL = os.environ.get('OAUTH2_ACCESS_TOKEN_URL', '')
+OAUTH_ACCESS_API_URL = os.environ.get('OAUTH2_ACCESS_API_URL', '')
+OAUTH_CLIENT_ID = os.environ.get('OAUTH2_CLIENT_ID', '')
+OAUTH_CLIENT_SECRET = os.environ.get('OAUTH2_CLIENT_SECRET', '')
+OAUTH_CACHE_USER_TIME = os.environ.get('OAUTH_CACHE_USER_TIME', 30 * 60)
+if OAUTH_ACCESS_TOKEN_URL:
+    AUTHENTICATION_BACKENDS = ("api.backend.DryccOauthBackend",) + AUTHENTICATION_BACKENDS
+
 # LDAP settings taken from environment variables.
 LDAP_ENDPOINT = os.environ.get('LDAP_ENDPOINT', '')
 LDAP_BIND_DN = os.environ.get('LDAP_BIND_DN', '')
@@ -537,7 +545,6 @@
 DRYCC_INFLUXDB_ORG = os.environ.get('DRYCC_INFLUXDB_ORG', 'root')
 DRYCC_INFLUXDB_TOKEN = os.environ.get('DRYCC_INFLUXDB_TOKEN', 'root')
 
-
 # Workflow-manager Configuration Options
 WORKFLOW_MANAGER_URL = os.environ.get('DRYCC_WORKFLOW_MANAGER_URL', None)
 WORKFLOW_MANAGER_TOKEN = os.environ.get('DRYCC_WORKFLOW_MANAGER_TOKEN', None)
diff --git a/rootfs/api/views.py b/rootfs/api/views.py
index 98ed69a2f..84de70e31 100644
--- a/rootfs/api/views.py
+++ b/rootfs/api/views.py
@@ -22,7 +22,6 @@
 from api.models import AlreadyExists, ServiceUnavailable, DryccException, \
     UnprocessableEntity
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -64,6 +63,8 @@ class UserRegistrationViewSet(GenericViewSet,
     def create(self, request, *args, **kwargs):
         if settings.LDAP_ENDPOINT:
             raise DryccException("You cannot register user when ldap is enabled.")
+        if settings.OAUTH_ACCESS_TOKEN_URL:
+            raise DryccException("You cannot register user when oauth2 is enabled.")
         return super(UserRegistrationViewSet, self).create(request, *args, **kwargs)
 
 
@@ -84,6 +85,8 @@ def list(self, request, **kwargs):
     def destroy(self, request, **kwargs):
         if settings.LDAP_ENDPOINT:
             raise DryccException("You cannot destroy user when ldap is enabled.")
+        if settings.OAUTH_ACCESS_TOKEN_URL:
+            raise DryccException("You cannot destroy user when oauth2 is enabled.")
         calling_obj = self.get_object()
         target_obj = calling_obj
 
@@ -110,7 +113,8 @@ def passwd(self, request, **kwargs):
             raise DryccException("new_password is a required field")
         if settings.LDAP_ENDPOINT:
             raise DryccException("You cannot change password when ldap is enabled.")
-
+        if settings.OAUTH_ACCESS_TOKEN_URL:
+            raise DryccException("You cannot change user when oauth2 is enabled.")
         caller_obj = self.get_object()
         target_obj = self.get_object()
         if request.data.get('username'):
diff --git a/rootfs/requirements.txt b/rootfs/requirements.txt
index b91766e74..4208e1d8c 100644
--- a/rootfs/requirements.txt
+++ b/rootfs/requirements.txt
@@ -24,3 +24,4 @@ celery==5.0.2
 django_redis==4.12.1
 influxdb-client==1.13.0
 dj-database-url
+Authlib
