Merge cc6ceb9 into 3256e81

docs/source/apps.rst

@@ -150,7 +150,7 @@ The recommended process is to deploy both the `apps` function and the `classifie
 
 .. code-block:: bash
 
-  python manage.py deploy --function classifier apps
+  python manage.py deploy --functions classifier apps
 
 
 Authorizing the Slack App

docs/source/deployment.rst

@@ -35,20 +35,20 @@ To deploy new changes for all AWS Lambda functions:
 
 .. code-block:: bash
 
-  python manage.py deploy --function all
+  python manage.py deploy
 
 Optionally, to deploy changes for only a specific AWS Lambda function:
 
 .. code-block:: bash
 
-  python manage.py deploy --function alert
-  python manage.py deploy --function alert_merger
-  python manage.py deploy --function apps
-  python manage.py deploy --function athena
-  python manage.py deploy --function classifier
-  python manage.py deploy --function rule
-  python manage.py deploy --function rule_promo
-  python manage.py deploy --function threat_intel_downloader
+  python manage.py deploy --functions alert
+  python manage.py deploy --functions alert_merger
+  python manage.py deploy --functions apps
+  python manage.py deploy --functions athena
+  python manage.py deploy --functions classifier
+  python manage.py deploy --functions rule
+  python manage.py deploy --functions rule_promo
+  python manage.py deploy --functions threat_intel_downloader
 
 To apply infrastructure level changes (additional Kinesis Shards, new CloudTrails, etc), run:
 
@@ -95,8 +95,8 @@ to point to the previous version:
 
 .. code-block:: bash
 
-  python manage.py rollback --function rule
-  python manage.py rollback --function alert
-  python manage.py rollback --function all
+  python manage.py rollback --functions rule
+  python manage.py rollback --functions alert
+  python manage.py rollback
 
 This is helpful to quickly revert changes to Lambda functions, e.g. if a bad rule was deployed.

docs/source/getting-started.rst

@@ -237,7 +237,7 @@ alerts on any usage of the root AWS account. Change the rule decorator to:
   python manage.py build
 
   # Deploy a new version of all of the Lambda functions with the updated rule and config files
-  python manage.py deploy --function all
+  python manage.py deploy
 
 .. note:: Use ``build`` and ``deploy`` to apply any changes to StreamAlert's
    configuration or Lambda functions, respectively. Some changes (like this example) require both.

docs/source/historical-search.rst

@@ -73,7 +73,7 @@ Alerts Search
 
   .. code-block:: bash
 
-    python manage.py deploy --function athena
+    python manage.py deploy --functions athena
 
 * Search alerts in `Athena Console <https://console.aws.amazon.com/athena>`_
 
@@ -99,7 +99,7 @@ It is optional to store data in S3 bucket and available for search in Athena tab
 
   .. code-block:: bash
 
-    python manage.py deploy --function classifier
+    python manage.py deploy --functions classifier
 
 * Search data `Athena Console <https://console.aws.amazon.com/athena>`_
 

docs/source/rule-promotion.rst

@@ -76,7 +76,7 @@ function code.
 
 .. code-block:: bash
 
-  python manage.py deploy --function rule_promo
+  python manage.py deploy --functions rule_promo
 
 .. note::
 

docs/source/rule-staging.rst

@@ -83,7 +83,7 @@ staged during a deploy. To allow for this, the Rules Engine can be deployed with
 
 .. code-block:: bash
 
-  python manage.py deploy --function rule --skip-rule-staging
+  python manage.py deploy --functions rule --skip-rule-staging
 
 This will force all new rules to send to user-defined outputs immediately upon deploy, bypassing
 the default staging period. Alternatively, the ``--stage-rules`` and ``--unstage-rules`` flags

streamalert_cli/manage_lambda/deploy.py

@@ -93,26 +93,11 @@ def handler(cls, options, config):
         if not terraform_generate_handler(config=config):
             return False
 
