From feb70ffbb579bd4026e43913e2ff489a1affa270 Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Fri, 26 Aug 2022 14:54:52 -0400 Subject: [PATCH 1/5] Refactoring auth code --- src/sparsezoo/utils/authentication.py | 126 +++++++------------ tests/sparsezoo/utils/test_authentication.py | 86 +++++++++++++ 2 files changed, 131 insertions(+), 81 deletions(-) create mode 100644 tests/sparsezoo/utils/test_authentication.py diff --git a/src/sparsezoo/utils/authentication.py b/src/sparsezoo/utils/authentication.py index 8765fe44..329e9e02 100644 --- a/src/sparsezoo/utils/authentication.py +++ b/src/sparsezoo/utils/authentication.py @@ -45,77 +45,8 @@ ) -class SparseZooCredentials: - """ - Class wrapping around the sparse zoo credentials file. - """ - - def __init__(self): - if os.path.exists(CREDENTIALS_YAML): - _LOGGER.debug(f"Loading sparse zoo credentials from {CREDENTIALS_YAML}") - with open(CREDENTIALS_YAML) as credentials_file: - credentials_yaml = yaml.safe_load(credentials_file) - if credentials_yaml and CREDENTIALS_YAML_TOKEN_KEY in credentials_yaml: - self._token = credentials_yaml[CREDENTIALS_YAML_TOKEN_KEY]["token"] - self._created = credentials_yaml[CREDENTIALS_YAML_TOKEN_KEY][ - "created" - ] - else: - self._token = None - self._created = None - else: - _LOGGER.debug( - f"No sparse zoo credentials files found at {CREDENTIALS_YAML}" - ) - self._token = None - self._created = None - - def save_token(self, token: str, created: float): - """ - Save the jwt for accessing sparse zoo APIs. Will create the credentials file - if it does not exist already. - - :param token: the jwt for accessing sparse zoo APIs - :param created: the approximate time the token was created - """ - _LOGGER.debug(f"Saving sparse zoo credentials at {CREDENTIALS_YAML}") - if not os.path.exists(CREDENTIALS_YAML): - create_parent_dirs(CREDENTIALS_YAML) - with open(CREDENTIALS_YAML, "w+") as credentials_file: - credentials_yaml = yaml.safe_load(credentials_file) - if credentials_yaml is None: - credentials_yaml = {} - credentials_yaml[CREDENTIALS_YAML_TOKEN_KEY] = { - "token": token, - "created": created, - } - self._token = token - self._created = created - - yaml.safe_dump(credentials_yaml, credentials_file) - - @property - def token(self): - """ - :return: obtain the token if under 1 day old, else return None - """ - _LOGGER.debug(f"Obtaining sparse zoo credentials from {CREDENTIALS_YAML}") - if self._token and self._created is not None: - creation_date = datetime.fromtimestamp(self._created, tz=timezone.utc) - creation_difference = datetime.now(tz=timezone.utc) - creation_date - if creation_difference.days < 30: - return self._token - else: - _LOGGER.debug(f"Expired sparse zoo credentials at {CREDENTIALS_YAML}") - return None - else: - _LOGGER.debug(f"No sparse zoo credentials found at {CREDENTIALS_YAML}") - return None - - def get_auth_header( - authentication_type: str = PUBLIC_AUTH_TYPE, - force_token_refresh: bool = False, + force_token_refresh: bool = False, path: str = CREDENTIALS_YAML ) -> Dict: """ Obtain an authentication header token from either credentials file or from APIs @@ -124,24 +55,57 @@ def get_auth_header( Currently only 'public' authentication type is supported. - :param authentication_type: authentication type for generating token :param force_token_refresh: forces a new token to be generated :return: An authentication header with key 'nm-token-header' containing the header token """ - credentials = SparseZooCredentials() - token = credentials.token - if token and not force_token_refresh: - return {NM_TOKEN_HEADER: token} - elif authentication_type.lower() == PUBLIC_AUTH_TYPE: + token = _maybe_load_token(path) + if token is None or force_token_refresh: _LOGGER.info("Obtaining new sparse zoo credentials token") - created = time.time() response = requests.post( url=AUTH_API, data=json.dumps({"authentication_type": PUBLIC_AUTH_TYPE}) ) response.raise_for_status() token = response.json()["token"] - credentials.save_token(token, created) - return {NM_TOKEN_HEADER: token} - else: - raise Exception(f"Authentication type {PUBLIC_AUTH_TYPE} not supported.") + created = time.time() + _save_token(token, created, path) + return {NM_TOKEN_HEADER: token} + + +def _maybe_load_token(path: str): + if not os.path.exists(path): + return None + + with open(path) as fp: + creds = yaml.safe_load(fp) + + if creds is None or CREDENTIALS_YAML_TOKEN_KEY not in creds: + return None + + info = creds[CREDENTIALS_YAML_TOKEN_KEY] + if "token" not in info or "created" not in info: + return None + + date_created = datetime.fromtimestamp(info["created"], tz=timezone.utc) + creation_difference = datetime.now(tz=timezone.utc) - date_created + + if creation_difference.days > 30: + return None + + return info["token"] + + +def _save_token(token: str, created: float, path: str): + """ + Save the jwt for accessing sparse zoo APIs. Will create the credentials file + if it does not exist already. + + :param token: the jwt for accessing sparse zoo APIs + :param created: the approximate time the token was created + """ + _LOGGER.debug(f"Saving sparse zoo credentials at {CREDENTIALS_YAML}") + if not os.path.exists(path): + create_parent_dirs(path) + with open(path, "w+") as fp: + auth = {CREDENTIALS_YAML_TOKEN_KEY: dict(token=token, created=created)} + yaml.safe_dump(auth, fp) diff --git a/tests/sparsezoo/utils/test_authentication.py b/tests/sparsezoo/utils/test_authentication.py new file mode 100644 index 00000000..2e7b459f --- /dev/null +++ b/tests/sparsezoo/utils/test_authentication.py @@ -0,0 +1,86 @@ +# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from sparsezoo.utils.authentication import ( + get_auth_header, + _maybe_load_token, + _save_token, + CREDENTIALS_YAML_TOKEN_KEY, + NM_TOKEN_HEADER, +) +import pytest +from datetime import datetime, timedelta +import yaml +from unittest.mock import patch, MagicMock + + +def test_load_token_no_path(tmp_path): + path = str(tmp_path / "token.yaml") + assert _maybe_load_token(path) is None + + +def test_load_token_yaml_fail(tmp_path): + path = str(tmp_path / "token.yaml") + with open(path, "w") as fp: + fp.write("asdf") + assert _maybe_load_token(path) is None + + +@pytest.mark.parametrize( + "content", + [ + {}, + {CREDENTIALS_YAML_TOKEN_KEY: {}}, + {CREDENTIALS_YAML_TOKEN_KEY: {"token": "asdf"}}, + {CREDENTIALS_YAML_TOKEN_KEY: {"created": "asdf"}}, + { + CREDENTIALS_YAML_TOKEN_KEY: { + "created": (datetime.now() - timedelta(days=40)).timestamp() + } + }, + ], +) +def test_load_token_failure_cases(tmp_path, content): + path = str(tmp_path / "token.yaml") + with open(path, "w") as fp: + yaml.dump(content, fp) + assert _maybe_load_token(path) is None + + +def test_load_token_valid(tmp_path): + auth = { + CREDENTIALS_YAML_TOKEN_KEY: { + "created": datetime.now().timestamp(), + "token": "asdf", + } + } + path = str(tmp_path / "token.yaml") + with open(path, "w") as fp: + yaml.dump(auth, fp) + assert _maybe_load_token(path) == "asdf" + + +def test_load_saved_token(tmp_path): + path = str(tmp_path / "some" / "dirs" / "token.yaml") + _save_token("asdf", datetime.now().timestamp(), path) + assert _maybe_load_token(path) == "asdf" + + +@patch("requests.post", return_value=MagicMock(json=lambda: {"token": "qwer"})) +def test_get_auth_token(post_mock, tmp_path): + path = tmp_path / "creds.yaml" + assert not path.exists() + assert get_auth_header(path=str(path)) == {NM_TOKEN_HEADER: "qwer"} + assert path.exists() + post_mock.assert_called() From f59f939402c7c0f5fdca074142cb06cf284ffa0f Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Fri, 26 Aug 2022 14:58:29 -0400 Subject: [PATCH 2/5] Adding debug logs back into auth --- src/sparsezoo/utils/authentication.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/sparsezoo/utils/authentication.py b/src/sparsezoo/utils/authentication.py index 329e9e02..d223e705 100644 --- a/src/sparsezoo/utils/authentication.py +++ b/src/sparsezoo/utils/authentication.py @@ -74,22 +74,28 @@ def get_auth_header( def _maybe_load_token(path: str): if not os.path.exists(path): + _LOGGER.debug(f"No sparse zoo credentials files found at {path}") return None + _LOGGER.debug(f"Loading sparse zoo credentials from {path}") + with open(path) as fp: creds = yaml.safe_load(fp) if creds is None or CREDENTIALS_YAML_TOKEN_KEY not in creds: + _LOGGER.debug(f"No sparse zoo credentials found at {path}") return None info = creds[CREDENTIALS_YAML_TOKEN_KEY] if "token" not in info or "created" not in info: + _LOGGER.debug(f"No sparse zoo credentials found at {path}") return None date_created = datetime.fromtimestamp(info["created"], tz=timezone.utc) creation_difference = datetime.now(tz=timezone.utc) - date_created if creation_difference.days > 30: + _LOGGER.debug(f"Expired sparse zoo credentials at {path}") return None return info["token"] From f90b33f934a748bd23d0fb2d578cb704e318edfa Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Fri, 26 Aug 2022 15:02:25 -0400 Subject: [PATCH 3/5] Minifying tests --- tests/sparsezoo/utils/test_authentication.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/tests/sparsezoo/utils/test_authentication.py b/tests/sparsezoo/utils/test_authentication.py index 2e7b459f..6ec085bf 100644 --- a/tests/sparsezoo/utils/test_authentication.py +++ b/tests/sparsezoo/utils/test_authentication.py @@ -1,11 +1,11 @@ # Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved. -# +# # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at -# +# # http://www.apache.org/licenses/LICENSE-2.0 -# +# # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. @@ -37,6 +37,9 @@ def test_load_token_yaml_fail(tmp_path): assert _maybe_load_token(path) is None +_OLD_DATE = (datetime.now() - timedelta(days=40)).timestamp() + + @pytest.mark.parametrize( "content", [ @@ -44,11 +47,7 @@ def test_load_token_yaml_fail(tmp_path): {CREDENTIALS_YAML_TOKEN_KEY: {}}, {CREDENTIALS_YAML_TOKEN_KEY: {"token": "asdf"}}, {CREDENTIALS_YAML_TOKEN_KEY: {"created": "asdf"}}, - { - CREDENTIALS_YAML_TOKEN_KEY: { - "created": (datetime.now() - timedelta(days=40)).timestamp() - } - }, + {CREDENTIALS_YAML_TOKEN_KEY: {"created": _OLD_DATE}}, ], ) def test_load_token_failure_cases(tmp_path, content): From 5ae4d18253f6cf771a543fb00c6f1a9bbba2bb12 Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Fri, 9 Sep 2022 11:32:48 -0400 Subject: [PATCH 4/5] Fixing style --- tests/sparsezoo/utils/test_authentication.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tests/sparsezoo/utils/test_authentication.py b/tests/sparsezoo/utils/test_authentication.py index 6ec085bf..aff526d4 100644 --- a/tests/sparsezoo/utils/test_authentication.py +++ b/tests/sparsezoo/utils/test_authentication.py @@ -12,17 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +from datetime import datetime, timedelta +from unittest.mock import MagicMock, patch + +import pytest +import yaml + from sparsezoo.utils.authentication import ( - get_auth_header, - _maybe_load_token, - _save_token, CREDENTIALS_YAML_TOKEN_KEY, NM_TOKEN_HEADER, + _maybe_load_token, + _save_token, + get_auth_header, ) -import pytest -from datetime import datetime, timedelta -import yaml -from unittest.mock import patch, MagicMock def test_load_token_no_path(tmp_path): From 42059a420801ffd1bcc68864a299fa4767b1b6f5 Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Fri, 9 Sep 2022 11:33:28 -0400 Subject: [PATCH 5/5] Fixing docstring of get_auth_header --- src/sparsezoo/utils/authentication.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sparsezoo/utils/authentication.py b/src/sparsezoo/utils/authentication.py index d223e705..5694cdae 100644 --- a/src/sparsezoo/utils/authentication.py +++ b/src/sparsezoo/utils/authentication.py @@ -50,7 +50,7 @@ def get_auth_header( ) -> Dict: """ Obtain an authentication header token from either credentials file or from APIs - if token is over 1 day old. Location of credentials file can be changed by setting + if token is over 30 day old. Location of credentials file can be changed by setting the environment variable `NM_SPARSE_ZOO_CREDENTIALS`. Currently only 'public' authentication type is supported.