From d371710648c0fdead26da8bbba0fd437b8b677f0 Mon Sep 17 00:00:00 2001
From: nh916 <navidhariri25@gmail.com>
Date: Wed, 20 Sep 2023 16:37:34 -0700
Subject: [PATCH] refactored `get_s3_client` into a function in a
 aws_s3_utils.py file

* removed some bulk from the `cript.API` class
* wrote docstrings for `get_s3_client`
* tested the code and it is working fine
* formatted the whole thing with black
* changed log to have `file` be capitalized for better UI
---
 src/cript/api/api.py                | 25 ++++++------------
 src/cript/api/utils/aws_s3_utils.py | 40 +++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 17 deletions(-)
 create mode 100644 src/cript/api/utils/aws_s3_utils.py

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