cornflow 2.0.0a11__py3-none-any.whl → 2.0.0a12__py3-none-any.whl

cornflow/tests/unit/test_log_in.py

@@ -5,6 +5,9 @@ Unit test for the log in endpoint
 import json
 import logging as log
 from unittest import mock
+import requests
+import jwt
+from datetime import datetime, timedelta
 
 from flask import current_app
 
@@ -14,7 +17,7 @@ from cornflow.commands.access import access_init_command
 from cornflow.commands.dag import register_deployed_dags_command_test
 from cornflow.models import UserModel
 from cornflow.shared import db
-from cornflow.shared.const import SERVICE_ROLE, OID_GOOGLE, OID_AZURE, OID_NONE
+from cornflow.shared.const import SERVICE_ROLE, INTERNAL_TOKEN_ISSUER, AUTH_OID
 from cornflow.tests.const import LOGIN_URL
 from cornflow.tests.custom_test_case import CustomTestCase, LoginTestCases
 
@@ -65,7 +68,7 @@ class TestLogIn(LoginTestCases.LoginEndpoint):
         self.assertIn("Error in generating user token", response.json["error"])
 
 
-class TestLogInOpenAuthNoConfig(CustomTestCase):
+class TestLogInOpenAuth(CustomTestCase):
     def create_app(self):
         """
         Creates and configures a Flask application for testing.
@@ -92,9 +95,9 @@ class TestLogInOpenAuthNoConfig(CustomTestCase):
         test_user.save()
 
         self.user_data.pop("email")
-
         self.test_user_id = test_user.id
 
+        # Setup service user
         self.service_data = {
             "username": "service_user",
             "email": "service@test.com",
@@ -120,283 +123,294 @@ class TestLogInOpenAuthNoConfig(CustomTestCase):
         )
 
         self.assertEqual(400, response.status_code)
-        self.assertEqual(response.json["error"], "Invalid request")
-
-    def test_other_user_password(self):
-        """
-        Tests that a user can not log in with username and password
-        """
-        response = self.client.post(
-            LOGIN_URL,
-            data=json.dumps(self.user_data),
-            headers={"Content-Type": "application/json"},
+        self.assertEqual(
+            response.json["error"],
+            "Must provide a token in Authorization header. Cannot log in with username and password",
         )
 
-        self.assertEqual(400, response.status_code)
-        self.assertEqual(response.json["error"], "Invalid request")
-
-
-class TestLogInOpenAuthAzure(CustomTestCase):
-    def create_app(self):
+    @mock.patch("cornflow.shared.authentication.auth.Auth.get_public_keys")
+    @mock.patch("cornflow.shared.authentication.auth.jwt")
+    def test_kid_not_in_public_keys(self, mock_jwt, mock_get_public_keys):
         """
-        Creates and configures a Flask application for testing.
-
-        :returns: A configured Flask application instance
-        :rtype: Flask
+        Tests token validation failure when the kid is not found in public keys
         """
-        app = create_app("testing-oauth")
-        app.config["SERVICE_USER_ALLOW_PASSWORD_LOGIN"] = 0
-        app.config["OID_PROVIDER"] = OID_AZURE
-        app.config["OID_CLIENT_ID"] = "SOME_SECRET"
-        app.config["OID_TENANT_ID"] = "SOME_SECRET"
-        app.config["OID_ISSUER"] = "SOME_SECRET"
-        return app
+        # Import the real exceptions to ensure they are preserved
+        from jwt.exceptions import InvalidTokenError, ExpiredSignatureError
 
-    def setUp(self):
-        log.root.setLevel(current_app.config["LOG_LEVEL"])
-        db.create_all()
-        access_init_command(verbose=False)
-        register_deployed_dags_command_test(verbose=False)
+        # Keep the real exception classes in the mock
+        mock_jwt.ExpiredSignatureError = ExpiredSignatureError
+        mock_jwt.InvalidTokenError = InvalidTokenError
 
-        self.service_data = {
-            "username": "service_user",
-            "email": "service@test.com",
-            "password": "Testpassword1!",
-        }
+        mock_jwt.get_unverified_header.return_value = {"kid": "test_kid"}
 
-        service_user = UserModel(data=self.service_data)
-        service_user.save()
+        # Mock jwt.decode to return different results based on arguments
+        def decode_side_effect(*args, **kwargs):
+            if kwargs.get("options", {}).get("verify_signature") is False:
+                return {"iss": current_app.config.get("OID_PROVIDER", "valid_issuer")}
+            raise InvalidTokenError("Invalid token")
 
-        self.service_data.pop("email")
-        self.service_user_id = service_user.id
+        mock_jwt.decode.side_effect = decode_side_effect
 
-        self.assign_role(self.service_user_id, SERVICE_ROLE)
+        # Mock get_public_keys to return keys that don't contain our kid
+        mock_get_public_keys.return_value = {"different_kid": "some_key"}
 
-    def test_service_user_login_first_fail(self):
-        """
-        Tests that a service user can not log in with username and password
-        """
         response = self.client.post(
             LOGIN_URL,
-            data=json.dumps({"token": "some_token"}),
-            headers={"Content-Type": "application/json"},
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer some_token",
+            },
         )
 
         self.assertEqual(400, response.status_code)
-        self.assertEqual(response.json["error"], "Token is not valid")
+        self.assertEqual(
+            response.json["error"], "Invalid token: Unknown key identifier (kid)"
+        )
 
     @mock.patch("cornflow.shared.authentication.auth.jwt")
-    def test_service_user_login_second_fail(self, mock_header):
+    def test_missing_kid_in_token(self, mock_jwt):
         """
