diff --git a/src/cript/api/api.py b/src/cript/api/api.py index 41b4e5bc2..3615818a4 100644 --- a/src/cript/api/api.py +++ b/src/cript/api/api.py @@ -26,6 +26,7 @@ from cript.api.paginator import Paginator from cript.api.utils.get_host_token import resolve_host_and_token from cript.api.utils.helper_functions import _get_node_type_from_json +from cript.api.utils.aws_s3_utils import get_s3_client from cript.api.utils.save_helper import ( _fix_node_save, _get_uuid_from_error_message, @@ -347,29 +348,19 @@ def _prepare_host(self, host: str) -> str: @property def _s3_client(self) -> boto3.client: # type: ignore """ - creates or returns a fully authenticated and ready s3 client + Property to use when wanting to interact with AWS S3. + + Gets a fully authenticated AWS S3 client if it was never created and stash it, + if the AWS S3 client has been created before, then returns the client that it has Returns ------- s3_client: boto3.client fully prepared and authenticated s3 client ready to be used throughout the script """ - if self._internal_s3_client is None: - auth = boto3.client("cognito-identity", region_name=self._REGION_NAME) - identity_id = auth.get_id(IdentityPoolId=self._IDENTITY_POOL_ID, Logins={self._COGNITO_LOGIN_PROVIDER: self._storage_token}) - # TODO remove this temporary fix to the token, by getting is from back end. - aws_token = self._storage_token - - aws_credentials = auth.get_credentials_for_identity(IdentityId=identity_id["IdentityId"], Logins={self._COGNITO_LOGIN_PROVIDER: aws_token}) - aws_credentials = aws_credentials["Credentials"] - s3_client = boto3.client( - "s3", - aws_access_key_id=aws_credentials["AccessKeyId"], - aws_secret_access_key=aws_credentials["SecretKey"], - aws_session_token=aws_credentials["SessionToken"], - ) - self._internal_s3_client = s3_client + self._internal_s3_client = get_s3_client(region_name=self._REGION_NAME, identity_pool_id=self._IDENTITY_POOL_ID, cognito_login_provider=self._COGNITO_LOGIN_PROVIDER, storage_token=self._storage_token) + return self._internal_s3_client def __enter__(self): @@ -869,7 +860,7 @@ def upload_file(self, file_path: Union[Path, str]) -> str: # upload file to AWS S3 self._s3_client.upload_file(Filename=file_path, Bucket=self._BUCKET_NAME, Key=object_name) # type: ignore - self.logger.info(f"Uploaded file: '{file_path}' to CRIPT storage") + self.logger.info(f"Uploaded File: '{file_path}' to CRIPT storage") # return the object_name within AWS S3 for easy retrieval return object_name diff --git a/src/cript/api/utils/aws_s3_utils.py b/src/cript/api/utils/aws_s3_utils.py new file mode 100644 index 000000000..a3d507d19 --- /dev/null +++ b/src/cript/api/utils/aws_s3_utils.py @@ -0,0 +1,40 @@ +from typing import Any + +import boto3 +from beartype import beartype + + +@beartype +def get_s3_client(region_name: str, identity_pool_id: str, cognito_login_provider: str, storage_token: str) -> Any: + """ + Creates an AWS S3 client and returns it to be used in the `cript.API` class. + + Parameters + ---------- + region_name: str + AWS S3 region name + identity_pool_id: str + AWS S3 identity pool id + cognito_login_provider: str + AWS S3 cognito login provider + storage_token: str + AWS S3 storage token gotten from the CRIPT frontend + + Returns + ------- + boto3.client + fully working AWS S3 client + """ + auth = boto3.client("cognito-identity", region_name=region_name) + identity_id = auth.get_id(IdentityPoolId=identity_pool_id, Logins={cognito_login_provider: storage_token}) + aws_token = storage_token + + aws_credentials = auth.get_credentials_for_identity(IdentityId=identity_id["IdentityId"], Logins={cognito_login_provider: aws_token}) + aws_credentials = aws_credentials["Credentials"] + s3_client = boto3.client( + "s3", + aws_access_key_id=aws_credentials["AccessKeyId"], + aws_secret_access_key=aws_credentials["SecretKey"], + aws_session_token=aws_credentials["SessionToken"], + ) + return s3_client