Significantly improves storage engines including (1) converting comma…

…nd line flags to ENV variables (fixes #15), (2) a way to generate YAML files for branches of the SSM tree (closes #11), and (3) the ability to ignore SecureString keys if they are not necessary (closes #13), and (4) the introduction of metadata in the YAML files to permit compatibility checking (more general fix for #15 with support for new features)

ssm-diff

@@ -10,7 +10,8 @@ from states import *
 def configure_endpoints(args):
     # configure() returns a DiffBase class (whose constructor may be wrapped in `partial` to pre-configure it)
     diff_class = DiffBase.get_plugin(args.engine).configure(args)
-    return storage.ParameterStore(args.profile, diff_class, paths=args.path), storage.YAMLFile(args.filename, paths=args.path)
+    return storage.ParameterStore(args.profile, diff_class, paths=args.paths, no_secure=args.no_secure), \
+           storage.YAMLFile(args.filename, paths=args.paths, no_secure=args.no_secure, root_path=args.yaml_root)
 
 
 def init(args):
@@ -49,8 +50,7 @@ def plan(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('-f', help='local state yml file', action='store', dest='filename', default='parameters.yml')
-    parser.add_argument('--path', '-p', action='append', help='filter SSM path')
+    parser.add_argument('-f', help='local state yml file', action='store', dest='filename')
     parser.add_argument('--engine', '-e', help='diff engine to use when interacting with SSM', action='store', dest='engine', default='DiffResolver')
     parser.add_argument('--profile', help='AWS profile name', action='store', dest='profile')
     subparsers = parser.add_subparsers(dest='func', help='commands')
@@ -70,12 +70,29 @@ if __name__ == "__main__":
     parser_apply.set_defaults(func=apply)
 
     args = parser.parse_args()
-    args.path = args.path if args.path else ['/']
-
-    if args.filename == 'parameters.yml':
-        if not args.profile:
-            if 'AWS_PROFILE' in os.environ:
-                args.filename = os.environ['AWS_PROFILE'] + '.yml'
-        else:
-            args.filename = args.profile + '.yml'
+
+    args.no_secure = os.environ.get('SSM_NO_SECURE', 'false').lower() in ['true', '1']
+    args.yaml_root = os.environ.get('SSM_YAML_ROOT', '/')
+    args.paths = os.environ.get('SSM_PATHS', None)
+    if args.paths is not None:
+        args.paths = args.paths.split(';:')
+    else:
+        # this defaults to '/'
+        args.paths = args.yaml_root
+
+    # root filename
+    if args.filename is not None:
+        filename = args.filename
+    elif args.profile:
+        filename = args.profile
+    elif 'AWS_PROFILE' in os.environ:
+        filename = os.environ['AWS_PROFILE']
+    else:
+        filename = 'parameters'
+
+    # remove extension (will be restored by storage classes)
+    if filename[-4:] == '.yml':
+        filename = filename[:-4]
+    args.filename = filename
+
     args.func(args)

states/engine.py

@@ -66,15 +66,15 @@ def describe_diff(cls, plan):
         description = ""
         for k, v in plan['add'].items():
             # { key: new_value }
-            description += colored("+", 'green'), "{} = {}".format(k, v) + '\n'
+            description += colored("+", 'green') + "{} = {}".format(k, v) + '\n'
 
         for k in plan['delete']:
             # { key: old_value }
-            description += colored("-", 'red'), k + '\n'
+            description += colored("-", 'red') + k + '\n'
 
         for k, v in plan['change'].items():
             # { key: {'old': value, 'new': value} }
-            description += colored("~", 'yellow'), "{}:\n\t< {}\n\t> {}".format(k, v['old'], v['new']) + '\n'
+            description += colored("~", 'yellow') + "{}:\n\t< {}\n\t> {}".format(k, v['old'], v['new']) + '\n'
 
         return description
 

states/helpers.py

@@ -5,24 +5,31 @@ def add(obj, path, value, sep='/'):
     """Add value to the `obj` dict at the specified path"""
     parts = path.strip(sep).split(sep)
     last = len(parts) - 1
+    current = obj
     for index, part in enumerate(parts):
         if index == last:
-            obj[part] = value
+            current[part] = value
         else:
-            obj = obj.setdefault(part, {})
+            current = current.setdefault(part, {})
+    # convenience return, object is mutated
+    return obj
 
 
 def search(state, path):
-    result = state
+    """Get value in `state` at the specified path, returning {} if the key is absent"""
+    if path.strip("/") == '':
+        return state
     for p in path.strip("/").split("/"):
-        if result.clone(p):
-            result = result[p]
-        else:
-            result = {}
-            break
-    output = {}
-    add(output, path, result)
-    return output
+        if p not in state:
+            return {}
+        state = state[p]
+    return state
+
+
+def filter(state, path):
+    if path.strip("/") == '':
+        return state
+    return add({}, path, search(state, path))
 
 
 def merge(a, b):

states/storage.py

@@ -7,7 +7,7 @@
 import yaml
 from botocore.exceptions import ClientError, NoCredentialsError
 
-from .helpers import merge, add, search
+from .helpers import merge, add, filter, search
 
 
 def str_presenter(dumper, data):
@@ -62,18 +62,34 @@ def to_yaml(cls, dumper, data):
 
 class YAMLFile(object):
     """Encodes/decodes a dictionary to/from a YAML file"""
-    def __init__(self, filename, paths=('/',)):
-        self.filename = filename
+    METADATA_CONFIG = 'ssm-diff:config'
+    METADATA_PATHS = 'ssm-diff:paths'
+    METADATA_ROOT = 'ssm:root'
+    METADATA_NO_SECURE = 'ssm:no-secure'
+
+    def __init__(self, filename, paths=('/',), root_path='/', no_secure=False):
+        self.filename = '{}.yml'.format(filename)
+        self.root_path = root_path
         self.paths = paths
+        self.validate_paths()
+        self.no_secure = no_secure
+
+    def validate_paths(self):
+        length = len(self.root_path)
+        for path in self.paths:
+            if path[:length] != self.root_path:
+                raise ValueError('Root path {} does not contain path {}'.format(self.root_path, path))
 
     def get(self):
         try:
             output = {}
             with open(self.filename, 'rb') as f:
                 local = yaml.safe_load(f.read())
+            self.validate_config(local)
+            local = self.nest_root(local)
             for path in self.paths:
                 if path.strip('/'):
-                    output = merge(output, search(local, path))
+                    output = merge(output, filter(local, path))
                 else:
                     return local
             return output
@@ -87,7 +103,55 @@ def get(self):
                 return dict()
             raise
 
+    def validate_config(self, local):
+        """YAML files may contain a special ssm:config tag that stores information about the file when it was generated.
+        This information can be used to ensure the file is compatible with future calls.  For example, a file created
+        with a particular subpath (e.g. /my/deep/path) should not be used to overwrite the root path since this would
+        delete any keys not in the original scope.  This method does that validation (with permissive defaults for
+        backwards compatibility)."""
+        config = local.pop(self.METADATA_CONFIG, {})
+
+        # strict requirement that the no_secure setting is equal
+        config_no_secure = config.get(self.METADATA_NO_SECURE, False)
+        if config_no_secure != self.no_secure:
+            raise ValueError("YAML file generated with no_secure={} but current class set to no_secure={}".format(
+                config_no_secure, self.no_secure,
+            ))
+        # strict requirement that root_path is equal
+        config_root = config.get(self.METADATA_ROOT, '/')
+        if config_root != self.root_path:
+            raise ValueError("YAML file generated with root_path={} but current class set to root_path={}".format(
+                config_root, self.root_path,
+            ))
+        # make sure all paths are subsets of file paths
+        config_paths = config.get(self.METADATA_PATHS, ['/'])
+        for path in self.paths:
+            for config_path in config_paths:
+                # if path is not found in a config path, it could look like we've deleted values
+                if path[:len(config_path)] == config_path:
+                    break
+            else:
+                raise ValueError("Path {} was not included in this file when it was created.".format(path))
+
+    def unnest_root(self, state):
+        if self.root_path == '/':
+            return state
+        return search(state, self.root_path)
+
+    def nest_root(self, state):
+        if self.root_path == '/':
+            return state
+        return add({}, self.root_path, state)
+
     def save(self, state):
+        state = self.unnest_root(state)
+        # inject state information so we can validate the file on load
+        # colon is not allowed in SSM keys so this namespace cannot collide with keys at any depth
+        state[self.METADATA_CONFIG] = {
+            self.METADATA_PATHS: self.paths,
+            self.METADATA_ROOT: self.root_path,
+            self.METADATA_NO_SECURE: self.no_secure
+        }
         try:
             with open(self.filename, 'wb') as f:
                 content = yaml.safe_dump(state, default_flow_style=False)
@@ -99,12 +163,21 @@ def save(self, state):
 
 class ParameterStore(object):
     """Encodes/decodes a dict to/from the SSM Parameter Store"""
-    def __init__(self, profile, diff_class, paths=('/',)):
+    def __init__(self, profile, diff_class, paths=('/',), no_secure=False):
         if profile:
             boto3.setup_default_session(profile_name=profile)
         self.ssm = boto3.client('ssm')
         self.diff_class = diff_class
         self.paths = paths
+        self.parameter_filters = []
+        if no_secure:
+            self.parameter_filters.append({
+                'Key': 'Type',
+                'Option': 'Equals',
+                'Values': [
+                    'String', 'StringList',
+                ]
+            })
 
     def clone(self):
         p = self.ssm.get_paginator('get_parameters_by_path')
@@ -114,7 +187,9 @@ def clone(self):
                 for page in p.paginate(
                         Path=path,
                         Recursive=True,
-                        WithDecryption=True):
+                        WithDecryption=True,
+                        ParameterFilters=self.parameter_filters,
+                ):
                     for param in page['Parameters']:
                         add(obj=output,
                             path=param['Name'],
@@ -136,10 +211,10 @@ def pull(self, local):
         return diff.merge()
 
     def dry_run(self, local):
-        return self.diff_class(self.clone(), local).plan
+        return self.diff_class(self.clone(), local)
 
     def push(self, local):
-        plan = self.dry_run(local)
+        plan = self.dry_run(local).plan
 
         # plan
         for k, v in plan['add'].items():