Add ParameterStore client with tests

Author: Luca Valentini

python_aws_ssm/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.1.0'

python_aws_ssm/parameters.py

@@ -0,0 +1,88 @@
+from typing import Dict, List, Optional, Union, Any
+
+import boto3
+
+
+class ParameterStore:
+    def __init__(self, client: Optional[boto3.client] = None):
+        self.client = client or boto3.client("ssm")
+
+    def get_parameters(self, ssm_key_names: List[str]) -> Dict[str, Optional[str]]:
+        """
+        Retrieve keys from SSM.
+        The keys are mapped to a dictionary for easy querying:
+            * Keys that exist in SSM should have a matching key in the result dict
+            and a matching value.
+            * Keys that do not exist in SSM should also have a matching key, but
+            have a matching value of None.
+        If SSM somehow returns keys that are not requested, these keys are not
+        returned in the result dict.
+        """
+
+        parameters = self.client.get_parameters(
+            Names=ssm_key_names, WithDecryption=True
+        ).get("Parameters")
+        return {
+            parameter.get("Name"): parameter.get("Value")
+            for parameter in parameters
+            if parameter.get("Name") in ssm_key_names
+        }
+
+    def get_parameters_by_path(
+        self,
+        ssm_base_path: str,
+        with_decryption: bool = True,
+        recursive: bool = False,
+        nested: bool = False,
+    ) -> Dict[str, Optional[str]]:
+        """
+        Retrieve all the keys under a certain path on SSM.
+        * Wnen recursive is set to False, SSM doesn't parameters under a nested path.
+            e.g.: /{ssm_base_path}/foo/bar will not return 'bar' nor '/foo/bar'.
+        * When recursive and nested are set to True, a nested dictionary is returned.
+            e.g.: /{ssm_base_path}/foo/bar will return {"foo": {"bar": "value"}}
+        * When nested is set to False, the full subpath is returned as key.
+            e.g.: /{ssm_base_path}/foo/bar will return {"foo/bar": "value"}}}
+        """
+
+        parameters = self.client.get_parameters_by_path(
+            Path=ssm_base_path, Recursive=recursive, WithDecryption=with_decryption
+        ).get("Parameters")
+
+        parameters = {
+            parameter.get("Name").replace(ssm_base_path, ""): parameter.get("Value")
+            for parameter in parameters
+        }
+
+        return (
+            self._parse_parameters(parameters) if recursive and nested else parameters
+        )
+
+    @staticmethod
+    def _parse_parameters(
+        parameters: Dict[str, Optional[str]]
+    ) -> Dict[Union[Dict, str], Optional[Union[Dict, str]]]:
+        parsed_dict: Dict[Union[Dict, str], Optional[Union[Dict, str]]] = {}
+        for key, value in parameters.items():
+            nested_dict = ParameterStore._tree_dict(key.split("/"), value)
+            parsed_dict = ParameterStore._deep_merge(parsed_dict, nested_dict)
+        return parsed_dict
+
+    @staticmethod
+    def _tree_dict(key_list: List[Any], value: Optional[Any]) -> Dict[Any, Any]:
+        tree_dict: Dict[Any, Any] = {key_list[-1]: value}
+        for key in reversed(key_list[:-1]):
+            tree_dict = {key: tree_dict}
+        return tree_dict
+
+    @staticmethod
+    def _deep_merge(a, b):
+        # NOTE: Thanks to: https://stackoverflow.com/a/56177639/9563578
+        if not isinstance(a, dict) or not isinstance(b, dict):
+            return a if b is None else b
+        else:
+            keys = set(a.keys()) | set(b.keys())
+            return {
+                key: ParameterStore._deep_merge(a.get(key), b.get(key)) for key in keys
+            }
+

tests/test_python_aws_ssm.py