-        Tests second exit point on token validation
+        Tests token validation failure when the token header is missing the kid
         """
-        mock_header.get_unverified_header.return_value = None
+        # Import the real exceptions to ensure they are preserved
+        from jwt.exceptions import InvalidTokenError, ExpiredSignatureError
+
+        # Keep the real exception classes in the mock
+        mock_jwt.ExpiredSignatureError = ExpiredSignatureError
+        mock_jwt.InvalidTokenError = InvalidTokenError
+
+        # Mock jwt.get_unverified_header to return a header without kid
+        mock_jwt.get_unverified_header.return_value = {"alg": "RS256"}
+
         response = self.client.post(
             LOGIN_URL,
-            data=json.dumps({"token": "some_token"}),
-            headers={"Content-Type": "application/json"},
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer some_token",
+            },
         )
 
         self.assertEqual(400, response.status_code)
-        self.assertEqual(response.json["error"], "Token is missing the headers")
+        self.assertEqual(
+            response.json["error"],
+            "Invalid token: Missing key identifier (kid) in token header",
+        )
 
+    @mock.patch("cornflow.shared.authentication.auth.requests.get")
     @mock.patch("cornflow.shared.authentication.auth.jwt")
-    def test_service_user_login_third_fail(self, mock_header):
+    def test_public_keys_fetch_fail(self, mock_jwt, mock_get):
         """
