Maor/security/csrf fix (#62)

* Fix CSRF security issue * CSRF for tests * Disable CSRF for tests * Disable CSRF check on token login * Fix circular imports * Fix attribute * Enable csrf in tests * Revert some changes * Fixes * Fixes * Revert * csrf in api_client * csrf cookie * api_request * no_prefix for csrf request * csrf route without assets * Add additional call for csrf request * Add additional call_count for csrf request * Add additional call_count for csrf request * Add additional call_count for csrf request * exempt internal ml trainer apis * exempt tracking endpoints * Lock exceptiongroup version * Change endpoint for csrf token * Remove unneeded param * Cookie max-age * Linter * Add CSRF token handling through cookies (#202) * Fix CSRF security issue * CSRF for tests * Disable CSRF for tests * Disable CSRF check on token login * Fix circular imports * Fix attribute * Enable csrf in tests * Revert some changes * Fixes * Fixes * Revert * csrf in api_client * csrf cookie * api_request * no_prefix for csrf request * csrf route without assets * Add additional call for csrf request * Add additional call_count for csrf request * Add additional call_count for csrf request * Add additional call_count for csrf request * exempt internal ml trainer apis * exempt tracking endpoints * Lock exceptiongroup version * Change endpoint for csrf token * Remove unneeded param * Add CSRF token handling through cookies * Replace custom getCookie function with Quasar's Cookies Co-authored-by: Maor Katzav <maor.katzav@databand.ai> Co-authored-by: Maor Katzav <maor.katzav@ibm.com> * CR fixes * CR fixes * Comment * replace cy.request with custom cy.request that supports csrf token * fix for the request overwrite * remove cookie before login * Exempt seeding api * Linter * remove cookies on logout * remove token after invitation activation * don't remove token after invitation activation * logout before user activation * Fix CSRF security issue * menu logout before user activation * Linter * skip 'invite user' test * Add csrf token to each ml trainer request for webserver. * Add logs * Exempt more internal bp Co-authored-by: Maor Katzav <maor.katzav@databand.ai> Co-authored-by: Illia Keba <IlliaKeba@ibm.com> Co-authored-by: Roza Prag <rose.prag@databand.ai>
databand-ai · Jan 2, 2024 · 2903ec3 · 2903ec3
1 parent 58c0d53
commit 2903ec3
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 5 deletions.
diff --git a/modules/dbnd/src/dbnd/utils/api_client.py b/modules/dbnd/src/dbnd/utils/api_client.py
@@ -196,6 +196,13 @@ def _authenticate(self):
                 self.default_headers["Authorization"] = "Bearer {}".format(token)
                 return
 
+            # get the csrf token cookie (if enabled on the server)
+            self.api_request("auth/csrf", None, method="GET")
+            csrf_token = self.session.cookies.get("X-CSRF-TOKEN")
+            if csrf_token:
+                logger.debug("Got csrf token from session")
+                self.default_headers["X-CSRFToken"] = csrf_token
+
             if "username" in credentials and "password" in credentials:
                 logger.debug("Attempting to login to webserver")
                 self.api_request("auth/login", method="POST", data=credentials)

diff --git a/modules/dbnd/test_dbnd/utils/test_api_client.py b/modules/dbnd/test_dbnd/utils/test_api_client.py
@@ -27,6 +27,7 @@ def test_dont_fail_after_first_error(self):
             self.connection_error,
             self.ok_response,
             self.ok_response,
+            self.ok_response,
         ]
         sut = self.create_sut(retries=2)
         try:
@@ -53,11 +54,15 @@ def test_after_anon_session_created_should_login_on_next_request(self):
             self.fail(
                 f"Api client should be resilient to failures, but instead got: {e}"
             )
-        assert self.network_request_mock.call_count == 3
-        assert self.network_request_mock.mock_calls[1] == self.login_call()
+        assert self.network_request_mock.call_count == 4
+        assert self.network_request_mock.mock_calls[2] == self.login_call()
 
     def test_api_client_should_retry_login_3_times(self):
-        self.network_request_mock.side_effect = self.connection_error
+        self.network_request_mock.side_effect = [
+            self.ok_response,
+            self.connection_error,
+            self.connection_error,
+        ]
         sut = self.create_sut(retries=3)
         try:
             with self.dont_sleep():
@@ -66,9 +71,10 @@ def test_api_client_should_retry_login_3_times(self):
             # Request will fail, this is expected
             pass
 
+        _csrf_call = self.csrf_call()
         _call = self.login_call()
-        assert self.network_request_mock.call_count == 3
-        self.network_request_mock.assert_has_calls([_call, _call, _call])
+        assert self.network_request_mock.call_count == 4
+        self.network_request_mock.assert_has_calls([_csrf_call, _call, _call, _call])
 
     @mock.patch("requests.session")
     def test_api_client_with_key_error_should_retry_3_times(self, mock_session):
@@ -106,6 +112,17 @@ def login_call(self):
             session=mock.ANY,
         )
 
+    def csrf_call(self):
+        return call(
+            "/api/v1/auth/csrf",
+            method="GET",
+            data=mock.ANY,
+            headers=mock.ANY,
+            query=mock.ANY,
+            request_timeout=mock.ANY,
+            session=mock.ANY,
+        )
+
     @contextmanager
     def dont_sleep(self):
         with patch.object(LinearRetryPolicy, "seconds_to_sleep") as secs: