Change minimize to a boolean flag with an option minimize-length argument

policy_sentry/bin/cli.py

@@ -2,7 +2,7 @@
 """
     Policy Sentry is a tool for generating least-privilege IAM Policies.
 """
-__version__ = "0.9.1"
+__version__ = "0.10.0"
 import click
 from policy_sentry import command
 

policy_sentry/command/write_policy.py

@@ -12,8 +12,42 @@
 
 logger = logging.getLogger(__name__)
 
+# adapted from
+# https://stackoverflow.com/questions/40753999/python-click-make-option-value-optional
+
+class RegisterLengthOption(click.Option):
+    """ Mark this option as getting a _length option """
+    register_length = True
+
+class RegisterLengthOptionHelp(click.Option):
+    """ Translate help for the hidden _length suffix """
+    def get_help_record(self, ctx):
+        help_text = super().get_help_record(ctx)
+        return (help_text[0].replace('_length ', ' '),) + help_text[1:]
+
+class RegisterMinimizeLengthCommand(click.Command):
+    """ Translate any opt= to opt_length= as needed """
+    def parse_args(self, ctx, args):
+        options = [o for o in ctx.command.params
+                   if getattr(o, 'register_length', None)]
+        prefixes = {p for p in sum([o.opts for o in options], [])
+                    if p.startswith('--')}
+        for i, a in enumerate(args):
+            a = a.split('=')
+            if a[0] in prefixes:
+                if len(a) > 1:
+                    args[i] = a[0]
+                    args.insert(i+1, a[0] + '_length=' + a[1])
+                else:
+                    # check if next argument is naked
+                    if len(args) > i+1 and not args[i+1].startswith('--'):
+                        value = args[i+1]
+                        args[i+1] = a[0] + '_length=' + value
+        return super().parse_args(ctx, args)
+
 
 @click.command(
+    cls=RegisterMinimizeLengthCommand,
     short_help="Write least-privilege IAM policies, restricting all actions to resource ARNs."
 )
 # pylint: disable=duplicate-code
@@ -24,10 +58,19 @@
 )
 @click.option(
     "--minimize",
+    cls=RegisterLengthOption,
+    is_flag=True,
     required=False,
-    type=int,
-    help="Minimize the resulting statement with *safe* usage of wildcards to reduce policy length. "
-    "Set this to the character length you want - for example, 4",
+    default=False,
+    help="Minimize the resulting statement with *safe* usage of wildcards to reduce policy length."
+)
+@click.option(
+    "--minimize_length",
+    cls=RegisterLengthOptionHelp,
+    help="Sets the minimum character length for minimization. Defaults to zero.",
+    type=click.IntRange(0),
+    required=False,
+    default=0
 )
 @click.option(
     "--fmt",
@@ -40,7 +83,7 @@
     '--verbose', '-v',
     type=click.Choice(['critical', 'error', 'warning', 'info', 'debug'],
     case_sensitive=False))
-def write_policy(input_file, minimize, fmt, verbose):
+def write_policy(input_file, minimize, minimize_length, fmt, verbose):
     """
     Write least-privilege IAM policies, restricting all actions to resource ARNs.
     """
@@ -56,7 +99,11 @@ def write_policy(input_file, minimize, fmt, verbose):
         except yaml.YAMLError as exc:
             logger.critical(exc)
             sys.exit()
-    policy = write_policy_with_template(cfg, minimize)
+
+    min_length = None
+    if minimize:
+        min_length = minimize_length
+    policy = write_policy_with_template(cfg, min_length)
 
     if fmt == "yaml":
         policy_str = yaml.dump(policy, sort_keys=False)
@@ -79,8 +126,6 @@ def write_policy_with_template(cfg, minimize=None):
     Returns:
         Dictionary: The JSON policy
     """
-    if minimize is not None and minimize < 0:
-        minimize = None
     sid_group = SidGroup()
     policy = sid_group.process_template(cfg, minimize)
     return policy