Merge pull request #64 from stac-labs/nested-secrets

Nested secrets & S3 secret loading

Author: ffaristocrat

src/stac_utils/aws.py

@@ -1,11 +1,15 @@
 import json
+import logging
 import os.path as op
 from tempfile import TemporaryDirectory
 
 import boto3
 from botocore.exceptions import ClientError
 
 
+logger = logging.getLogger(__name__)
+
+
 def get_secret(region_name: str, secret_name: str) -> dict:
     # Create a Secrets Manager client
     """
@@ -43,7 +47,18 @@ def write_secret(region_name: str, secret_name: str, secret: dict):
     )
 
 
-def load_from_s3(bucket: str, path: str, file_name: str) -> dict:
+def split_s3_url(url: str) -> tuple[str, str, str]:
+    prefix = "s3://"
+    if url.startswith(prefix):
+        url = url[len(prefix) :]
+
+    bucket, _, fpath = url.partition("/")
+    path, _, file_name = fpath.rpartition("/")
+
+    return bucket, path, file_name
+
+
+def load_from_s3(bucket: str, path: [str, None], file_name: str) -> dict:
     """
     Returns data from s3 given bucket, path, and file name
 
@@ -52,8 +67,9 @@ def load_from_s3(bucket: str, path: str, file_name: str) -> dict:
     :param file_name: Name of file to load
     :return: Data from specified file
     """
+    path = path or ""
     s3 = boto3.resource("s3").Bucket(bucket)
-    key = path.strip("/") + "/" + file_name
+    key = (path.strip("/") + "/" + file_name).lstrip("/")
 
     data = {}
     with TemporaryDirectory() as temp:
@@ -66,12 +82,13 @@ def load_from_s3(bucket: str, path: str, file_name: str) -> dict:
                 raise e
             data = {}
         except json.JSONDecodeError:
+            logger.warning(f"{key} is not a JSON file!")
             data = {}
 
     return data
 
 
-def save_to_s3(data: dict, bucket: str, path: str, file_name: str):
+def save_to_s3(data: dict, bucket: str, path: [str, None], file_name: str):
     """
     Saves data to s3 in specified location
 
@@ -81,8 +98,9 @@ def save_to_s3(data: dict, bucket: str, path: str, file_name: str):
     :param file_name: Desired file name
     :return: Data
     """
+    path = path or ""
     s3 = boto3.resource("s3").Bucket(bucket)
-    key = path.strip("/") + "/" + file_name
+    key = (path.strip("/") + "/" + file_name).lstrip("/")
 
     with TemporaryDirectory() as temp:
         temp_file = op.join(temp, file_name)

src/stac_utils/secret_context.py

@@ -2,17 +2,23 @@
 import os
 from unittest.mock import patch
 
-from .aws import get_secret
+from .aws import get_secret, split_s3_url, load_from_s3
 
 
 def secrets(
     file_name: str = None,
     secret_name: str = None,
     aws_region: str = None,
     dictionary: dict = None,
+    s3_url: str = None,
 ):
     """
-    Takes either a file, a Python dict, or an AWS secret in Secrets Manager and loads it all into the os.environ as a context.
+    Takes any combination of
+        * local JSON file
+        * Python dictionary
+        * AWS secret in Secrets Manager
+        * JSON file on S3
+    and loads it all into the os.environ as a context.
 
     Usage is typically:
 
