Refactor out common functionality into separate function.

3 files changed, 59 insertions, 67 deletions

diff --git a/uploader/oauth2/tokens.py b/uploader/oauth2/tokens.py
new file mode 100644
index 0000000..eb650f6
--- /dev/null
+++ b/uploader/oauth2/tokens.py
@@ -0,0 +1,47 @@
+"""Utilities for dealing with tokens."""
+import uuid
+from typing import Union
+from urllib.parse import urljoin
+from datetime import datetime, timedelta
+
+from authlib.jose import jwt
+from flask import current_app as app
+
+from uploader import monadic_requests as mrequests
+
+from . import jwks
+from .client import (SCOPE, authserver_uri, oauth2_clientid)
+
+
+def request_token(token_uri: str, user_id: Union[uuid.UUID, str], **kwargs):
+    """Request token from the auth server."""
+    issued = datetime.now()
+    jwtkey = jwks.newest_jwk_with_rotation(
+        jwks.jwks_directory(app, "UPLOADER_SECRETS"),
+        int(app.config["JWKS_ROTATION_AGE_DAYS"]))
+    _mins2expiry = kwargs.get("minutes_to_expiry", 5)
+    return mrequests.post(
+        token_uri,
+        json={
+            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+            "scope": kwargs.get("scope", SCOPE),
+            "assertion": jwt.encode(
+                header={
+                    "alg": "RS256",
+                    "typ": "JWT",
+                    "kid": jwtkey.as_dict()["kid"]
+                },
+                payload={
+                    "iss": str(oauth2_clientid()),
+                    "sub": str(user_id),
+                    "aud": urljoin(authserver_uri(), "auth/token"),
+                    "exp": (issued + timedelta(minutes=_mins2expiry)).timestamp(),
+                    "nbf": int(issued.timestamp()),
+                    "iat": int(issued.timestamp()),
+                    "jti": str(uuid.uuid4())
+                },
+                key=jwtkey).decode("utf8"),
+            "client_id": oauth2_clientid(),
+            **kwargs.get("extra_params", {})
+        }
+    )
diff --git a/uploader/oauth2/views.py b/uploader/oauth2/views.py
index db4ef61..1ee4257 100644
--- a/uploader/oauth2/views.py
+++ b/uploader/oauth2/views.py
@@ -1,9 +1,6 @@
 """Views for OAuth2 related functionality."""
-import uuid
-from datetime import datetime, timedelta
 from urllib.parse import urljoin, urlparse, urlunparse
 
