From 13e798b526a974dd301afbd790325ce25f1ee070 Mon Sep 17 00:00:00 2001
From: lijianguo <lijianguo@zhumengyuan.com>
Date: Thu, 11 Mar 2021 17:20:35 +0800
Subject: [PATCH] feat(oauth2): add oauth2 support

---
 rootfs/api/authentication.py      | 13 ++++++-
 rootfs/api/exceptions.py          |  5 +++
 rootfs/api/oauth2/__init__.py     | 61 +++++++++++++++++++++++++++++++
 rootfs/api/oauth2/manager.py      | 24 ++++++++++++
 rootfs/api/serializers.py         | 24 ++++++++++++
 rootfs/api/settings/production.py | 14 ++++++-
 rootfs/api/urls.py                | 11 +++++-
 rootfs/api/views.py               | 37 +++++++++++++++++--
 8 files changed, 181 insertions(+), 8 deletions(-)
 create mode 100644 rootfs/api/oauth2/__init__.py
 create mode 100644 rootfs/api/oauth2/manager.py

diff --git a/rootfs/api/authentication.py b/rootfs/api/authentication.py
index d1c72cf2d..c1cedc1e5 100644
--- a/rootfs/api/authentication.py
+++ b/rootfs/api/authentication.py
@@ -1,8 +1,9 @@
 import logging
+from django.core.cache import cache
 from django.contrib.auth.models import AnonymousUser
 from rest_framework import authentication
 from rest_framework.authentication import TokenAuthentication
-
+from rest_framework import exceptions
 
 logger = logging.getLogger(__name__)
 
@@ -27,3 +28,13 @@ def authenticate(self, request):
         except Exception as e:
             logger.debug(e)
             return AnonymousUser(), None
+
+
+class ExpiringTokenAuthentication(TokenAuthentication):
+    def authenticate_credentials(self, key):
+        model = self.get_model()
+        has_key = cache.get('token_' + key)
+        if not has_key:
+            model.objects.select_related('user').filter(key=key).delete()
+            raise exceptions.AuthenticationFailed('Token has expired, please login')
+        return super(ExpiringTokenAuthentication, self).authenticate_credentials(key)
diff --git a/rootfs/api/exceptions.py b/rootfs/api/exceptions.py
index 07cf20dec..8b6187847 100644
--- a/rootfs/api/exceptions.py
+++ b/rootfs/api/exceptions.py
@@ -53,3 +53,8 @@ def custom_exception_handler(exc, context):
         logging.exception(exc.__cause__, exc_info=exc)
 
     return response