@@ -23,8 +29,10 @@ def secrets(
     :param secret_name: Desired secret_name
     :param aws_region: Desired AWS region for secret
     :param dictionary: Specified dictionary
+    :param s3_url: S3 URL to JSON file
     :return:
     """
+
     values = {}
     if not secret_name and os.environ.get("SECRET_NAME"):
         secret_name = os.environ.get("SECRET_NAME")
@@ -39,12 +47,66 @@ def secrets(
                 secret_name or os.environ["SECRET_NAME"],
             )
         )
+    if not s3_url and os.environ.get("SECRET_S3_URL"):
+        s3_url = os.environ.get("SECRET_S3_URL")
+        # blank secret s3 url in the context, so it doesn't get loaded a second time
+        # if we nest secrets
+        values["SECRET_S3_URL"] = ""
+
+    if s3_url:
+        values.update(load_from_s3(*split_s3_url(s3_url)))
 
     if file_name:
         values.update(json.load(open(file_name, "rt")))
     if dictionary:
         values.update(dictionary)
 
     # the patcher doesn't like non-string keys OR values
-    values = {str(k): str(v) if v is not None else "" for k, v in values.items()}
+    values = {
+        str(k): safe_dump_json_to_string(v) if v is not None else ""
+        for k, v in values.items()
+    }
+
     return patch.dict(os.environ, values=values)
+
+
+def safe_dump_json_to_string(value: [list, dict, str, int, tuple, float, None]) -> str:
+    """Utility function to encode values to string, working through nested dictionaries & list"""
+
+    if type(value) in [dict]:
+        return json.dumps(
+            {str(k): safe_dump_json_to_string(v) for k, v in value.items()}
+        )
+
+    if type(value) in [list, tuple]:
+        return json.dumps([safe_dump_json_to_string(v) for v in value])
+
+    if value is None:
+        return "null"
+
+    return str(value)
+
+
+def safe_load_string_to_json(value: str) -> [list, dict, str, int, float, None]:
+    """Utility function to decode values from string, restoring all nested dictionaries & list"""
+
+    try:
+        loaded = json.loads(value)
+    except (json.decoder.JSONDecodeError, TypeError):
+        loaded = value
+
+    if type(loaded) in [dict]:
+        return {k: safe_load_string_to_json(v) for k, v in loaded.items()}
+
+    if type(loaded) in [list]:
+        return [safe_load_string_to_json(v) for v in loaded]
+
+    return loaded
+
+
+def get_env(key: str, default=None) -> [list, dict, str]:
+    """Utility function combining os.environ.get & safe_load_from_json"""
+
+    value = os.environ.get(key, default=default)
+
+    return safe_load_string_to_json(value)

src/tests/test_aws.py

@@ -1,9 +1,15 @@
 import json
 import unittest
-from unittest.mock import MagicMock, patch, call
+from unittest.mock import MagicMock, patch
 from botocore.exceptions import ClientError
 
-from src.stac_utils.aws import get_secret, write_secret, load_from_s3, save_to_s3
+from src.stac_utils.aws import (
+    get_secret,
+    write_secret,
+    load_from_s3,
+    save_to_s3,
+    split_s3_url,
+)
 
 
 class TestAWS(unittest.TestCase):
@@ -62,6 +68,27 @@ def test_load_from_s3(
         result_data = load_from_s3("foo", "bar", "spam")
         self.assertEqual(mock_data, result_data)
         mock_boto.return_value.Bucket.return_value.download_file.assert_called_once()
+        self.assertEqual(
+            mock_boto.return_value.Bucket.return_value.download_file.call_args[0][0],
+            "bar/spam",
+        )
+
+    @patch("json.load")
+    @patch("src.stac_utils.aws.open")
+    @patch("boto3.resource")
+    def test_load_from_s3_no_path(
+        self, mock_boto: MagicMock, mock_open: MagicMock, mock_load: MagicMock
+    ):
+        """Test load from s3 when path is None"""
+        mock_data = {"foo": "bar"}
+        mock_load.return_value = mock_data
+        result_data = load_from_s3("foo", None, "spam")
+        self.assertEqual(mock_data, result_data)
+        mock_boto.return_value.Bucket.return_value.download_file.assert_called_once()
+        self.assertEqual(
+            mock_boto.return_value.Bucket.return_value.download_file.call_args[0][0],
+            "spam",
+        )
 
     @patch("json.load")
     @patch("src.stac_utils.aws.open")
@@ -128,6 +155,30 @@ def test_save_to_s3_client_error(
         )
         self.assertRaises(ClientError, save_to_s3, mock_data, "foo", "bar", "spam")
 
+    def test_split_s3_url(self):
+        """Test split S3 url"""
+
+        test_url = "s3://foo-bucket/bar-path/spam-key.json"
+        self.assertTupleEqual(
+            split_s3_url(test_url), ("foo-bucket", "bar-path", "spam-key.json")
+        )
+
+    def test_split_s3_url_no_prefix(self):
+        """Test split S3 url with no prefix"""
+
+        test_url = "foo-bucket/bar-path/spam-key.json"
+        self.assertTupleEqual(
+            split_s3_url(test_url), ("foo-bucket", "bar-path", "spam-key.json")
+        )
+
+    def test_split_s3_url_no_path(self):
+        """Test split S3 url with no path"""
+
+        test_url = "s3://foo-bucket/spam-key.json"
+        self.assertTupleEqual(
+            split_s3_url(test_url), ("foo-bucket", "", "spam-key.json")
+        )
+
 
 if __name__ == "__main__":
     unittest.main()