@@ -0,0 +1,132 @@
+from unittest import TestCase, mock
+from unittest.mock import MagicMock
+
+from python_aws_ssm import __version__
+from python_aws_ssm.parameters import ParameterStore
+
+
+def test_version():
+    assert __version__ == "0.1.0"
+
+
+class TestGetParameters(TestCase):
+    def setUp(self):
+        self.parameter_store = ParameterStore(client=MagicMock())
+
+    def tearDown(self):
+        pass
+
+    def test_get_parameters_keys_are_mapped(self):
+        self.parameter_store.client.get_parameters.return_value = {
+            "Parameters": [
+                {"Name": "foo_ssm_key_1", "Value": "foo_ssm_value_1"},
+                # Note: foo_ssm_key_2 does not exist so is not returned from SSM.
+                {"Name": "foo_ssm_key_3", "Value": "foo_ssm_value_3"},
+            ]
+        }
+
+        secrets = self.parameter_store.get_parameters(
+            ["foo_ssm_key_1", "foo_ssm_key_2", "foo_ssm_key_3"]
+        )
+
+        self.assertEqual(
+            {"foo_ssm_key_1": "foo_ssm_value_1", "foo_ssm_key_3": "foo_ssm_value_3"},
+            secrets,
+        )
+
+        self.parameter_store.client.get_parameters.assert_called_once_with(
+            Names=["foo_ssm_key_1", "foo_ssm_key_2", "foo_ssm_key_3"],
+            WithDecryption=True,
+        )
+
+    def test_get_parameters_unknown_keys_are_ignored(self):
+        self.parameter_store.client.get_parameters.return_value = {
+            "Parameters": [
+                {"Name": "foo_ssm_key_1", "Value": "foo_ssm_value_1"},
+                {"Name": "some_other_key", "Value": "value"},
+            ]
+        }
+
+        secrets = self.parameter_store.get_parameters(["foo_ssm_key_1"])
+
+        self.assertEqual({"foo_ssm_key_1": "foo_ssm_value_1"}, secrets)
+
+    def test_get_parameters_aws_errors_are_not_caught(self):
+        expected_error = Exception("Unexpected AWS error!")
+        self.parameter_store.client.get_parameters.side_effect = expected_error
+
+        with self.assertRaises(Exception, msg="Unexpected AWS error!"):
+            self.parameter_store.get_parameters(["/key"])
+
+    def test_get_parameters_by_path_keys_are_mapped(self):
+        self.parameter_store.client.get_parameters_by_path.return_value = {
+            "Parameters": [
+                {"Name": "/bar/env/foo_ssm_key_1", "Value": "foo_ssm_value_1"},
+                {"Name": "/bar/env/foo_ssm_key_2", "Value": "foo_ssm_value_2"},
+            ]
+        }
+        secrets = self.parameter_store.get_parameters_by_path("/bar/env/")
+
+        self.assertEqual(
+            {"foo_ssm_key_1": "foo_ssm_value_1", "foo_ssm_key_2": "foo_ssm_value_2"},
+            secrets,
+        )
+
+        self.parameter_store.client.get_parameters_by_path.assert_called_once_with(
+            Path="/bar/env/", Recursive=False, WithDecryption=True
+        )
+
+    def test_get_parameters_by_path_recursive_not_nested(self):
+        self.parameter_store.client.get_parameters_by_path.return_value = {
+            "Parameters": [
+                {"Name": "/bar/env/foo_ssm_key_1", "Value": "foo_ssm_value_1"},
+                {"Name": "/bar/env/foo_ssm_key_2", "Value": "foo_ssm_value_2"},
+            ]
+        }
+        secrets = self.parameter_store.get_parameters_by_path(
+            "/bar/", recursive=True, nested=False
+        )
+
+        self.assertEqual(
+            {
+                "env/foo_ssm_key_1": "foo_ssm_value_1",
+                "env/foo_ssm_key_2": "foo_ssm_value_2",
+            },
+            secrets,
+        )
+
+        self.parameter_store.client.get_parameters_by_path.assert_called_once_with(
+            Path="/bar/", Recursive=True, WithDecryption=True
+        )
+
+    def test_get_parameters_by_path_recursive_nested(self):
+        self.parameter_store.client.get_parameters_by_path.return_value = {
+            "Parameters": [
+                {"Name": "/bar/env/foo_ssm_key_1", "Value": "foo_ssm_value_1"},
+                {"Name": "/bar/env/foo_ssm_key_2", "Value": "foo_ssm_value_2"},
+            ]
+        }
+        secrets = self.parameter_store.get_parameters_by_path(
+            "/bar/", recursive=True, nested=True
+        )
+
+        self.assertEqual(
+            {
+                "env": {
+                    "foo_ssm_key_1": "foo_ssm_value_1",
+                    "foo_ssm_key_2": "foo_ssm_value_2",
+                }
+            },
+            secrets,
+        )
+
+        self.parameter_store.client.get_parameters_by_path.assert_called_once_with(
+            Path="/bar/", Recursive=True, WithDecryption=True
+        )
+
+    def test_get_parameter_by_path_aws_errors_are_not_caught(self):
+        expected_error = Exception("Unexpected AWS error!")
+        self.parameter_store.client.get_parameters_by_path.side_effect = expected_error
+
+        with self.assertRaises(Exception, msg="Unexpected AWS error!"):
+            self.parameter_store.get_parameters_by_path(["/key"])