-        functions = options.function
-
-        if 'all' in options.function:
-            functions = {
-                'alert',
-                'alert_merger',
-                'apps',
-                'athena',
-                'classifier',
-                'rule',
-                'rule_promo',
-                'scheduled_queries',
-                'threat_intel_downloader'
-            }
-
-        if not deploy(functions, config, options.clusters):
+        if not deploy(options.functions, config, options.clusters):
             return False
 
         # Update the rule table now if the rules engine is being deployed
-        if 'rule' in functions:
+        if 'rule' in set(options.functions):
             _update_rule_table(options, config)
 
         return True
@@ -129,8 +114,7 @@ def deploy(functions, config, clusters=None):
     Returns:
         bool: False if errors occurred, True otherwise
     """
-
-    LOGGER.info('Deploying: %s', ' '.join(sorted(functions)))
+    LOGGER.info('Deploying: %s', ', '.join(sorted(functions)))
 
     # Terraform apply only to the module which contains our lambda functions
     deploy_targets = set()

streamalert_cli/manage_lambda/rollback.py

@@ -21,6 +21,7 @@
 from streamalert_cli.utils import (
     add_default_lambda_args,
     CLICommand,
+    function_map,
     set_parser_epilog,
 )
 
@@ -97,55 +98,37 @@ def handler(cls, options, config):
         if not terraform_generate_handler(config=config):
             return False
 
-        LOGGER.info('Rolling back: %s', ' '.join(options.function))
+        functions = function_map()
+        targeted_funcs = set(options.functions)
+        functions = {
+            key: value for key, value in functions.items() if key in targeted_funcs
+        }
+
+        LOGGER.info('Rolling back: %s', ', '.join(sorted(functions)))
 
-        rollback_all = 'all' in options.function
         prefix = config['global']['account']['prefix']
         clusters = sorted(options.clusters or config.clusters())
         client = boto3.client('lambda')
 
         # Track the success of rolling back the functions
         success = True
-        if rollback_all or 'alert' in options.function:
-            success = success and _rollback_production(
-                client,
-                '{}_streamalert_alert_processor'.format(prefix)
-            )
-
-        if rollback_all or 'alert_merger' in options.function:
-            success = success and _rollback_production(
-                client,
-                '{}_streamalert_alert_merger'.format(prefix)
-            )
-
-        if rollback_all or 'apps' in options.function:
-            for cluster in clusters:
-                apps_config = config['clusters'][cluster]['modules'].get('streamalert_apps', {})
-                for lambda_name in sorted(apps_config):
-                    success = success and _rollback_production(client, lambda_name)
-
-        if rollback_all or 'athena' in options.function:
-            success = success and _rollback_production(
-                client,
-                '{}_streamalert_athena_partitioner'.format(prefix)
-            )
-
-        if rollback_all or 'classifier' in options.function:
-            for cluster in clusters:
+        for func, suffix in functions.items():
+            if suffix:  # A suffix implies this is a standard function naming convention
                 success = success and _rollback_production(
                     client,
-                    '{}_{}_streamalert_classifier'.format(prefix, cluster)
+                    '{}_streamalert_{}'.format(prefix, suffix)
                 )
-
-        if rollback_all or 'rule' in options.function:
-            success = success and _rollback_production(
-                client, '{}_streamalert_rules_engine'.format(prefix)
-            )
-
-        if rollback_all or 'threat_intel_downloader' in options.function:
-            success = success and _rollback_production(
-                client,
-                '{}_streamalert_threat_intel_downloader'.format(prefix)
-            )
+            elif func == 'apps':  # Apps need special handling due to unique naming
+                for cluster in clusters:
+                    cluster_modules = config['clusters'][cluster]['modules']
+                    apps_config = cluster_modules.get('streamalert_apps', {})
+                    for lambda_name in sorted(apps_config):
+                        success = success and _rollback_production(client, lambda_name)
+            elif func == 'classifier':  # Classifers need special handling due to clustering
+                for cluster in clusters:
+                    success = success and _rollback_production(
+                        client,
+                        '{}_{}_streamalert_{}'.format(prefix, cluster, func)
+                    )
 
         return success

streamalert_cli/utils.py

@@ -35,6 +35,26 @@
 ]
 
 
+def function_map():
+    """Provide a map of CLI function name to their expected actual function suffix
+
+    Returns:
+        dict: Mapping of CLI function name to expected actual function suffix
+    """
+    # This is purposely not a constant because it is expected to be modifiable
+    return {
+        'alert': 'alert_processor',
+        'alert_merger': 'alert_merger',
+        'apps': None,  # needs special handling
+        'athena': 'athena_partitioner',
+        'classifier': None,  # needs special handling
+        'rule': 'rules_engine',
+        'rule_promo': 'rule_promotion',
+        'scheduled_queries': 'scheduled_queries_runner',
+        'threat_intel_downloader': 'threat_intel_downloader',
+    }
+
+
 class CLICommand:
     """
     An abstract class that encapsulates the logic of a single manage.py CLI command.