-        Tests third exit point on token validation
+        Tests failure when trying to fetch public keys from the OIDC provider
         """
-        mock_header.get_unverified_header.return_value = {"kid": "some_value"}
+        # Import the real exceptions to ensure they are preserved
+        from jwt.exceptions import InvalidTokenError, ExpiredSignatureError
+
+        # Keep the real exception classes in the mock
+        mock_jwt.ExpiredSignatureError = ExpiredSignatureError
+        mock_jwt.InvalidTokenError = InvalidTokenError
+
+        # Clear the cache
+        from cornflow.shared.authentication.auth import public_keys_cache
+
+        public_keys_cache.clear()
+
+        # Mock jwt to pass initial validation
+        mock_jwt.get_unverified_header.return_value = {"kid": "test_kid"}
+        mock_jwt.decode.side_effect = lambda *args, **kwargs: (
+            {"iss": current_app.config["OID_PROVIDER"]}
+            if kwargs.get("options", {}).get("verify_signature") is False
+            else {}
+        )
+
+        # Mock get to fail
+        mock_get.side_effect = requests.exceptions.RequestException(
+            "Failed to get keys"
+        )
+
         response = self.client.post(
             LOGIN_URL,
-            data=json.dumps({"token": "some_token"}),
-            headers={"Content-Type": "application/json"},
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer some_token",
+            },
         )
 
         self.assertEqual(400, response.status_code)
-        self.assertIn(
-            "Error getting issuer discovery meta from", response.json["error"]
+        self.assertEqual(
+            response.json["error"],
+            "Failed to fetch public keys from authentication provider",
         )
 
-    @mock.patch("cornflow.endpoints.login.Auth.validate_oid_token")
-    def test_service_user_login_no_fail(self, mock_auth):
-        mock_auth.return_value = {"preferred_username": "service_user"}
+    @mock.patch("cornflow.shared.authentication.Auth.decode_token")
+    def test_service_user_login_no_fail(self, mock_decode):
+        """
+        Tests successful login for an existing service user with valid token
+        """
+        mock_decode.return_value = {"sub": "service_user"}
         response = self.client.post(
             LOGIN_URL,
-            data=json.dumps({"token": "some_token"}),
-            headers={"Content-Type": "application/json"},
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer some_token",
+            },
         )
 
-        print(response.json)
-
         self.assertEqual(200, response.status_code)
         self.assertEqual(self.service_user_id, response.json["id"])
 
-    @mock.patch("cornflow.endpoints.login.Auth.validate_oid_token")
-    def test_new_user_login_no_fail(self, mock_auth):
-        mock_auth.return_value = {"preferred_username": "test_user"}
+    @mock.patch("cornflow.shared.authentication.Auth.decode_token")
+    def test_new_user_login_no_fail(self, mock_decode):
+        """
+        Tests successful login and creation of a new user with valid token
+        """
+        mock_decode.return_value = {"sub": "test_user"}
         response = self.client.post(
             LOGIN_URL,
-            data=json.dumps({"token": "some_token"}),
-            headers={"Content-Type": "application/json"},
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer some_token",
+            },
         )
 
         self.assertEqual(200, response.status_code)
         self.assertEqual(self.service_user_id + 1, response.json["id"])
 
-
-class TestLogInOpenAuthGoogle(CustomTestCase):
-    def create_app(self):
+    @mock.patch("cornflow.shared.authentication.Auth.verify_token")
+    @mock.patch("cornflow.shared.authentication.auth.jwt")
+    def test_public_keys_caching(self, mock_jwt, mock_verify_token):
         """
-        Creates and configures a Flask application for testing.
-
-        :returns: A configured Flask application instance
-        :rtype: Flask
+        Tests that public keys are cached and reused for subsequent requests.
+        Also verifies that when a new kid is encountered that's not in the cache,
+        the system fetches fresh keys from the provider.
         """
-        app = create_app("testing-oauth")
-        app.config["SERVICE_USER_ALLOW_PASSWORD_LOGIN"] = 0
-        app.config["OID_PROVIDER"] = OID_GOOGLE
-        app.config["OID_CLIENT_ID"] = "SOME_SECRET"
-        app.config["OID_TENANT_ID"] = "SOME_SECRET"
-        app.config["OID_ISSUER"] = "SOME_SECRET"
-        return app
+        # Import the real exceptions to ensure they are preserved
+        from jwt.exceptions import InvalidTokenError, ExpiredSignatureError
 
-    def setUp(self):
-        log.root.setLevel(current_app.config["LOG_LEVEL"])
-        db.create_all()
-        access_init_command(verbose=False)
-        register_deployed_dags_command_test(verbose=False)
+        # Keep the real exception classes in the mock
+        mock_jwt.ExpiredSignatureError = ExpiredSignatureError
+        mock_jwt.InvalidTokenError = InvalidTokenError
 
-        self.service_data = {
-            "username": "service_user",
-            "email": "service@test.com",
-            "password": "Testpassword1!",
-        }
+        # Mock jwt to return valid unverified header and payload
+        mock_jwt.get_unverified_header.return_value = {"kid": "test_kid"}
+        mock_jwt.decode.side_effect = lambda *args, **kwargs: (
+            {"iss": current_app.config["OID_PROVIDER"]}
+            if kwargs.get("options", {}).get("verify_signature") is False
+            else {"sub": "test_user", "email": "test_user@test.com"}
+        )
 
-        service_user = UserModel(data=self.service_data)
-        service_user.save()
+        # Mock verify_token to always return a valid token payload
+        mock_verify_token.return_value = {
+            "sub": "test_user",
+            "email": "test_user@test.com",
+        }
 
-        self.service_data.pop("email")
-        self.service_user_id = service_user.id
+        # Make first request
+        response = self.client.post(
+            LOGIN_URL,
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer some_token",
+            },
+        )
 
-        self.assign_role(self.service_user_id, SERVICE_ROLE)
+        self.assertEqual(200, response.status_code)
 
-    def test_service_user_login(self):
-        """
-        Tests that a service user can not log in with username and password
-        """
+        # Make second request
         response = self.client.post(
             LOGIN_URL,
-            data=json.dumps({"token": "some_token"}),
-            headers={"Content-Type": "application/json"},
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer some_token",
+            },
         )
 
-        self.assertEqual(501, response.status_code)
-        self.assertEqual(
-            response.json["error"], "The selected OID provider is not implemented"
-        )
+        self.assertEqual(200, response.status_code)
 
+        # Verify token was verified twice
+        self.assertEqual(2, mock_verify_token.call_count)
 
-class TestLogInOpenAuthNone(CustomTestCase):
-    def create_app(self):
-        """
-        Creates and configures a Flask application for testing.
+        # Now test with a different token
+        response = self.client.post(
+            LOGIN_URL,
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer some_different_token",
+            },
+        )
 
