From b39517745b0b295eca6791a449dbf369bac85587 Mon Sep 17 00:00:00 2001 From: Yury Dzerin Date: Wed, 23 Feb 2022 16:16:52 +0100 Subject: [PATCH 01/13] Implement mechanism for caching token on login --- .github/workflows/code-check.yml | 4 + src/firebolt/client/auth.py | 17 ++- src/firebolt/client/resource_manager_hooks.py | 12 +- src/firebolt/common/token_storage.py | 114 ++++++++++++++++++ tests/unit/client/test_auth.py | 11 +- tests/unit/client/test_auth_async.py | 12 +- tests/unit/client/test_client.py | 3 + tests/unit/client/test_client_async.py | 3 + tests/unit/common/test_token_storage.py | 104 ++++++++++++++++ 9 files changed, 261 insertions(+), 19 deletions(-) create mode 100644 src/firebolt/common/token_storage.py create mode 100644 tests/unit/common/test_token_storage.py diff --git a/.github/workflows/code-check.yml b/.github/workflows/code-check.yml index e474926b1c2..759ebfa992e 100644 --- a/.github/workflows/code-check.yml +++ b/.github/workflows/code-check.yml @@ -26,3 +26,7 @@ jobs: uses: pre-commit/action@v2.0.3 with: extra_args: --all-files + + - name: Type check with mypy + run: | + mypy src diff --git a/src/firebolt/client/auth.py b/src/firebolt/client/auth.py index 95d3c721813..78697601e00 100644 --- a/src/firebolt/client/auth.py +++ b/src/firebolt/client/auth.py @@ -6,6 +6,7 @@ from firebolt.client.constants import _REQUEST_ERRORS, DEFAULT_API_URL from firebolt.common.exception import AuthenticationError +from firebolt.common.token_storage import TokenSecureStorage from firebolt.common.urls import AUTH_URL from firebolt.common.util import fix_url_schema @@ -34,13 +35,18 @@ def from_token(token: str) -> "Auth": return a def __init__( - self, username: str, password: str, api_endpoint: str = DEFAULT_API_URL + self, + username: str, + password: str, + api_endpoint: str = DEFAULT_API_URL, ): self.username = username self.password = password + self._token_storage = TokenSecureStorage(username=username, password=password) + # Add schema to url if it's missing self._api_endpoint = fix_url_schema(api_endpoint) - self._token: Optional[str] = None + self._token: Optional[str] = self._token_storage.get_cached_token() self._expires: Optional[int] = None def copy(self) -> "Auth": @@ -56,7 +62,6 @@ def expired(self) -> Optional[int]: def get_new_token_generator(self) -> Generator[Request, Response, None]: """Get new token using username and password""" - try: response = yield Request( "POST", @@ -74,6 +79,9 @@ def get_new_token_generator(self) -> Generator[Request, Response, None]: self._token = parsed["access_token"] self._expires = int(time()) + int(parsed["expires_in"]) + + self._token_storage.cache_token(parsed["access_token"]) + except _REQUEST_ERRORS as e: raise AuthenticationError(repr(e), self._api_endpoint) @@ -83,8 +91,11 @@ def auth_flow(self, request: Request) -> Generator[Request, Response, None]: if not self.token or self.expired: yield from self.get_new_token_generator() + request.headers["Authorization"] = f"Bearer {self.token}" + response = yield request + if response.status_code == codes.UNAUTHORIZED: yield from self.get_new_token_generator() request.headers["Authorization"] = f"Bearer {self.token}" diff --git a/src/firebolt/client/resource_manager_hooks.py b/src/firebolt/client/resource_manager_hooks.py index 1682bbb1fc7..c4318530804 100644 --- a/src/firebolt/client/resource_manager_hooks.py +++ b/src/firebolt/client/resource_manager_hooks.py @@ -1,7 +1,7 @@ from json import JSONDecodeError from logging import getLogger -from httpx import HTTPStatusError, Request, RequestError, Response +from httpx import HTTPStatusError, Request, RequestError, Response, codes logger = getLogger(__name__) @@ -40,14 +40,24 @@ def raise_on_4xx_5xx(response: Response) -> None: parsed_response = exc.response.json() except JSONDecodeError: parsed_response = {"_raw": exc.response.text} + + if exc.response.status_code == codes.UNAUTHORIZED: + # Do not raise an exception on UNAUTHORIZED + # It should be responsibility of authentication mechanism + return + debug_message = ( f"Error response {exc.response.status_code} " f"while requesting {exc.request.url!r}. " f"Response: {parsed_response}. " ) + if "message" in parsed_response: + error_message = parsed_response["message"] logger.debug(f"{debug_message} Message: {error_message}") + raise RuntimeError(error_message) from exc + logger.debug(debug_message) raise exc diff --git a/src/firebolt/common/token_storage.py b/src/firebolt/common/token_storage.py new file mode 100644 index 00000000000..a77159347aa --- /dev/null +++ b/src/firebolt/common/token_storage.py @@ -0,0 +1,114 @@ +import base64 +import json +import os +from json import JSONDecodeError +from typing import Optional + +from appdirs import user_data_dir +from cryptography.fernet import Fernet, InvalidToken +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC + + +class TokenSecureStorage: + def __init__(self, username: str, password: str): + """ + Class for permanent storage of token in the filesystem in encrypted way + + :param username: username used for toke encryption + :param password: password used for toke encryption + """ + self._data_dir = user_data_dir(appname="firebolt") + os.makedirs(self._data_dir, exist_ok=True) + + self._token_file = os.path.join(self._data_dir, "token.json") + + self.salt = self._get_salt() + self.encrypter = FernetEncrypter(self.salt, username, password) + + def _get_salt(self) -> str: + """ + Get salt from the file if exists, or generate a new one + + :return: salt + """ + res = self._read_data_json() + if "salt" not in res: + return FernetEncrypter.generate_salt() + else: + return res["salt"] + + def _read_data_json(self) -> dict: + """ + Read json file token.json + + :return: json object as dict + """ + if not os.path.exists(self._token_file): + return {} + + with open(self._token_file) as f: + try: + return json.load(f) + except JSONDecodeError: + return {} + + def get_cached_token(self) -> Optional[str]: + """ + Get decrypted token using username and password + If the token not found or token cannot be decrypted using username, password + None will be returned + + :return: token or None + """ + res = self._read_data_json() + if "token" not in res: + return None + + return self.encrypter.decrypt(res["token"]) + + def cache_token(self, token: str) -> None: + """ + + :param token: + :return: + """ + token = self.encrypter.encrypt(token) + + with open(self._token_file, "w") as f: + json.dump({"token": token, "salt": self.salt}, f) + + +class FernetEncrypter: + def __init__(self, salt: str, username: str, password: str): + """ + + :param salt: + :param username: + :param password: + """ + + kdf = PBKDF2HMAC( + algorithm=hashes.SHA256(), + salt=base64.b64decode(salt), + length=32, + iterations=39000, + ) + self.fernet = Fernet( + base64.urlsafe_b64encode( + kdf.derive(bytes(f"{username}{password}", encoding="utf-8")) + ) + ) + + @staticmethod + def generate_salt() -> str: + return base64.b64encode(os.urandom(16)).decode("ascii") + + def encrypt(self, data: str) -> str: + return self.fernet.encrypt(bytes(data, encoding="utf-8")).decode("utf-8") + + def decrypt(self, data: str) -> Optional[str]: + try: + return self.fernet.decrypt(bytes(data, encoding="utf-8")).decode("utf-8") + except InvalidToken: + return None diff --git a/tests/unit/client/test_auth.py b/tests/unit/client/test_auth.py index e2d663ada0f..a5e32261c81 100644 --- a/tests/unit/client/test_auth.py +++ b/tests/unit/client/test_auth.py @@ -2,6 +2,7 @@ import pytest from httpx import Client, Request, StreamError, codes +from pyfakefs.fake_filesystem import FakeFilesystem from pytest_httpx import HTTPXMock from pytest_mock import MockerFixture @@ -30,9 +31,7 @@ def test_auth_basic( def test_auth_refresh_on_expiration( - httpx_mock: HTTPXMock, - test_token: str, - test_token2: str, + httpx_mock: HTTPXMock, test_token: str, test_token2: str, fs: FakeFilesystem ): """Auth refreshes the token on expiration.""" @@ -56,9 +55,7 @@ def test_auth_refresh_on_expiration( def test_auth_uses_same_token_if_valid( - httpx_mock: HTTPXMock, - test_token: str, - test_token2: str, + httpx_mock: HTTPXMock, test_token: str, test_token2: str, fs: FakeFilesystem ): """Auth refreshes the token on expiration""" @@ -92,7 +89,7 @@ def test_auth_uses_same_token_if_valid( httpx_mock.reset(False) -def test_auth_error_handling(httpx_mock: HTTPXMock): +def test_auth_error_handling(httpx_mock: HTTPXMock, fs: FakeFilesystem): """Auth handles various errors properly.""" for api_endpoint in ("https://host", "host"): diff --git a/tests/unit/client/test_auth_async.py b/tests/unit/client/test_auth_async.py index 23809b830fe..6ea3efd64ca 100644 --- a/tests/unit/client/test_auth_async.py +++ b/tests/unit/client/test_auth_async.py @@ -1,4 +1,5 @@ from httpx import AsyncClient, Request, codes +from pyfakefs.fake_filesystem import FakeFilesystem from pytest import mark from pytest_httpx import HTTPXMock @@ -8,9 +9,7 @@ @mark.asyncio async def test_auth_refresh_on_expiration( - httpx_mock: HTTPXMock, - test_token: str, - test_token2: str, + httpx_mock: HTTPXMock, test_token: str, test_token2: str, fs: FakeFilesystem ): """Auth refreshes the token on expiration.""" @@ -39,9 +38,7 @@ async def test_auth_refresh_on_expiration( @mark.asyncio async def test_auth_uses_same_token_if_valid( - httpx_mock: HTTPXMock, - test_token: str, - test_token2: str, + httpx_mock: HTTPXMock, test_token: str, test_token2: str, fs: FakeFilesystem ): """Auth refreshes the token on expiration""" @@ -81,8 +78,7 @@ async def test_auth_uses_same_token_if_valid( @mark.asyncio async def test_auth_adds_header( - httpx_mock: HTTPXMock, - test_token: str, + httpx_mock: HTTPXMock, test_token: str, fs: FakeFilesystem ): """Auth adds required authentication headers to httpx.Request.""" httpx_mock.add_response( diff --git a/tests/unit/client/test_client.py b/tests/unit/client/test_client.py index 2f51de549de..37101582d45 100644 --- a/tests/unit/client/test_client.py +++ b/tests/unit/client/test_client.py @@ -1,6 +1,7 @@ from typing import Callable from httpx import codes +from pyfakefs.fake_filesystem import FakeFilesystem from pytest import raises from pytest_httpx import HTTPXMock @@ -15,6 +16,7 @@ def test_client_retry( test_username: str, test_password: str, test_token: str, + fs: FakeFilesystem, ): """ Client retries with new auth token @@ -56,6 +58,7 @@ def test_client_different_auths( test_username: str, test_password: str, test_token: str, + fs: FakeFilesystem, ): """ Client properly handles such auth types: diff --git a/tests/unit/client/test_client_async.py b/tests/unit/client/test_client_async.py index d821a590fd1..b42a6fa1707 100644 --- a/tests/unit/client/test_client_async.py +++ b/tests/unit/client/test_client_async.py @@ -1,6 +1,7 @@ from typing import Callable from httpx import codes +from pyfakefs.fake_filesystem import FakeFilesystem from pytest import mark, raises from pytest_httpx import HTTPXMock @@ -16,6 +17,7 @@ async def test_client_retry( test_username: str, test_password: str, test_token: str, + fs: FakeFilesystem, ): """ Client retries with new auth token @@ -58,6 +60,7 @@ async def test_client_different_auths( test_username: str, test_password: str, test_token: str, + fs: FakeFilesystem, ): """ Client properly handles such auth types: diff --git a/tests/unit/common/test_token_storage.py b/tests/unit/common/test_token_storage.py new file mode 100644 index 00000000000..68931b499ad --- /dev/null +++ b/tests/unit/common/test_token_storage.py @@ -0,0 +1,104 @@ +import os + +from appdirs import user_config_dir +from pyfakefs.fake_filesystem import FakeFilesystem + +from firebolt.common.token_storage import FernetEncrypter, TokenSecureStorage + + +def test_encrypter_happy_path(): + """ + Simple encrypt/decrypt using FernetEncrypter + """ + salt = FernetEncrypter.generate_salt() + encrypter1 = FernetEncrypter(salt, username="username", password="password") + encrypter2 = FernetEncrypter(salt, username="username", password="password") + + token = "some string to encrypt" + encrypted_token = encrypter1.encrypt(token) + + assert token == encrypter2.decrypt(encrypted_token) + + +def test_encrypter_wrong_parameter(): + """ + Test that decryption only works, if the correct salt, + username and password is provided, otherwise None is returned + """ + salt1 = FernetEncrypter.generate_salt() + salt2 = FernetEncrypter.generate_salt() + + encrypter1 = FernetEncrypter(salt1, username="username", password="password") + + token = "some string to encrypt" + encrypted_token = encrypter1.encrypt(token) + + encrypter2 = FernetEncrypter(salt2, username="username", password="password") + assert encrypter2.decrypt(encrypted_token) is None + + encrypter2 = FernetEncrypter(salt1, username="username1", password="password") + assert encrypter2.decrypt(encrypted_token) is None + + encrypter2 = FernetEncrypter(salt1, username="username", password="password1") + assert encrypter2.decrypt(encrypted_token) is None + + encrypter2 = FernetEncrypter(salt1, username="username", password="password") + assert encrypter2.decrypt(encrypted_token) == token + + +def test_token_storage_happy_path(fs: FakeFilesystem): + """ + Test storage happy path cache token and get token + """ + settings = {"username": "username", "password": "password"} + assert TokenSecureStorage(**settings).get_cached_token() is None + + token = "some string to encrypt" + TokenSecureStorage(**settings).cache_token(token) + + assert token == TokenSecureStorage(**settings).get_cached_token() + token = "some new string to encrypt" + + TokenSecureStorage(**settings).cache_token(token) + assert ( + token + == TokenSecureStorage( + username="username", password="password" + ).get_cached_token() + ) + + +def test_token_storage_wrong_parameter(fs: FakeFilesystem): + """ + Test getting token with different username or password + """ + settings = {"username": "username", "password": "password"} + token = "some string to encrypt" + TokenSecureStorage(**settings).cache_token(token) + + assert ( + TokenSecureStorage( + username="username", password="wrong_password" + ).get_cached_token() + is None + ) + assert ( + TokenSecureStorage( + username="wrong_username", password="password" + ).get_cached_token() + is None + ) + assert TokenSecureStorage(**settings).get_cached_token() == token + + +def test_token_storage_json_broken(fs: FakeFilesystem): + """ + Check that the TokenSecureStorage properly handles broken json + """ + settings = {"username": "username", "password": "password"} + + data_dir = os.path.join(user_config_dir(), "firebolt") + fs.create_dir(data_dir) + fs.create_file(os.path.join(data_dir, "token.json"), contents="{Not a valid json") + + assert TokenSecureStorage(**settings).get_cached_token() is None From cc9a72b5e0f6de8484345fc475fabeb9cc396924 Mon Sep 17 00:00:00 2001 From: Yury Dzerin Date: Wed, 23 Feb 2022 16:54:48 +0100 Subject: [PATCH 02/13] Add missing dependency --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index 314d1e2f0b6..fb9d51242f7 100755 --- a/setup.cfg +++ b/setup.cfg @@ -24,6 +24,7 @@ project_urls = packages = find: install_requires = aiorwlock==1.1.0 + appdirs>=1.4.4 async-property==0.2.1 httpx[http2]==0.21.3 pydantic[dotenv]==1.8.2 From 12d215282c6abe96d91bcf85e5f6dfed14672596 Mon Sep 17 00:00:00 2001 From: Yury Dzerin Date: Wed, 23 Feb 2022 16:57:13 +0100 Subject: [PATCH 03/13] Add missing cryptography --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index fb9d51242f7..aab5fbde079 100755 --- a/setup.cfg +++ b/setup.cfg @@ -26,6 +26,7 @@ install_requires = aiorwlock==1.1.0 appdirs>=1.4.4 async-property==0.2.1 + cryptography>=36.0.1 httpx[http2]==0.21.3 pydantic[dotenv]==1.8.2 readerwriterlock==1.0.9 From 5bbe3415c3ab8b21b75248aa2bb160149f2b9865 Mon Sep 17 00:00:00 2001 From: Yury Dzerin Date: Wed, 23 Feb 2022 17:01:57 +0100 Subject: [PATCH 04/13] Add appdirs stubs --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index aab5fbde079..b79450e79f1 100755 --- a/setup.cfg +++ b/setup.cfg @@ -25,6 +25,7 @@ packages = find: install_requires = aiorwlock==1.1.0 appdirs>=1.4.4 + appdirs-stubs>=0.1.0 async-property==0.2.1 cryptography>=36.0.1 httpx[http2]==0.21.3 From 9eefe8a7bc757e4f16705599a2addc81f76a2bf2 Mon Sep 17 00:00:00 2001 From: Yury Dzerin Date: Wed, 23 Feb 2022 17:03:59 +0100 Subject: [PATCH 05/13] Add missing pyfakefs for testing --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index b79450e79f1..fdf219d51ea 100755 --- a/setup.cfg +++ b/setup.cfg @@ -47,6 +47,7 @@ dev = devtools==0.7.0 mypy==0.910 pre-commit==2.15.0 + pyfakefs>=4.5.3 pytest==6.2.5 pytest-asyncio pytest-cov==3.0.0 From 168fd27ae037e8de11ccb8cf0e6ad79bc90b5de6 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Thu, 3 Mar 2022 13:15:25 +0100 Subject: [PATCH 06/13] remove redundant code check --- .github/workflows/code-check.yml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.github/workflows/code-check.yml b/.github/workflows/code-check.yml index 759ebfa992e..e474926b1c2 100644 --- a/.github/workflows/code-check.yml +++ b/.github/workflows/code-check.yml @@ -26,7 +26,3 @@ jobs: uses: pre-commit/action@v2.0.3 with: extra_args: --all-files - - - name: Type check with mypy - run: | - mypy src From 9a3edb8b75cebabb404cbf23a797212a391faec0 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Thu, 3 Mar 2022 13:23:33 +0100 Subject: [PATCH 07/13] return UNAUTHORIZED error handling --- src/firebolt/client/resource_manager_hooks.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/src/firebolt/client/resource_manager_hooks.py b/src/firebolt/client/resource_manager_hooks.py index c4318530804..1682bbb1fc7 100644 --- a/src/firebolt/client/resource_manager_hooks.py +++ b/src/firebolt/client/resource_manager_hooks.py @@ -1,7 +1,7 @@ from json import JSONDecodeError from logging import getLogger -from httpx import HTTPStatusError, Request, RequestError, Response, codes +from httpx import HTTPStatusError, Request, RequestError, Response logger = getLogger(__name__) @@ -40,24 +40,14 @@ def raise_on_4xx_5xx(response: Response) -> None: parsed_response = exc.response.json() except JSONDecodeError: parsed_response = {"_raw": exc.response.text} - - if exc.response.status_code == codes.UNAUTHORIZED: - # Do not raise an exception on UNAUTHORIZED - # It should be responsibility of authentication mechanism - return - debug_message = ( f"Error response {exc.response.status_code} " f"while requesting {exc.request.url!r}. " f"Response: {parsed_response}. " ) - if "message" in parsed_response: - error_message = parsed_response["message"] logger.debug(f"{debug_message} Message: {error_message}") - raise RuntimeError(error_message) from exc - logger.debug(debug_message) raise exc From cf413845464e35397bd31f4ed0e681bf0dcc5ceb Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Thu, 3 Mar 2022 13:34:09 +0100 Subject: [PATCH 08/13] minor imports refactoring --- src/firebolt/common/token_storage.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/firebolt/common/token_storage.py b/src/firebolt/common/token_storage.py index a77159347aa..9a24a34122c 100644 --- a/src/firebolt/common/token_storage.py +++ b/src/firebolt/common/token_storage.py @@ -1,6 +1,6 @@ -import base64 -import json -import os +from json import load as json_load, dump as json_dump +from base64 import urlsafe_b64encode, b64decode, b64encode +from os import urandom, makedirs, path from json import JSONDecodeError from typing import Optional @@ -19,9 +19,9 @@ def __init__(self, username: str, password: str): :param password: password used for toke encryption """ self._data_dir = user_data_dir(appname="firebolt") - os.makedirs(self._data_dir, exist_ok=True) + makedirs(self._data_dir, exist_ok=True) - self._token_file = os.path.join(self._data_dir, "token.json") + self._token_file = path.join(self._data_dir, "token.json") self.salt = self._get_salt() self.encrypter = FernetEncrypter(self.salt, username, password) @@ -44,12 +44,12 @@ def _read_data_json(self) -> dict: :return: json object as dict """ - if not os.path.exists(self._token_file): + if not path.exists(self._token_file): return {} with open(self._token_file) as f: try: - return json.load(f) + return json_load(f) except JSONDecodeError: return {} @@ -76,7 +76,7 @@ def cache_token(self, token: str) -> None: token = self.encrypter.encrypt(token) with open(self._token_file, "w") as f: - json.dump({"token": token, "salt": self.salt}, f) + json_dump({"token": token, "salt": self.salt}, f) class FernetEncrypter: @@ -90,19 +90,19 @@ def __init__(self, salt: str, username: str, password: str): kdf = PBKDF2HMAC( algorithm=hashes.SHA256(), - salt=base64.b64decode(salt), + salt=b64decode(salt), length=32, iterations=39000, ) self.fernet = Fernet( - base64.urlsafe_b64encode( + urlsafe_b64encode( kdf.derive(bytes(f"{username}{password}", encoding="utf-8")) ) ) @staticmethod def generate_salt() -> str: - return base64.b64encode(os.urandom(16)).decode("ascii") + return b64encode(urandom(16)).decode("ascii") def encrypt(self, data: str) -> str: return self.fernet.encrypt(bytes(data, encoding="utf-8")).decode("utf-8") From e5c8a96018673b1011e49a4f56062bd2583a4472 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Thu, 3 Mar 2022 13:49:00 +0100 Subject: [PATCH 09/13] minor refactoring --- src/firebolt/common/token_storage.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/src/firebolt/common/token_storage.py b/src/firebolt/common/token_storage.py index 9a24a34122c..93810f96c2a 100644 --- a/src/firebolt/common/token_storage.py +++ b/src/firebolt/common/token_storage.py @@ -1,7 +1,6 @@ -from json import load as json_load, dump as json_dump +from json import load as json_load, dump as json_dump, JSONDecodeError from base64 import urlsafe_b64encode, b64decode, b64encode from os import urandom, makedirs, path -from json import JSONDecodeError from typing import Optional from appdirs import user_data_dir @@ -10,6 +9,13 @@ from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC +def generate_salt() -> str: + return b64encode(urandom(16)).decode("ascii") + + +APPNAME = "firebolt" + + class TokenSecureStorage: def __init__(self, username: str, password: str): """ @@ -18,7 +24,7 @@ def __init__(self, username: str, password: str): :param username: username used for toke encryption :param password: password used for toke encryption """ - self._data_dir = user_data_dir(appname="firebolt") + self._data_dir = user_data_dir(appname=APPNAME) makedirs(self._data_dir, exist_ok=True) self._token_file = path.join(self._data_dir, "token.json") @@ -33,10 +39,7 @@ def _get_salt(self) -> str: :return: salt """ res = self._read_data_json() - if "salt" not in res: - return FernetEncrypter.generate_salt() - else: - return res["salt"] + return res.get("salt", generate_salt()) def _read_data_json(self) -> dict: """ @@ -100,10 +103,6 @@ def __init__(self, salt: str, username: str, password: str): ) ) - @staticmethod - def generate_salt() -> str: - return b64encode(urandom(16)).decode("ascii") - def encrypt(self, data: str) -> str: return self.fernet.encrypt(bytes(data, encoding="utf-8")).decode("utf-8") From 5272e58e9471e9bc95c5e4549ede07e1e03fdc49 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Thu, 3 Mar 2022 15:54:19 +0100 Subject: [PATCH 10/13] allow multiple tokens, add global fs mock for tests --- src/firebolt/common/token_storage.py | 16 +++++++-- tests/unit/common/test_token_storage.py | 47 +++++++++++++++++++------ tests/unit/conftest.py | 4 +-- 3 files changed, 52 insertions(+), 15 deletions(-) diff --git a/src/firebolt/common/token_storage.py b/src/firebolt/common/token_storage.py index 93810f96c2a..6f5039f12cf 100644 --- a/src/firebolt/common/token_storage.py +++ b/src/firebolt/common/token_storage.py @@ -1,5 +1,6 @@ from json import load as json_load, dump as json_dump, JSONDecodeError from base64 import urlsafe_b64encode, b64decode, b64encode +from hashlib import sha256 from os import urandom, makedirs, path from typing import Optional @@ -9,11 +10,18 @@ from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC +APPNAME = "firebolt" + + def generate_salt() -> str: return b64encode(urandom(16)).decode("ascii") -APPNAME = "firebolt" +def generate_file_name(username: str, password: str) -> str: + username_hash = sha256(username.encode("utf-8")).hexdigest()[:32] + password_hash = sha256(password.encode("utf-8")).hexdigest()[:32] + + return f"{username_hash}{password_hash}.json" class TokenSecureStorage: @@ -27,7 +35,9 @@ def __init__(self, username: str, password: str): self._data_dir = user_data_dir(appname=APPNAME) makedirs(self._data_dir, exist_ok=True) - self._token_file = path.join(self._data_dir, "token.json") + self._token_file = path.join( + self._data_dir, generate_file_name(username, password) + ) self.salt = self._get_salt() self.encrypter = FernetEncrypter(self.salt, username, password) @@ -43,7 +53,7 @@ def _get_salt(self) -> str: def _read_data_json(self) -> dict: """ - Read json file token.json + Read json token file :return: json object as dict """ diff --git a/tests/unit/common/test_token_storage.py b/tests/unit/common/test_token_storage.py index 68931b499ad..39f5b56f678 100644 --- a/tests/unit/common/test_token_storage.py +++ b/tests/unit/common/test_token_storage.py @@ -3,14 +3,18 @@ from appdirs import user_config_dir from pyfakefs.fake_filesystem import FakeFilesystem -from firebolt.common.token_storage import FernetEncrypter, TokenSecureStorage +from firebolt.common.token_storage import ( + FernetEncrypter, + TokenSecureStorage, + generate_salt, +) def test_encrypter_happy_path(): """ Simple encrypt/decrypt using FernetEncrypter """ - salt = FernetEncrypter.generate_salt() + salt = generate_salt() encrypter1 = FernetEncrypter(salt, username="username", password="password") encrypter2 = FernetEncrypter(salt, username="username", password="password") @@ -25,8 +29,8 @@ def test_encrypter_wrong_parameter(): Test that decryption only works, if the correct salt, username and password is provided, otherwise None is returned """ - salt1 = FernetEncrypter.generate_salt() - salt2 = FernetEncrypter.generate_salt() + salt1 = generate_salt() + salt2 = generate_salt() encrypter1 = FernetEncrypter(salt1, username="username", password="password") @@ -60,12 +64,7 @@ def test_token_storage_happy_path(fs: FakeFilesystem): token = "some new string to encrypt" TokenSecureStorage(**settings).cache_token(token) - assert ( - token - == TokenSecureStorage( - username="username", password="password" - ).get_cached_token() - ) + assert token == TokenSecureStorage(**settings).get_cached_token() def test_token_storage_wrong_parameter(fs: FakeFilesystem): @@ -102,3 +101,31 @@ def test_token_storage_json_broken(fs: FakeFilesystem): fs.create_file(os.path.join(data_dir, "token.json"), contents="{Not a valid json") assert TokenSecureStorage(**settings).get_cached_token() is None + + +def test_multiple_tokens(fs: FakeFilesystem) -> None: + """ + Check that the TokenSecureStorage properly handles multiple tokens hashed + """ + settings1 = {"username": "username1", "password": "password1"} + settings2 = {"username": "username2", "password": "password2"} + token1 = "token1" + token2 = "token2" + token3 = "token3" + + st1 = TokenSecureStorage(**settings1) + st2 = TokenSecureStorage(**settings2) + + st1.cache_token(token1) + + assert st1.get_cached_token() == token1 + assert st2.get_cached_token() is None + + st2.cache_token(token2) + + assert st1.get_cached_token() == token1 + assert st2.get_cached_token() == token2 + + st1.cache_token(token3) + assert st1.get_cached_token() == token3 + assert st2.get_cached_token() == token2 diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 608e62ffaed..704448604ca 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -108,7 +108,7 @@ def do_mock( assert request.url == auth_url return Response( status_code=httpx.codes.OK, - json={"access_token": "", "expires_in": 2 ** 32}, + json={"access_token": "", "expires_in": 2**32}, ) return do_mock @@ -284,7 +284,7 @@ def check_credentials( return Response( status_code=httpx.codes.OK, - json={"expires_in": 2 ** 32, "access_token": access_token}, + json={"expires_in": 2**32, "access_token": access_token}, ) return check_credentials From 1a85c75726277bfa66ad07787b692826c4274c4c Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Thu, 3 Mar 2022 16:21:54 +0100 Subject: [PATCH 11/13] add token expiration handlin --- src/firebolt/client/auth.py | 2 +- src/firebolt/common/token_storage.py | 11 ++++++++-- tests/unit/common/test_token_storage.py | 27 +++++++++++++++++++------ 3 files changed, 31 insertions(+), 9 deletions(-) diff --git a/src/firebolt/client/auth.py b/src/firebolt/client/auth.py index 78697601e00..add0ab02ae2 100644 --- a/src/firebolt/client/auth.py +++ b/src/firebolt/client/auth.py @@ -80,7 +80,7 @@ def get_new_token_generator(self) -> Generator[Request, Response, None]: self._token = parsed["access_token"] self._expires = int(time()) + int(parsed["expires_in"]) - self._token_storage.cache_token(parsed["access_token"]) + self._token_storage.cache_token(parsed["access_token"], self._expires) except _REQUEST_ERRORS as e: raise AuthenticationError(repr(e), self._api_endpoint) diff --git a/src/firebolt/common/token_storage.py b/src/firebolt/common/token_storage.py index 6f5039f12cf..a083a435228 100644 --- a/src/firebolt/common/token_storage.py +++ b/src/firebolt/common/token_storage.py @@ -3,6 +3,7 @@ from hashlib import sha256 from os import urandom, makedirs, path from typing import Optional +from time import time from appdirs import user_data_dir from cryptography.fernet import Fernet, InvalidToken @@ -78,9 +79,13 @@ def get_cached_token(self) -> Optional[str]: if "token" not in res: return None + # Ignore expired tokens + if "expiration" in res and res["expiration"] <= int(time()): + return None + return self.encrypter.decrypt(res["token"]) - def cache_token(self, token: str) -> None: + def cache_token(self, token: str, expiration_ts: int) -> None: """ :param token: @@ -89,7 +94,9 @@ def cache_token(self, token: str) -> None: token = self.encrypter.encrypt(token) with open(self._token_file, "w") as f: - json_dump({"token": token, "salt": self.salt}, f) + json_dump( + {"token": token, "salt": self.salt, "expiration": expiration_ts}, f + ) class FernetEncrypter: diff --git a/tests/unit/common/test_token_storage.py b/tests/unit/common/test_token_storage.py index 39f5b56f678..22fd4bceb44 100644 --- a/tests/unit/common/test_token_storage.py +++ b/tests/unit/common/test_token_storage.py @@ -2,6 +2,7 @@ from appdirs import user_config_dir from pyfakefs.fake_filesystem import FakeFilesystem +from unittest.mock import patch from firebolt.common.token_storage import ( FernetEncrypter, @@ -50,6 +51,7 @@ def test_encrypter_wrong_parameter(): assert encrypter2.decrypt(encrypted_token) == token +@patch("firebolt.common.token_storage.time", return_value=0) def test_token_storage_happy_path(fs: FakeFilesystem): """ Test storage happy path cache token and get token @@ -58,22 +60,23 @@ def test_token_storage_happy_path(fs: FakeFilesystem): assert TokenSecureStorage(**settings).get_cached_token() is None token = "some string to encrypt" - TokenSecureStorage(**settings).cache_token(token) + TokenSecureStorage(**settings).cache_token(token, 1) assert token == TokenSecureStorage(**settings).get_cached_token() token = "some new string to encrypt" - TokenSecureStorage(**settings).cache_token(token) + TokenSecureStorage(**settings).cache_token(token, 1) assert token == TokenSecureStorage(**settings).get_cached_token() +@patch("firebolt.common.token_storage.time", return_value=0) def test_token_storage_wrong_parameter(fs: FakeFilesystem): """ Test getting token with different username or password """ settings = {"username": "username", "password": "password"} token = "some string to encrypt" - TokenSecureStorage(**settings).cache_token(token) + TokenSecureStorage(**settings).cache_token(token, 1) assert ( TokenSecureStorage( @@ -103,6 +106,7 @@ def test_token_storage_json_broken(fs: FakeFilesystem): assert TokenSecureStorage(**settings).get_cached_token() is None +@patch("firebolt.common.token_storage.time", return_value=0) def test_multiple_tokens(fs: FakeFilesystem) -> None: """ Check that the TokenSecureStorage properly handles multiple tokens hashed @@ -116,16 +120,27 @@ def test_multiple_tokens(fs: FakeFilesystem) -> None: st1 = TokenSecureStorage(**settings1) st2 = TokenSecureStorage(**settings2) - st1.cache_token(token1) + st1.cache_token(token1, 1) assert st1.get_cached_token() == token1 assert st2.get_cached_token() is None - st2.cache_token(token2) + st2.cache_token(token2, 1) assert st1.get_cached_token() == token1 assert st2.get_cached_token() == token2 - st1.cache_token(token3) + st1.cache_token(token3, 1) assert st1.get_cached_token() == token3 assert st2.get_cached_token() == token2 + + +@patch("firebolt.common.token_storage.time", return_value=0) +def test_expired_token(fs: FakeFilesystem) -> None: + """ + Check that TokenSecureStorage ignores expired tokens + """ + tss = TokenSecureStorage(username="username", password="password") + tss.cache_token("token", 0) + + assert tss.get_cached_token() is None From 73a4cfcfadec6c821ea35ca5f3cf65e0e4e636ed Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Thu, 3 Mar 2022 17:12:59 +0100 Subject: [PATCH 12/13] fix code checks --- src/firebolt/common/token_storage.py | 11 ++++++----- tests/unit/common/test_token_storage.py | 2 +- tests/unit/conftest.py | 4 ++-- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/firebolt/common/token_storage.py b/src/firebolt/common/token_storage.py index a083a435228..7643ac326df 100644 --- a/src/firebolt/common/token_storage.py +++ b/src/firebolt/common/token_storage.py @@ -1,16 +1,17 @@ -from json import load as json_load, dump as json_dump, JSONDecodeError -from base64 import urlsafe_b64encode, b64decode, b64encode +from base64 import b64decode, b64encode, urlsafe_b64encode from hashlib import sha256 -from os import urandom, makedirs, path -from typing import Optional +from json import JSONDecodeError +from json import dump as json_dump +from json import load as json_load +from os import makedirs, path, urandom from time import time +from typing import Optional from appdirs import user_data_dir from cryptography.fernet import Fernet, InvalidToken from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC - APPNAME = "firebolt" diff --git a/tests/unit/common/test_token_storage.py b/tests/unit/common/test_token_storage.py index 22fd4bceb44..cefeb00bbf2 100644 --- a/tests/unit/common/test_token_storage.py +++ b/tests/unit/common/test_token_storage.py @@ -1,8 +1,8 @@ import os +from unittest.mock import patch from appdirs import user_config_dir from pyfakefs.fake_filesystem import FakeFilesystem -from unittest.mock import patch from firebolt.common.token_storage import ( FernetEncrypter, diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 704448604ca..608e62ffaed 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -108,7 +108,7 @@ def do_mock( assert request.url == auth_url return Response( status_code=httpx.codes.OK, - json={"access_token": "", "expires_in": 2**32}, + json={"access_token": "", "expires_in": 2 ** 32}, ) return do_mock @@ -284,7 +284,7 @@ def check_credentials( return Response( status_code=httpx.codes.OK, - json={"expires_in": 2**32, "access_token": access_token}, + json={"expires_in": 2 ** 32, "access_token": access_token}, ) return check_credentials From b9dada1aba99fea20581afcd6882e4e522582134 Mon Sep 17 00:00:00 2001 From: Stepan Burlakov Date: Fri, 4 Mar 2022 10:37:42 +0100 Subject: [PATCH 13/13] patch fs for each test --- tests/conftest.py | 8 ++++++++ 1 file changed, 8 insertions(+) create mode 100644 tests/conftest.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000000..0aed98b7256 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,8 @@ +from pyfakefs.fake_filesystem_unittest import Patcher +from pytest import fixture + + +@fixture(autouse=True) +def global_fake_fs() -> None: + with Patcher(): + yield