Merge pull request #4 from blacklocus/check-options-more-thoroughly

Add more input validation
blacklocus · Oct 17, 2016 · 0661c05 · 0661c05
2 parents e7a6b69 + 6d3fa57
commit 0661c05
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 18 deletions.
diff --git a/aurora_echo/echo_new.py b/aurora_echo/echo_new.py
@@ -29,7 +29,7 @@
 import click
 
 from aurora_echo.echo_const import ECHO_NEW_STAGE, ECHO_NEW_COMMAND
-from aurora_echo.echo_util import EchoUtil, log_prefix_factory
+from aurora_echo.echo_util import EchoUtil, log_prefix_factory, validate_input_param
 from aurora_echo.entry import root
 
 rds = boto3.client('rds')
@@ -132,12 +132,12 @@ def create_cluster_and_instance(cluster_params: dict, instance_params: dict, int
 
 
 @root.command()
-@click.option('--aws-account-number', '-a', required=True)
-@click.option('--region', '-r', required=True)
-@click.option('--cluster-snapshot-name', '-s', required=True)
-@click.option('--managed-name', '-n', required=True)
-@click.option('--db-subnet-group-name', '-sub', required=True)
-@click.option('--db-instance-class', '-c', required=True)
+@click.option('--aws-account-number', '-a', callback=validate_input_param, required=True)
+@click.option('--region', '-r', callback=validate_input_param, required=True)
+@click.option('--cluster-snapshot-name', '-s', callback=validate_input_param, required=True)
+@click.option('--managed-name', '-n', callback=validate_input_param, required=True)
+@click.option('--db-subnet-group-name', '-sub', callback=validate_input_param, required=True)
+@click.option('--db-instance-class', '-c', callback=validate_input_param, required=True)
 @click.option('--engine', '-e', default='aurora')
 @click.option('--availability-zone', '-az')
 @click.option('--vpc-security-group-id', '-sg', multiple=True)

diff --git a/aurora_echo/echo_promote.py b/aurora_echo/echo_promote.py
@@ -28,7 +28,7 @@
 import click
 
 from aurora_echo.echo_const import ECHO_NEW_STAGE, ECHO_PROMOTE_COMMAND, ECHO_PROMOTE_STAGE, ECHO_RETIRE_STAGE
-from aurora_echo.echo_util import EchoUtil, log_prefix_factory
+from aurora_echo.echo_util import EchoUtil, log_prefix_factory, validate_input_param
 from aurora_echo.entry import root
 
 rds = boto3.client('rds')
@@ -88,11 +88,11 @@ def update_dns(hosted_zone_id: str, record_set: dict, cluster_endpoint: str, ttl
 
 
 @root.command()
-@click.option('--aws-account-number', '-a', required=True)
-@click.option('--region', '-r', required=True)
-@click.option('--managed-name', '-n', required=True)
-@click.option('--hosted-zone-id', '-z', required=True)
-@click.option('--record-set', '-rs', required=True)
+@click.option('--aws-account-number', '-a', callback=validate_input_param, required=True)
+@click.option('--region', '-r', callback=validate_input_param, required=True)
+@click.option('--managed-name', '-n', callback=validate_input_param, required=True)
+@click.option('--hosted-zone-id', '-z', callback=validate_input_param, required=True)
+@click.option('--record-set', '-rs', callback=validate_input_param, required=True)
 @click.option('--ttl', default=60)
 @click.option('--interactive', '-i', default=True, type=bool)
 def promote(aws_account_number: str, region: str, managed_name: str, hosted_zone_id: str, record_set: str, ttl: str,

diff --git a/aurora_echo/echo_retire.py b/aurora_echo/echo_retire.py
@@ -28,7 +28,7 @@
 import click
 
 from aurora_echo.echo_const import ECHO_RETIRE_COMMAND, ECHO_RETIRE_STAGE
-from aurora_echo.echo_util import EchoUtil, log_prefix_factory
+from aurora_echo.echo_util import EchoUtil, log_prefix_factory, validate_input_param
 from aurora_echo.entry import root
 
 rds = boto3.client('rds')
@@ -63,9 +63,9 @@ def delete_instance(instance: dict, interactive: bool):
 
 
 @root.command()
-@click.option('--aws-account-number', '-a', required=True)
-@click.option('--region', '-r', required=True)
-@click.option('--managed-name', '-n', required=True)
+@click.option('--aws-account-number', '-a', callback=validate_input_param, required=True)
+@click.option('--region', '-r', callback=validate_input_param, required=True)
+@click.option('--managed-name', '-n', callback=validate_input_param, required=True)
 @click.option('--interactive', '-i', default=True, type=bool)
 def retire(aws_account_number: str, region: str, managed_name: str, interactive: bool):
     click.echo('{} Starting aurora-echo for {}'.format(log_prefix(), managed_name))

diff --git a/aurora_echo/echo_util.py b/aurora_echo/echo_util.py
@@ -26,6 +26,7 @@
 
 import boto3
 import click
+from botocore.exceptions import ClientError
 from dateutil.relativedelta import relativedelta
 
 from aurora_echo.echo_const import ECHO_MANAGEMENT_TAG_INDICATOR
@@ -39,6 +40,12 @@ def log_prefix():
     return log_prefix
 
 
+def validate_input_param(ctx, param, value):
+    if not value:
+        raise click.BadParameter('parameter must not be empty')
+    return value
+
+
 class EchoUtil(object):
     """
      General utilities, such as constructing tags, finding DB instances under a given managed name or stage,
@@ -91,7 +98,11 @@ def find_managed_instances(self, managed_name: str):
 
         # get all their tags
         for instance in response['DBInstances']:
-            tags = rds.list_tags_for_resource(ResourceName=self.construct_arn(instance['DBInstanceIdentifier']))
+            try:
+                arn = self.construct_arn(instance['DBInstanceIdentifier'])
+                tags = rds.list_tags_for_resource(ResourceName=arn)
+            except ClientError:
+                raise click.UsageError('Unable to list tags for resource at {!r}. Check your account number and region and try again.'.format(arn))
             tag_list = tags['TagList']
             for tag in tag_list:
                 # does it have our managed tag?