-        :returns: A configured Flask application instance
-        :rtype: Flask
-        """
-        app = create_app("testing-oauth")
-        app.config["SERVICE_USER_ALLOW_PASSWORD_LOGIN"] = 0
-        app.config["OID_PROVIDER"] = OID_NONE
-        app.config["OID_CLIENT_ID"] = "SOME_SECRET"
-        app.config["OID_TENANT_ID"] = "SOME_SECRET"
-        app.config["OID_ISSUER"] = "SOME_SECRET"
-        return app
+        self.assertEqual(200, response.status_code)
 
-    def setUp(self):
-        log.root.setLevel(current_app.config["LOG_LEVEL"])
-        db.create_all()
-        access_init_command(verbose=False)
-        register_deployed_dags_command_test(verbose=False)
+        # Verify token was verified a third time
+        self.assertEqual(3, mock_verify_token.call_count)
 
-        self.service_data = {
-            "username": "service_user",
-            "email": "service@test.com",
-            "password": "Testpassword1!",
+    def test_old_token(self):
+        """
+        Tests using an expired token.
+        """
+        # Generate an expired token
+        payload = {
+            # Token expired 1 hour ago
+            "exp": datetime.utcnow() - timedelta(hours=1),
+            # Token created 2 hours ago
+            "iat": datetime.utcnow() - timedelta(hours=2),
+            "sub": "testname",
+            "iss": INTERNAL_TOKEN_ISSUER,
         }
 
-        service_user = UserModel(data=self.service_data)
-        service_user.save()
-
-        self.service_data.pop("email")
-        self.service_user_id = service_user.id
-
-        self.assign_role(self.service_user_id, SERVICE_ROLE)
+        expired_token = jwt.encode(
+            payload, current_app.config["SECRET_TOKEN_KEY"], algorithm="HS256"
+        )
 
-    def test_service_user_login(self):
-        """
-        Tests that a service user can not log in with username and password
-        """
         response = self.client.post(
             LOGIN_URL,
-            data=json.dumps({"token": "some_token"}),
-            headers={"Content-Type": "application/json"},
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {expired_token}",
+            },
         )
 
-        self.assertEqual(501, response.status_code)
+        self.assertEqual(400, response.status_code)
         self.assertEqual(
-            response.json["error"], "The OID provider configuration is not valid"
+            "The token has expired, please login again", response.json["error"]
         )
 
-
-class TestLogInOpenAuthOther(CustomTestCase):
-    def create_app(self):
+    def test_missing_auth_header(self):
         """
