[cli] updating cli to support toggling metrics per cluster/function

* Adding some flags to the `manage.py metrics` command to accept a cluster and a function name
* By default, function name is required while cluster is not (will default to all clusters)

conf/clusters/prod.json

@@ -25,15 +25,14 @@
         "current_version": "$LATEST",
         "log_level": "info",
         "memory": 128,
-        "timeout": 10,
-        "enable_metrics": true
+        "timeout": 10
       },
       "rule_processor": {
         "current_version": "$LATEST",
+        "enable_metrics": true,
         "log_level": "info",
         "memory": 128,
-        "timeout": 30,
-        "enable_metrics": true
+        "timeout": 30
       }
     }
   },

manage.py

@@ -23,7 +23,7 @@
 
 terraform <cmd> -var-file=../terraform.tfvars -var-file=../variables.json
 """
-from argparse import ArgumentParser, RawTextHelpFormatter, SUPPRESS as ARGPARSE_SUPPRESS
+from argparse import Action, ArgumentParser, RawTextHelpFormatter, SUPPRESS as ARGPARSE_SUPPRESS
 import os
 import re
 
@@ -33,6 +33,13 @@
 from stream_alert_cli.runner import cli_runner
 
 
+class UniqueSetAction(Action):
+    """Subclass of argparse.Action to avoid multiple of the same choice from a list"""
+    def __call__(self, parser, namespace, values, option_string=None):
+        unique_items = set(values)
+        setattr(namespace, self.dest, unique_items)
+
+
 def _add_output_subparser(subparsers):
     """Add the output subparser: manage.py output [subcommand] [options]"""
     output_usage = 'manage.py output [subcommand] [options]'
@@ -196,6 +203,12 @@ def _add_metrics_subparser(subparsers):
     """Add the metrics subparser: manage.py metrics [options]"""
     metrics_usage = 'manage.py metrics [options]'
 
+    # get cluster choices from available files
+    clusters = [os.path.splitext(cluster)[0] for _, _, files
+                in os.walk('conf/clusters') for cluster in files]
+
+    cluster_choices_block = ('\n').join('{:>28}{}'.format('', cluster) for cluster in clusters)
+
     metrics_description = ("""
 StreamAlertCLI v{}
 Enable or disable metrics for all lambda functions. This toggles the creation of metric filters.
@@ -204,13 +217,23 @@ def _add_metrics_subparser(subparsers):
 
     -e/--enable         Enable CloudWatch metrics through logging and metric filters
     -d/--disable        Disable CloudWatch metrics through logging and metric filters
+    -f/--functions      Space delimited list of functions to enable metrics for
+                          Choices are:
+                            rule
+                            alert (not implemented)
+                            athena (not implemented)
     --debug             Enable Debug logger output
 
+Optional Arguemnts:
+
+    -c/--clusters       Space delimited list of clusters to enable metrics for. If
+                          omitted, this will enable metrics for all clusters. Choices are:
+{}
 Examples:
 
     manage.py metrics --enable
 
-""".format(version))
+""".format(version, cluster_choices_block))
 
     metrics_parser = subparsers.add_parser(
         'metrics',
@@ -223,6 +246,16 @@ def _add_metrics_subparser(subparsers):
     # Set the name of this parser to 'metrics'
     metrics_parser.set_defaults(command='metrics')
 
+    # allow the user to select 1 or more functions to enable metrics for
+    metrics_parser.add_argument(
+        '-f', '--functions',
+        choices=['rule', 'alert', 'athena'],
+        help=ARGPARSE_SUPPRESS,
+        nargs='+',
+        action=UniqueSetAction,
+        required=True
+    )
+
     # get the metric toggle value
     toggle_group = metrics_parser.add_mutually_exclusive_group(required=True)
 
@@ -238,6 +271,16 @@ def _add_metrics_subparser(subparsers):
         action='store_false'
     )
 