+
+
+class Oauth2Exception(APIException):
+    """Exception class used for when the application's health check fails"""
+    pass
diff --git a/rootfs/api/oauth2/__init__.py b/rootfs/api/oauth2/__init__.py
new file mode 100644
index 000000000..d2d9e2962
--- /dev/null
+++ b/rootfs/api/oauth2/__init__.py
@@ -0,0 +1,61 @@
+import base64
+import logging
+import urllib
+import requests
+from django.conf import settings
+
+from api.exceptions import Oauth2Exception
+
+logger = logging.getLogger(__name__)
+
+
+class OAuthBase(object):
+    def __init__(self):
+        self.client_id = settings.OAUTH2_CLIENT_ID
+        self.client_key = settings.OAUTH2_CLIENT_SECRET
+        self.api_url = settings.OAUTH2_ACCESS_API_URL
+        self.access_token = None
+
+    def oauth_get(self, url, data):
+        request_url = '%s?%s' % (url, urllib.parse.urlencode(data))
+        headers = {
+            'Authorization': f'Bearer {self.access_token}'
+        }
+        try:
+            response = requests.get(request_url, headers=headers, timeout=5)
+        except requests.exceptions.ConnectionError as err:
+            message = "There was a problem retrieving data from " \
+                      "the oauth2 server. URL: {}, params: {}".format(url, data)
+            logger.error(message)
+            raise Oauth2Exception(message) from err
+        return response
+
+    def oauth_post(self, url, data):
+        token = base64.b64encode(
+            ('{}:{}').format(self.client_id, self.client_key).encode()).decode(
+            encoding='UTF-8')
+        headers = {
+            'Authorization': f'Basic {token}'
+        }
+        try:
+            response = requests.post(url, headers=headers, data=data, timeout=5)
+        except requests.exceptions.ConnectionError as err:
+            message = "There was a problem retrieving accecc_token from " \
+                      "the oauth2 server. URL: {}, params: {}".format(url, data)
+            logger.error(message)
+            raise Oauth2Exception(message) from err
+        return response
+
+    def get_auth_url(self):
+        pass
+
+    def get_user_info(self):
+        params = {'access_token': self.access_token}
+        response = self.oauth_get(f'{self.api_url}/users', params)
+        result = response.json()
+        return result
+
+    def get_email(self):
+        response = self.oauth_get(f'{self.api_url}/users/emails')
+        result = response.json()
+        return result['email']
diff --git a/rootfs/api/oauth2/manager.py b/rootfs/api/oauth2/manager.py
new file mode 100644
index 000000000..6b644c6eb
--- /dev/null
+++ b/rootfs/api/oauth2/manager.py
@@ -0,0 +1,24 @@
+from django.conf import settings
+from django.core.cache import cache
+
+from api.oauth2 import OAuthBase
+
+
+class OAuthManager(OAuthBase):
+
+    def get_password_token(self, username, password):
+        params = {
+            'grant_type': 'password',
+            'username': username,
+            'password': password,
+        }
+        response = self.oauth_post(settings.OAUTH2_ACCESS_TOKEN_URL, params)
+        result = response.json()
+        self.access_token = result['access_token']
+        cache.set('token_' + result['access_token'],
+                  True,
+                  int(result['expires_in'] - 60))
+        return result
+
+
+oauth_client = OAuthManager()
diff --git a/rootfs/api/serializers.py b/rootfs/api/serializers.py
index 710f23436..5eb7fadad 100644
--- a/rootfs/api/serializers.py
+++ b/rootfs/api/serializers.py
@@ -167,6 +167,30 @@ def create(self, validated_data):
         return user
 
 
+class AuthTokenSerializer(serializers.Serializer):
+    username = serializers.CharField(max_length=150, required=True)
+    password = serializers.CharField(required=True)
+
+    def update_or_create(self, 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', '')})
+        if created:
+            user.last_login = now,
+            user.date_joined = now
+            user.is_active = True
+        user.set_password(data['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..a2af3ce14 100644
--- a/rootfs/api/settings/production.py
+++ b/rootfs/api/settings/production.py
@@ -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)
@@ -537,7 +536,18 @@
 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)
+
+OAUTH2_ACCESS_TOKEN_URL = None
+if not LDAP_ENDPOINT:
+    OAUTH2_ACCESS_TOKEN_URL = os.environ.get('OAUTH2_ACCESS_TOKEN_URL', '')
+    OAUTH2_ACCESS_API_URL = os.environ.get('OAUTH2_ACCESS_API_URL', '')
+    OAUTH2_CLIENT_ID = os.environ.get('OAUTH2_CLIENT_ID', '')
+    OAUTH2_CLIENT_SECRET = os.environ.get('OAUTH2_CLIENT_SECRET', '')
+
+    if OAUTH2_ACCESS_TOKEN_URL:
+        REST_FRAMEWORK['DEFAULT_AUTHENTICATION_CLASSES'] = (
+            'api.authentication.ExpiringTokenAuthentication',
+        )
diff --git a/rootfs/api/urls.py b/rootfs/api/urls.py
index 5ed8ac65f..7c3aef449 100644
--- a/rootfs/api/urls.py
+++ b/rootfs/api/urls.py
@@ -127,8 +127,6 @@
         views.UserManagementViewSet.as_view({'post': 'passwd'})),
     url(r'^auth/whoami/?$',
         views.UserManagementViewSet.as_view({'get': 'list'})),