-from authlib.jose import jwt
 from flask import (
     flash,
     jsonify,
@@ -18,9 +15,8 @@ from uploader import monadic_requests as mrequests
 from uploader.monadic_requests import make_error_handler
 
 from . import jwks
+from .tokens import request_token
 from .client import (
-    SCOPE,
-    oauth2_get,
     user_logged_in,
     authserver_uri,
     oauth2_clientid,
@@ -33,8 +29,8 @@ oauth2 = Blueprint("oauth2", __name__)
 @oauth2.route("/code")
 def authorisation_code():
     """Receive authorisation code from auth server and use it to get token."""
-    def __process_error__(resp_or_exception):
-        app.logger.debug("ERROR: (%s)", resp_or_exception)
+    def __process_error__(error_response):
+        app.logger.debug("ERROR: (%s)", error_response.content)
         flash("There was an error retrieving the authorisation token.",
               "alert alert-danger")
         return redirect("/")
@@ -60,36 +56,14 @@ def authorisation_code():
         return redirect("/")
 
     baseurl = urlparse(request.base_url, scheme=request.scheme)
-    issued = datetime.now()
-    jwtkey = jwks.newest_jwk_with_rotation(
-        jwks.jwks_directory(app, "UPLOADER_SECRETS"),
-        int(app.config["JWKS_ROTATION_AGE_DAYS"]))
-    return mrequests.post(
-        urljoin(authserver_uri(), "auth/token"),
-        json={
-            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+    return request_token(
+        token_uri=urljoin(authserver_uri(), "auth/token"),
+        user_id=request.args["user_id"],
+        extra_params={
             "code": code,
-            "scope": SCOPE,
             "redirect_uri": urljoin(
                 urlunparse(baseurl),
                 url_for("oauth2.authorisation_code")),
-            "assertion": jwt.encode(
-                header={
-                    "alg": "RS256",
-                    "typ": "JWT",
-                    "kid": jwtkey.as_dict()["kid"]
-                },
-                payload={
-                    "iss": str(oauth2_clientid()),
-                    "sub": request.args["user_id"],
-                    "aud": urljoin(authserver_uri(),"auth/token"),
-                    "exp": (issued + timedelta(minutes=5)).timestamp(),
-                    "nbf": int(issued.timestamp()),
-                    "iat": int(issued.timestamp()),
-                    "jti": str(uuid.uuid4())
-                },
-                key=jwtkey).decode("utf8"),
-            "client_id": oauth2_clientid()
         }).either(__process_error__, __success__)
 
 @oauth2.route("/public-jwks")
diff --git a/uploader/phenotypes/views.py b/uploader/phenotypes/views.py
index 834a450..4a9df8b 100644
--- a/uploader/phenotypes/views.py
+++ b/uploader/phenotypes/views.py
@@ -13,7 +13,6 @@ from logging import INFO, ERROR, DEBUG, FATAL, CRITICAL, WARNING
 from urllib.parse import urljoin, urlparse, ParseResult, urlunparse, urlencode
 
 import datetime
-from datetime import timedelta
 
 from redis import Redis
 from pymonad.either import Left
@@ -25,9 +24,7 @@ from gn_libs import sqlite3
 from gn_libs import jobs as gnlibs_jobs
 from gn_libs.jobs.jobs import JobNotFound
 from gn_libs.mysqldb import database_connection
-from gn_libs import monadic_requests as mrequests
 
-from authlib.jose import jwt
 from flask import (flash,
                    request,
                    url_for,
@@ -47,8 +44,9 @@ from uploader import session
 from uploader.files import save_file#, fullpath
 from uploader.ui import make_template_renderer
 from uploader.oauth2.client import oauth2_post
+from uploader.oauth2.tokens import request_token
 from uploader.authorisation import require_login
-from uploader.oauth2 import jwks, client as oauth2client
+from uploader.oauth2 import client as oauth2client
 from uploader.route_utils import generic_select_population
 from uploader.datautils import safe_int, enumerate_sequence
 from uploader.species.models import all_species, species_by_id
@@ -676,37 +674,10 @@ def load_data_to_database(
             return redirect(url_for(
                 "background-jobs.job_status", job_id=load_job["job_id"]))
 
-        issued = datetime.datetime.now()
-        jwtkey = jwks.newest_jwk_with_rotation(
-            jwks.jwks_directory(app, "UPLOADER_SECRETS"),
-            int(app.config["JWKS_ROTATION_AGE_DAYS"]))
 
-        return mrequests.post(
-            urljoin(oauth2client.authserver_uri(), "auth/token"),
-            json={
-                "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
-                "scope": oauth2client.SCOPE,
-                "assertion": jwt.encode(
-                    header={
-                        "alg": "RS256",
-                        "typ": "JWT",
-                        "kid": jwtkey.as_dict()["kid"]
-                    },
-                    payload={
-                        "iss": str(oauth2client.oauth2_clientid()),
-                        "sub": str(session.user_details()["user_id"]),
-                        "aud": urljoin(oauth2client.authserver_uri(),
-                                       "auth/token"),
-                        # TODO: Update expiry time once fix is implemented in
-                        #       auth server.
-                        "exp": (issued + timedelta(minutes=5)).timestamp(),
-                        "nbf": int(issued.timestamp()),
-                        "iat": int(issued.timestamp()),
-                        "jti": str(uuid.uuid4())
-                    },
-                    key=jwtkey).decode("utf8"),
-                "client_id": oauth2client.oauth2_clientid()
-            }
+        return request_token(
+            token_uri=urljoin(oauth2client.authserver_uri(), "auth/token"),
+            user_id=session.user_details()["user_id"]
         ).then(
             lambda token: gnlibs_jobs.initialise_job(
                 conn,