+    # allow the user to select 0 or more clusters to enable metrics for
+    metrics_parser.add_argument(
+        '-c', '--clusters',
+        choices=clusters,
+        help=ARGPARSE_SUPPRESS,
+        nargs='+',
+        action=UniqueSetAction,
+        default=clusters
+    )
+
     # allow verbose output for the CLI with the --debug option
     metrics_parser.add_argument(
         '--debug',
@@ -256,6 +299,12 @@ def _add_metric_alarm_subparser(subparsers):
 
     metric_choices_block = ('\n').join('{:>35}{}'.format('', metric) for metric in all_metrics)
 
+    # get cluster choices from available files
+    clusters = [os.path.splitext(cluster)[0] for _, _, files
+                in os.walk('conf/clusters') for cluster in files]
+
+    cluster_choices_block = ('\n').join('{:>37}{}'.format('', cluster) for cluster in clusters)
+
     metric_alarm_description = ("""
 StreamAlertCLI v{}
 Add a CloudWatch alarm for predefined metrics. These are save in the config and
@@ -291,6 +340,8 @@ def _add_metric_alarm_subparser(subparsers):
     -ad/--alarm-description      The description for the alarm
     -c/--clusters                Space delimited list of clusters to apply this metric to. This is
                                    ignored if the --metric-target of 'aggregate' is used.
+                                   Choices are:
+{}
     -es/--extended-statistic     The percentile statistic for the metric associated with the alarm.
                                    Specify a value between p0.0 and p100.  Cannot be used in
                                    conjunction with the --statistic flag.
@@ -330,7 +381,7 @@ def _add_metric_alarm_subparser(subparsers):
     AWS:        https://docs.aws.amazon.com/AmazonCloudWatch/latest/APIReference/API_PutMetricAlarm.html
     Terraform:  https://www.terraform.io/docs/providers/aws/r/cloudwatch_metric_alarm.html
 
-""".format(version, metric_choices_block))
+""".format(version, metric_choices_block, cluster_choices_block))
 
     metric_alarm_parser = subparsers.add_parser(
         'create-alarm',
@@ -442,16 +493,13 @@ def _alarm_description_validator(val):
         type=_alarm_description_validator
     )
 
-    # get cluster choices from available files
-    clusters = [os.path.splitext(cluster)[0] for _, _, files
-                in os.walk('conf/clusters') for cluster in files]
-
     # allow the user to select 0 or more clusters to apply this alarm to
     metric_alarm_parser.add_argument(
         '-c', '--clusters',
         choices=clusters,
         help=ARGPARSE_SUPPRESS,
         nargs='+',
+        action=UniqueSetAction,
         default=[]
     )
 

stream_alert_cli/config.py

@@ -19,6 +19,7 @@
 import re
 import sys
 
+from stream_alert.shared import metrics
 from stream_alert_cli.helpers import continue_prompt
 from stream_alert_cli.logger import LOGGER_CLI
 
@@ -128,13 +129,27 @@ def set_aws_account_id(self, aws_account_id):
 
         LOGGER_CLI.info('AWS Account ID successfully configured')
 
-    def toggle_metrics(self, enabled):
-        """Toggle CloudWatch metric logging and filter creation"""
-        metrics = self.config['global']['infrastructure'].get('metrics')
-        if not metrics:
-            self.config['global']['infrastructure']['metrics'] = {}
+    def toggle_metrics(self, enabled, clusters, lambda_functions):
+        """Toggle CloudWatch metric logging and filter creation
 
-        metrics['enabled'] = enabled
+        Args:
+            enabled (bool): False if disabling metrics, true if enable_logging
+            clusters (list): Clusters to enable or disable metrics on
+            lambda_functions (list): Which lambda functions to enable or disable
+                metrics on (rule, alert, or athena)
+        """
+        for function in lambda_functions:
+            if function == 'athena':
+                if 'athena_partition_refresh_config' in self.config['lambda']:
+                    self.config['lambda']['athena_partition_refresh_config'] \
+                        ['enable_metrics'] = enabled
+                else:
+                    LOGGER_CLI.error('No Athena configuration found; please initialize first.')
+                continue
+
+            for cluster in clusters:
+                self.config['clusters'][cluster]['modules']['stream_alert'] \
+                    ['{}_processor'.format(function)]['enable_metrics'] = enabled
 
         self.write()
 

stream_alert_cli/runner.py

@@ -655,7 +655,7 @@ def _toggle_metrics(options):
     Args:
         options (argparser): Contains boolean necessary for toggling metrics
     """
-    CONFIG.toggle_metrics(options.enable_metrics)
+    CONFIG.toggle_metrics(options.enable_metrics, options.clusters, options.functions)
 
 
 def _create_alarm(options):