-    url(r'^auth/login/$',
-        views_obtain_auth_token),
     url(r'^auth/tokens/$',
         views.TokenManagementViewSet.as_view({'post': 'regenerate'})),
     url(r'^auth/tokens/(?P<username>[\w.@+-]+)/?$',
@@ -158,3 +156,12 @@
     url(r'^apps/(?P<id>{})/metrics/(?P<container_type>[a-z0-9]+(\-[a-z0-9]+)*)?$'.format(settings.APP_URL_REGEX),  # noqa
         views.MetricView.as_view({'get': 'status'})),
 ]
+
+if settings.OAUTH2_ACCESS_TOKEN_URL:
+    urlpatterns += [
+        url(r'auth/login/$', views.AuthToken.as_view())
+    ]
+else:
+    urlpatterns += [
+        url(r'^auth/login/$', views_obtain_auth_token),
+    ]
diff --git a/rootfs/api/views.py b/rootfs/api/views.py
index 98ed69a2f..81389976e 100644
--- a/rootfs/api/views.py
+++ b/rootfs/api/views.py
@@ -1,11 +1,12 @@
 """
 RESTful view classes for presenting Drycc API objects.
 """
+import json
 import logging
 from copy import deepcopy
-from django.http import Http404, HttpResponse
 from django.conf import settings
 from django.contrib.auth.models import User
+from django.http import Http404, HttpResponse
 from django.shortcuts import get_object_or_404
 from guardian.shortcuts import assign_perm, get_objects_for_user, \
     get_users_with_perms, remove_perm
@@ -15,13 +16,15 @@
 from rest_framework.exceptions import PermissionDenied, NotFound, AuthenticationFailed
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response
+from rest_framework.views import APIView
 from rest_framework.viewsets import GenericViewSet
 from rest_framework.authtoken.models import Token
 
 from api import influxdb, authentication, models, permissions, serializers, viewsets
 from api.models import AlreadyExists, ServiceUnavailable, DryccException, \
     UnprocessableEntity
-
+from api.oauth2.manager import oauth_client
+from api.permissions import IsAnonymous
 
 logger = logging.getLogger(__name__)
 
@@ -64,6 +67,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.OAUTH2_ACCESS_TOKEN_URL:
+            raise DryccException("You cannot register user when oauth2 is enabled.")
         return super(UserRegistrationViewSet, self).create(request, *args, **kwargs)
 
 
@@ -84,6 +89,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.OAUTH2_ACCESS_TOKEN_URL:
+            raise DryccException("You cannot destroy user when oauth2 is enabled.")
         calling_obj = self.get_object()
         target_obj = calling_obj
 
@@ -110,7 +117,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.OAUTH2_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'):
@@ -992,3 +1000,26 @@ def status(self, request, **kwargs):
             "networks": self._get_networks(
                 app_id, container_type, start, stop, every)
         })
+
+
+class AuthToken(APIView):
+    permission_classes = [IsAnonymous]
+
+    def post(self, request):
+        serializer = serializers.AuthTokenSerializer(data=self.request.data)
+        serializer.is_valid(raise_exception=True)
+        username = serializer.validated_data.get('username')
+        password = serializer.validated_data.get('password')
+        token_response = oauth_client.get_password_token(username, password)
+        user_info = oauth_client.get_user_info()
+        user_info['username'] = username
+        user_info['password'] = password
+        if not user_info.get('email'):
+            user_info['email'] = oauth_client.get_email()
+        user = serializer.update_or_create(user_info)
+        Token.objects.filter(user=user).delete()
+        token = Token.objects.create(
+            key=token_response['access_token'],
+            user=user
+        )
+        return HttpResponse(json.dumps({"token": token.key}))
