Manage loading in values from parameter store into the environment

pse.py

import logging
import os
from collections import defaultdict
from typing import Dict


LOG = logging.getLogger(__name__)


class ParameterStoreEnvirontment:
    detection_prefix = "ssm-pse:"

    def __init__(self, ssm_client):
        self.client = ssm_client

    def detect_and_inject(
        self,
        with_decryption: bool = True,
        default: str = None,
        injection_target: Dict[str, str] = os.environ,
    ) -> None:
        """Scans a dictionary for values that look like SSM PSE values, and loads
        them in from Parameter Store.
        """

        env_to_path_map = self.detect_env_for_loading(injection_target)
        new_environment = self.fetch(
            env_to_path_map, with_decryption=with_decryption, default=default
        )

        for key, value in new_environment.items():
            injection_target[key] = value

        return

    def fetch(
        self, input: Dict[str, str], with_decryption: bool = True, default: str = None,
    ) -> Dict[str, str]:

        # We need a list, as multiple keys can map to the same values
        inverted_input = defaultdict(list)
        for var, path in input.items():
            inverted_input[path].append(var)

        response = self.client.get_parameters(
            Names=list(inverted_input.keys()), WithDecryption=with_decryption,
        )

        self._validate_response(response, allow_missing=default is not None)

        return {
            var: param["Value"]
            for param in response["Parameters"]
            for var in inverted_input[param["Name"]]
        }

    def detect_env_for_loading(
        self, detection_target: Dict[str, str] = os.environ,
    ) -> Dict[str, str]:
        """
        Search an object for anything contains our expected prefix, and looks like
        a Parameter Store path.
        """
        output = {}
        for key, value in detection_target.items():
            if value.startswith(self.detection_prefix):
                ps_path = value[len(self.detection_prefix) :]
                # just some extra sanity to make sure it's kinda like a path (PS needs this)

                if ps_path.startswith("/"):
                    output[key] = ps_path
                else:
                    LOG.warning(
                        "Environment value detected with PSE Prefix: {}={}".format(
                            key, value
                        )
                    )

        return output

    @staticmethod
    def _validate_response(response, allow_missing: bool) -> None:
        if response["InvalidParameters"]:
            LOG.warning(
                "Invalid Parameters detected: %(parameters)r",
                {"parameters": response["InvalidParameters"]},
            )
            if default is None:
                raise ValueError("Invalid Parameters detected without a `default` set.")

        return