@@ -239,23 +259,19 @@ def generate_subparser(parser, name, description=None, subcommand=False, **kwarg
 
 def add_default_lambda_args(lambda_parser):
     """Add the default arguments to the deploy and rollback parsers"""
-
-    functions = sorted([
-        'alert', 'alert_merger', 'apps', 'athena', 'classifier',
-        'rule', 'rule_promo', 'scheduled_queries', 'threat_intel_downloader'
-    ])
-    # require the name of the function being deployed/rolled back
+    functions = sorted(function_map())
+    # optionally allow for the name of 1+ functions being deployed/rolled back
     lambda_parser.add_argument(
-        '-f', '--function',
-        choices=functions + ['all'],
-        metavar='FUNCTION',
+        '-f', '--functions',
+        choices=functions,
+        default=functions,
+        metavar='FUNCTIONS',
         help=(
             'One or more of the following functions to perform this action against: {}. '
-            'Use \'all\' to act against all functions.'
+            'If omitted, this action will be performed against all functions.'
         ).format(', '.join(functions)),
         nargs='+',
         action=UniqueSortedListAction,
-        required=True
     )
 
     # Add the option to specify cluster(s)

tests/unit/streamalert_cli/manage_lambda/test_rollback.py

@@ -15,7 +15,7 @@ class MockOptions:
 
     def __init__(self, clusters, function):
         self.clusters = clusters
-        self.function = function
+        self.functions = function
 
 
 @mock.patch.object(rollback, 'boto3', mock.MagicMock())
@@ -64,9 +64,13 @@ def test_rollback_production_error(self, mock_logger):
     def test_rollback_all(self, mock_helper):
         """CLI - Lambda rollback all"""
         mock_helper.return_value = True
+        funcs = [
+            'alert', 'alert_merger', 'apps', 'athena', 'classifier',
+            'rule', 'rule_promo', 'scheduled_queries', 'threat_intel_downloader'
+        ]
         assert_equal(
             rollback.RollbackCommand.handler(
-                MockOptions(None, ['all']),
+                MockOptions(None, funcs),
                 MockCLIConfig(config=basic_streamalert_config())
             ),
             True
@@ -80,6 +84,8 @@ def test_rollback_all(self, mock_helper):
             mock.call(mock.ANY, 'unit-test_corp_streamalert_classifier'),
             mock.call(mock.ANY, 'unit-test_prod_streamalert_classifier'),
             mock.call(mock.ANY, 'unit-test_streamalert_rules_engine'),
+            mock.call(mock.ANY, 'unit-test_streamalert_rule_promotion'),
+            mock.call(mock.ANY, 'unit-test_streamalert_scheduled_queries_runner'),
             mock.call(mock.ANY, 'unit-test_streamalert_threat_intel_downloader')
         ])