Repository: CERT-Polska/mwdb-core (language: python)

Issue:
Proposed changes in JWT tokens used for authorization
**Feature Category**

- [x] Correctness
- [ ] User Interface / User Experience
- [ ] Performance
- [ ] Other (please explain)

**Describe the problem**

After short review, I found that there are few issues with currently used JWT tokens:

1. Currently session token can be reused as set password token, because JWT doesn't explicitly limit the scope. Fortunately, that's the only case I found. The required JWT payload fields usually differ, but in set password token, required part is a subset of fields set in session token payload.

    JWTs are used in various ways in authorization process, so it would be nice to define a `scope` field that declares what action is authorized by JWT. Possible values are:
    - `session`
    - `api_key`
    - `set_password`
    - `download_file`
    Scope should be checked before action is authorized.

2. The payload part is non-standard and doesn't contain necessary parts that are usually required in JWT. RFC 7519 declares these registered claims as optional, but it would be nice to follow the specification

Good example is JWT used as session token:

<img width="213" alt="image" src="https://user-images.githubusercontent.com/8720367/152395228-ca83f170-5648-44a3-ae15-f377f1937c9e.png">

Standard claims are:
- `exp` (expiration time) for tokens with expiration time that is usually a part of **payload** not **header**
- `iat` (issued at) also is usually a part of **payload** not **header**
- `iss` (issuer) / `aud` (audience) that declares the issuer entity and the audience for token (e.g. "https://mwdb.cert.pl") so the token cannot be reused to authorize action in other service. In case of our tokens: issuer and audience are the same, so we can set only one the issuer. 
   Usually different MWDB Core instances differ in secret key that is random generated so it's not a huge security issue now, but it's nice to have well-defined issuer in JWT.
- `sub` (subject) that should be login of the user that is authorized for the operation.

See also https://datatracker.ietf.org/doc/html/rfc7519#section-4.1