-        Creates and configures a Flask application for testing.
-
-        :returns: A configured Flask application instance
-        :rtype: Flask
+        Tests that missing Authorization header raises proper error
         """
-        app = create_app("testing-oauth")
-        app.config["SERVICE_USER_ALLOW_PASSWORD_LOGIN"] = 0
-        app.config["OID_PROVIDER"] = 3
-        app.config["OID_CLIENT_ID"] = "SOME_SECRET"
-        app.config["OID_TENANT_ID"] = "SOME_SECRET"
-        app.config["OID_ISSUER"] = "SOME_SECRET"
-        return app
-
-    def setUp(self):
-        log.root.setLevel(current_app.config["LOG_LEVEL"])
-        db.create_all()
-        access_init_command(verbose=False)
-        register_deployed_dags_command_test(verbose=False)
-
-        self.service_data = {
-            "username": "service_user",
-            "email": "service@test.com",
-            "password": "Testpassword1!",
-        }
-
-        service_user = UserModel(data=self.service_data)
-        service_user.save()
-
-        self.service_data.pop("email")
-        self.service_user_id = service_user.id
+        response = self.client.post(
+            LOGIN_URL, data=json.dumps({}), headers={"Content-Type": "application/json"}
+        )
 
-        self.assign_role(self.service_user_id, SERVICE_ROLE)
+        self.assertEqual(400, response.status_code)
+        self.assertEqual(response.json["error"], "Authorization header is missing")
 
-    def test_service_user_login(self):
+    def test_invalid_auth_header(self):
         """
-        Tests that a service user can not log in with username and password
+        Tests that malformed Authorization header raises proper error
         """
         response = self.client.post(
             LOGIN_URL,
-            data=json.dumps({"token": "some_token"}),
-            headers={"Content-Type": "application/json"},
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Invalid Format",
+            },
         )
 
-        self.assertEqual(501, response.status_code)
+        self.assertEqual(400, response.status_code)
         self.assertEqual(
-            response.json["error"], "The OID provider configuration is not valid"
+            "Invalid Authorization header format. Must be 'Bearer <token>'",
+            response.json["error"],
         )
 
 
@@ -410,6 +424,7 @@ class TestLogInOpenAuthService(CustomTestCase):
         :rtype: Flask
         """
         app = create_app("testing-oauth")
+        app.config["SERVICE_USER_ALLOW_PASSWORD_LOGIN"] = 1
         return app
 
     def setUp(self):
@@ -448,7 +463,6 @@ class TestLogInOpenAuthService(CustomTestCase):
         """
         Tests that a service user can log in with username and password
         """
-
         response = self.client.post(
             LOGIN_URL,
             data=json.dumps(self.service_data),
@@ -458,7 +472,7 @@ class TestLogInOpenAuthService(CustomTestCase):
         self.assertEqual(200, response.status_code)
         self.assertEqual(self.service_user_id, response.json["id"])
 
-    def test_validation_error(self):
+    def test_no_credentials_error(self):
         """
         Tests that a user can not log in without token or username and password
         """
@@ -469,6 +483,7 @@ class TestLogInOpenAuthService(CustomTestCase):
         )
 
         self.assertEqual(400, response.status_code)
+        self.assertEqual(response.json["error"], "Authorization header is missing")
 
     def test_other_user_password(self):
         """
@@ -483,29 +498,37 @@ class TestLogInOpenAuthService(CustomTestCase):
         self.assertEqual(400, response.status_code)
         self.assertEqual(response.json["error"], "Invalid request")
 
-    def test_token_login(self):
+    @mock.patch("cornflow.shared.authentication.Auth.decode_token")
+    def test_token_login(self, mock_decode):
         """
-        Tests that a user can log in with a token
+        Tests that a user can successfully log in with a valid token
         """
+        mock_decode.return_value = {"sub": "testname"}
+
         response = self.client.post(
             LOGIN_URL,
-            data=json.dumps({"token": "test"}),
-            headers={"Content-Type": "application/json"},
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer valid_token",
+            },
         )
 
-        self.assertEqual(501, response.status_code)
-        self.assertEqual(
-            response.json["error"], "The OID provider configuration is not valid"
-        )
+        self.assertEqual(200, response.status_code)
+        self.assertEqual(self.test_user_id, response.json["id"])
 
-    def test_log_in_with_all_fields(self):
+    def test_invalid_token_login(self):
         """
-        Tests that a user can not log in with username and password and a token
+        Tests that login fails with an invalid token
         """
         response = self.client.post(
             LOGIN_URL,
-            data=json.dumps({**self.service_data, "token": "test"}),
-            headers={"Content-Type": "application/json"},
+            data=json.dumps({}),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": "Bearer invalid_token",
+            },
         )
 
         self.assertEqual(400, response.status_code)
+        self.assertEqual(response.json["error"], "Invalid token format or signature")