3. `itsdangerous`  JOSE functionality is deprecated and documentation recommends to use `authlib` for JWT generation (https://itsdangerous.palletsprojects.com/en/2.0.x/jws/). We already use it for OIDC integration, so it should not be a problem to move to the another library!

**Describe the solution you'd like**

Move to the `authlib` implementation and implement JWT in a way that comply with the standards.

**Note on backwards compatibility**

As always, we should be aware to not make any service disruption after upgrade to the new version. Session, set password and file download tokens are pretty volatile and recoverable. API keys are not and can be difficult to change.

That's why PR https://github.com/CERT-Polska/mwdb-core/pull/424 still waits for a good time to be merged. In this case we should also keep API keys backwards compatible, but mark them as `legacy` in the UI with proper information that they should be regenerated in the near future.



Discussion / hints:
Notes:

It would be good to consider on the list that could invalidate tokens or some system that determines the validity of tokens after their use.

Remember about password_ver and identity_ver values during implementation of new JWT tokens system in mwdb-core.

The following library looks better than authlib for this task:
https://pyjwt.readthedocs.io/en/stable/usage.html

Fix patch:
diff --git a/mwdb/core/auth.py b/mwdb/core/auth.py
new file mode 100644
index 000000000..bdd8b4397
--- /dev/null
+++ b/mwdb/core/auth.py
@@ -0,0 +1,66 @@
+import datetime
+from enum import Enum
+from typing import Any, Set
+
+import jwt
+
+from mwdb.core.config import app_config
+
+
+class AuthScope(Enum):
+    session = "session"
+    api_key = "api_key"
+    set_password = "set_password"
+    download_file = "download_file"
+
+
+def generate_token(fields, scope, expiration=24 * 3600):
+    issued_at = datetime.datetime.now(tz=datetime.timezone.utc)
+    token_claims = {
+        **fields,
+        "exp": (issued_at + datetime.timedelta(seconds=expiration)),
+        "iat": issued_at,
+        "aud": app_config.mwdb.base_url,
+        "scope": scope.value,
+    }
+
+    if "login" in fields.keys():
+        token_claims["sub"] = fields["login"]
+    payload = {**fields, **token_claims}
+    token = jwt.encode(payload, app_config.mwdb.secret_key, algorithm="HS512")
+    return token
+
+
+def verify_token(token: str, scope: AuthScope) -> Any:
+    try:
+        data = jwt.decode(
+            token,
+            key=app_config.mwdb.secret_key,
+            algorithms=["HS512"],
+            audience=app_config.mwdb.base_url,
+            options={"verify_aud": True},
+        )
+        if data.get("scope") != scope.value:
+            return None
+
+        if scope.value != "download_file" and "sub" not in data:
+            return None
+
+    except jwt.InvalidTokenError:
+        return None
+    return data
+
+
+def verify_legacy_token(token: str, required_fields: Set[str]) -> Any:
+    try:
+        data = jwt.decode(
+            token,
+            key=app_config.mwdb.secret_key,
+            algorithms=["HS512"],
+        )
+        if set(data.keys()) != required_fields:
+            return None
+
+    except jwt.InvalidTokenError:
+        return None
+    return data
diff --git a/mwdb/model/api_key.py b/mwdb/model/api_key.py
index d5a98f17d..b164180dd 100644
--- a/mwdb/model/api_key.py
+++ b/mwdb/model/api_key.py
@@ -1,11 +1,10 @@
 import datetime
 import uuid
 
-from itsdangerous import BadSignature, JSONWebSignatureSerializer, SignatureExpired
 from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.orm.exc import NoResultFound
 
-from mwdb.core.config import app_config
+from mwdb.core.auth import AuthScope, generate_token, verify_legacy_token, verify_token
 
 from . import db
 
@@ -29,16 +28,13 @@ def issuer_login(self):
 
     @staticmethod
     def verify_token(token):
-        s = JSONWebSignatureSerializer(app_config.mwdb.secret_key)
-        try:
-            data = s.loads(token)
-        except SignatureExpired:
-            return None
-        except BadSignature:
-            return None
+        data = verify_token(token, scope=AuthScope.api_key)
 
-        if "api_key_id" not in data:
-            return None
+        if data is None:
+            # check for legacy API Token
+            data = verify_legacy_token(token, required_fields={"login", "api_key_id"})
+            if data is None:
+                return None
 
         try:
             api_key_obj = APIKey.query.filter(
@@ -50,5 +46,7 @@ def verify_token(token):
         return api_key_obj.user
 
     def generate_token(self):
-        s = JSONWebSignatureSerializer(app_config.mwdb.secret_key)
-        return s.dumps({"login": self.user.login, "api_key_id": str(self.id)})
+        return generate_token(
+            {"login": self.user.login, "api_key_id": str(self.id)},
+            scope=AuthScope.api_key,
+        )
diff --git a/mwdb/model/file.py b/mwdb/model/file.py
index 0a9dc3c7a..d3f78831d 100644
--- a/mwdb/model/file.py
+++ b/mwdb/model/file.py
@@ -4,12 +4,12 @@
 import shutil
 import tempfile
 
-from itsdangerous import BadSignature, SignatureExpired, TimedJSONWebSignatureSerializer
 from sqlalchemy import or_
 from sqlalchemy.dialects.postgresql.array import ARRAY
 from sqlalchemy.ext.mutable import MutableList
 from werkzeug.utils import secure_filename
 
+from mwdb.core.auth import AuthScope, generate_token, verify_token
 from mwdb.core.config import StorageProviderType, app_config
 from mwdb.core.karton import send_file_to_karton
 from mwdb.core.util import (
@@ -280,21 +280,18 @@ def release_after_upload(self):
             self.upload_stream = None
 
     def generate_download_token(self):
-        serializer = TimedJSONWebSignatureSerializer(
-            app_config.mwdb.secret_key, expires_in=60
+        return generate_token(
+            {"identifier": self.sha256},
+            scope=AuthScope.download_file,
+            expiration=60,
         )
-        return serializer.dumps({"identifier": self.sha256})
 
     @staticmethod
     def get_by_download_token(download_token):
-        serializer = TimedJSONWebSignatureSerializer(app_config.mwdb.secret_key)
-        try:
-            download_req = serializer.loads(download_token)
-            return File.get(download_req["identifier"]).first()
-        except SignatureExpired:
-            return None
-        except BadSignature:
+        download_req = verify_token(download_token, scope=AuthScope.download_file)
+        if not download_req:
             return None
+        return File.get(download_req["identifier"]).first()
 
     def _send_to_karton(self):
         return send_file_to_karton(self)
diff --git a/mwdb/model/user.py b/mwdb/model/user.py
index 1831088e5..bdd54ad6b 100644
--- a/mwdb/model/user.py
+++ b/mwdb/model/user.py
@@ -3,12 +3,11 @@
 
 import bcrypt
 from flask import g
-from itsdangerous import BadSignature, SignatureExpired, TimedJSONWebSignatureSerializer
 from sqlalchemy import and_
 from sqlalchemy.ext.associationproxy import association_proxy
 from sqlalchemy.orm.exc import NoResultFound
 
-from mwdb.core.config import app_config
+from mwdb.core.auth import AuthScope, generate_token, verify_legacy_token, verify_token
 
 from . import db
 from .group import Group, Member
@@ -150,29 +149,25 @@ def create(
             db.session.commit()
         return user
 
-    def _generate_token(self, fields, expiration):
-        s = TimedJSONWebSignatureSerializer(
-            app_config.mwdb.secret_key, expires_in=expiration
-        )
-        return s.dumps(
+    def _generate_token(self, fields, scope, expiration):
+        token = generate_token(
             dict(
                 [("login", self.login)]
                 + [(field, getattr(self, field)) for field in fields]
-            )
+            ),
+            scope,
+            expiration,
         )
+        return token
 
     @staticmethod
-    def _verify_token(token, fields):
-        s = TimedJSONWebSignatureSerializer(app_config.mwdb.secret_key)
-        try:
-            data = s.loads(token)
-        except SignatureExpired:
-            return None
-        except BadSignature:
+    def _verify_token(token, fields, scope):
+        data = verify_token(token, scope)
+        if data is None:
             return None
 
         try:
-            user_obj = User.query.filter(User.login == data["login"]).one()
+            user_obj = User.query.filter(User.login == data["sub"]).one()
         except NoResultFound:
             return None
 
@@ -186,23 +181,47 @@ def _verify_token(token, fields):
 
     def generate_session_token(self):
         return self._generate_token(
-            ["password_ver", "identity_ver"], expiration=24 * 3600
+            ["password_ver", "identity_ver"],
+            scope=AuthScope.session,
+            expiration=24 * 3600,
         )
 
     def generate_set_password_token(self):
-        return self._generate_token(["password_ver"], expiration=14 * 24 * 3600)
+        return self._generate_token(
+            ["password_ver"],
+            scope=AuthScope.set_password,
+            expiration=14 * 24 * 3600,
+        )
 
     @staticmethod
     def verify_session_token(token):
-        return User._verify_token(token, ["password_ver", "identity_ver"])
+        return User._verify_token(
+            token, ["password_ver", "identity_ver"], scope=AuthScope.session
+        )
 
     @staticmethod
     def verify_set_password_token(token):
-        return User._verify_token(token, ["password_ver"])
+        return User._verify_token(
+            token,
+            ["password_ver"],
+            scope=AuthScope.set_password,
+        )
 
     @staticmethod
     def verify_legacy_token(token):
-        return User._verify_token(token, ["version_uid"])
+        data = verify_legacy_token(token, required_fields={"login", "version_uid"})
+        if data is None:
+            return None
+
+        try:
+            user_obj = User.query.filter(User.login == data["login"]).one()
+        except NoResultFound:
+            return None
+
+        if user_obj.version_uid != data["version_uid"]:
+            return None
+
+        return user_obj
 
     def is_member(self, group_id):
         groups = db.session.query(Member.group_id).filter(Member.user_id == self.id)
diff --git a/mwdb/resources/file.py b/mwdb/resources/file.py
index 5b48a27f3..b1bf3d5bd 100644
--- a/mwdb/resources/file.py
+++ b/mwdb/resources/file.py
@@ -454,4 +454,4 @@ def post(self, identifier):
 
         download_token = file.generate_download_token()
         schema = FileDownloadTokenResponseSchema()
-        return schema.dump({"token": download_token.decode()})
+        return schema.dump({"token": download_token})
diff --git a/requirements.txt b/requirements.txt
index 38467f45d..81f23531f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -25,5 +25,5 @@ karton-core==4.3.0
 Authlib==0.15.4
 prance==0.21.8.0  # apispec unpinned dependency
 regex==2021.7.6
-itsdangerous==2.0.1
-Flask-Limiter==2.1.3
+pyJWT==2.3.0
+Flask-Limiter==2.1.3
\ No newline at end of file


Test patch:
diff --git a/docker-compose-e2e.yml b/docker-compose-e2e.yml
index dbd0f1102..2e59e2d4e 100644
--- a/docker-compose-e2e.yml
+++ b/docker-compose-e2e.yml
@@ -14,6 +14,9 @@ services:
     env_file:
       # NOTE: use gen_vars.sh in order to generate this file
       - mwdb-vars.env
+    environment:
+      # Hardcoded secret key for consistent JWT testing
+      - MWDB_SECRET_KEY=e2e-testing-key
   mwdb-web:
     depends_on:
       - mwdb
diff --git a/tests/backend/requirements.txt b/tests/backend/requirements.txt
index b12d25d71..9ae59238d 100644
--- a/tests/backend/requirements.txt
+++ b/tests/backend/requirements.txt
@@ -1,4 +1,5 @@
 python-dateutil==2.7.0
+pyJWT==2.3.0
 pytest
 requests
-python-baseconv
\ No newline at end of file
+python-baseconv
diff --git a/tests/backend/test_auth.py b/tests/backend/test_auth.py
index 1b28bd80d..6375fb42b 100644
--- a/tests/backend/test_auth.py
+++ b/tests/backend/test_auth.py
@@ -1,4 +1,8 @@
-from .utils import MwdbTest, ShouldRaise, admin_login
+import os
+
+import jwt
+
+from .utils import MwdbTest, ShouldRaise, admin_login, rand_string
 
 
 def test_profile_change_invalidate(typical_session, admin_session):
@@ -17,9 +21,7 @@ def test_profile_change_invalidate(typical_session, admin_session):
     typical_session.recent_samples(1)
     typical_via_api.recent_samples(1)
 
-    admin_session.request("PUT", "/user/" + typical_login, json={
-        "disabled": True
-    })
+    admin_session.request("PUT", "/user/" + typical_login, json={"disabled": True})
 
     # Password-based session: invalidated, APIKey-based session: forbidden
     with ShouldRaise(status_code=403):
@@ -27,9 +29,7 @@ def test_profile_change_invalidate(typical_session, admin_session):
     with ShouldRaise(status_code=401):
         typical_session.recent_samples(1)
 
-    admin_session.request("PUT", "/user/" + typical_login, json={
-        "disabled": False
-    })
+    admin_session.request("PUT", "/user/" + typical_login, json={"disabled": False})
 
     # Password-based session: need to reauth, APIKey-based session: Ok
     typical_via_api.recent_samples(1)
@@ -44,32 +44,39 @@ def test_profile_password_change(typical_session, admin_session):
     typical_login = typical_session.userinfo["login"]
     typical_session.add_sample()
 
-    set_pass_token = admin_session.request("GET", f"/user/{typical_login}/change_password")["token"]
+    set_pass_token = admin_session.request(
+        "GET", f"/user/{typical_login}/change_password"
+    )["token"]
 
     # Shouldn't be able to use as session token
     typical_session.set_auth_token(set_pass_token)
     with ShouldRaise(status_code=401):
         typical_session.recent_samples(1)
 
-    typical_session.request("POST", "/auth/change_password", json={
-        "password": "very_new_password",
-        "token": set_pass_token
-    })
+    typical_session.request(
+        "POST",
+        "/auth/change_password",
+        json={"password": "very_new_password", "token": set_pass_token},
+    )
 
     typical_session.login_as(typical_login, "very_new_password")
 
     # Valid only once
     with ShouldRaise(status_code=403):
-        typical_session.request("POST", "/auth/change_password", json={
-            "password": typical_login,
-            "token": set_pass_token
-        })
-
-    set_pass_token = admin_session.request("GET", f"/user/{typical_login}/change_password")["token"]
-    typical_session.request("POST", "/auth/change_password", json={
-        "password": typical_login,
-        "token": set_pass_token
-    })
+        typical_session.request(
+            "POST",
+            "/auth/change_password",
+            json={"password": typical_login, "token": set_pass_token},
+        )
+
+    set_pass_token = admin_session.request(
+        "GET", f"/user/{typical_login}/change_password"
+    )["token"]
+    typical_session.request(
+        "POST",
+        "/auth/change_password",
+        json={"password": typical_login, "token": set_pass_token},
+    )
     typical_session.login_as(typical_login, typical_login)
     typical_session.recent_samples(1)
 
@@ -83,7 +90,9 @@ def test_api_key_management(typical_session, admin_session):
 
     # It should be possible to create API key without sending any payload
     # Should fall back to empty name
-    res = admin_session.session.post(admin_session.mwdb_url + f"/user/{typical_login}/api_key")
+    res = admin_session.session.post(
+        admin_session.mwdb_url + f"/user/{typical_login}/api_key"
+    )
     res.raise_for_status()
     assert res.json()["name"] == ""
 
@@ -105,3 +114,67 @@ def test_api_key_management(typical_session, admin_session):
 
     with ShouldRaise(status_code=401):
         typical.add_sample()
+
+
+def test_jwt_legacy_api_keys(admin_session):
+    secret_key = "e2e-testing-key"
+    api_key_id = admin_session.api_key_create("admin", "testing-key").json()["id"]
+
+    pre2_0_payload = {"login": "admin", "version_uid": None}
+    pre2_7_payload = {"login": "admin", "api_key_id": api_key_id}
+
+    pre2_0_token = jwt.encode(pre2_0_payload, secret_key, algorithm="HS512")
+    pre2_7_token = jwt.encode(pre2_7_payload, secret_key, algorithm="HS512")
+
+    session = MwdbTest()
+    response = session.request("get", "/server")
+    assert not response["is_authenticated"]
+
+    session.set_auth_token(pre2_0_token)
+    response = session.request("get", "/server")
+    assert response["is_authenticated"]
+
+    session.set_auth_token(pre2_7_token)
+    response = session.request("get", "/server")
+    assert response["is_authenticated"]
+
+
+def test_invalid_jwt(admin_session):
+    secret_key = "e2e-testing-key"
+    correct_claims = jwt.decode(admin_session.auth_token, options={"verify_signature": False})
+
+    session = MwdbTest()
+
+    # Test if JWT is correctly validated
+    incorrect_audience = {**correct_claims, "aud": "http://evil.com"}
+    session.set_auth_token(jwt.encode(incorrect_audience, secret_key, algorithm="HS512"))
+    response = session.request("get", "/server")
+    assert not response["is_authenticated"]
+
+    expired_jwt = {**correct_claims, "exp": correct_claims["iat"]}
+    session.set_auth_token(jwt.encode(expired_jwt, secret_key, algorithm="HS512"))
+    response = session.request("get", "/server")
+    assert not response["is_authenticated"]
+
+    wrong_scope_jwt = {**correct_claims, "scope": "420_no_scope"}
+    session.set_auth_token(jwt.encode(wrong_scope_jwt, secret_key, algorithm="HS512"))
+    response = session.request("get", "/server")
+    assert not response["is_authenticated"]
+
+    # Test missing fields to check if old token doesn't cause ISE 500
+    only_login_jwt = {"login": "admin"}
+    session.set_auth_token(jwt.encode(only_login_jwt, secret_key, algorithm="HS512"))
+    response = session.request("get", "/server")
+    assert not response["is_authenticated"]
+
+    missing_scope_jwt = {**correct_claims}
+    del missing_scope_jwt["scope"]
+    session.set_auth_token(jwt.encode(missing_scope_jwt, secret_key, algorithm="HS512"))
+    response = session.request("get", "/server")
+    assert not response["is_authenticated"]
+
+    missing_sub_jwt = {**correct_claims}
+    del missing_sub_jwt["sub"]
+    session.set_auth_token(jwt.encode(missing_sub_jwt, secret_key, algorithm="HS512"))
+    response = session.request("get", "/server")
+    assert not response["is_authenticated"]
diff --git a/tests/backend/utils.py b/tests/backend/utils.py
index 6cdc8b2d4..6e8da5383 100644
--- a/tests/backend/utils.py
+++ b/tests/backend/utils.py
@@ -53,11 +53,12 @@ def __init__(self):
         self.mwdb_url = os.environ.get("MWDB_URL")
         self.session = requests.session()
         self.userinfo = None
+        self.auth_token = None
 
     def clone_session(self):
         cloned = MwdbTest()
         cloned.userinfo = self.userinfo
-        cloned.session.headers.update({"Authorization": "Bearer " + self.userinfo["token"]})
+        cloned.set_auth_token(self.userinfo["token"])
         return cloned
 
     def login_as(self, username, password):
@@ -67,10 +68,11 @@ def login_as(self, username, password):
         )
         res.raise_for_status()
         self.userinfo = res.json()
-        self.session.headers.update({"Authorization": "Bearer " + self.userinfo["token"]})
+        self.set_auth_token(self.userinfo["token"])
         return self.userinfo
 
     def set_auth_token(self, token):
+        self.auth_token = token
         self.session.headers.update({"Authorization": "Bearer " + token})
 
     def request(self, method, endpoint, params=None, json=None):