Merge f3f1426 into a83f3e9

conf/outputs.json

@@ -2,6 +2,9 @@
   "aws-lambda": {
     "sample-lambda": "function-name:qualifier"
   },
+  "aws-lambda-v2": [
+    "sample-lambda"
+  ],
   "aws-s3": {
     "bucket": "aws-s3-bucket"
   },

streamalert/alert_processor/outputs/aws.py

@@ -158,7 +158,10 @@ def _firehose_request_wrapper(json_alert, delivery_stream):
 
 @StreamAlertOutput
 class LambdaOutput(AWSOutput):
-    """LambdaOutput handles all alert dispatching to AWS Lambda"""
+    """LambdaOutput handles all alert dispatching to AWS Lambda
+
+    This output is deprecated by the aws-lambda-v2 output
+    """
     __service__ = 'aws-lambda'
 
     @classmethod
@@ -256,6 +259,157 @@ def _dispatch(self, alert, descriptor):
         return True
 
 
+@StreamAlertOutput
+class LambdaOutputV2(AWSOutput):
+    """LambdaOutput handles all alert dispatching to AWS Lambda"""
+    __service__ = 'aws-lambda-v2'
+
+    @classmethod
+    def get_user_defined_properties(cls):
+        """Get properties that must be assigned by the user when configuring a new Lambda
+        output.  This should be sensitive or unique information for this use-case that needs
+        to come from the user.
+
+        Every output should return a dict that contains a 'descriptor' with a description of the
+        integration being configured.
+
+        Sending to Lambda also requires a user provided Lambda function name and optional qualifier
+        (if applicable for the user's use case). A fully-qualified AWS ARN is also acceptable for
+        this value. This value should not be masked during input and is not a credential requirement
+        that needs encrypted.
+
+        When invoking a Lambda function in a different AWS account, the Alert Processor will have
+        to first assume a role in the target account. Both the Alert Processor and the destination
+        role will need AssumeRole IAM policies to allow this:
+
+        @see https://aws.amazon.com/premiumsupport/knowledge-center/lambda-function-assume-iam-role/
+
+        Returns:
+            OrderedDict: Contains various OutputProperty items
+        """
+        return OrderedDict(
+            [
+                (
+                    'descriptor',
+                    OutputProperty(
+                        description='a short and unique descriptor for this Lambda function '
+                        'configuration (ie: abbreviated name)'
+                    )
+                ),
+                (
+                    'lambda_function_arn',
+                    OutputProperty(
+                        description='The ARN of the AWS Lambda function to Invoke',
+                        input_restrictions={' '}
+                    )
+                ),
+                (
+                    'function_qualifier',
+                    OutputProperty(
+                        description='The function qualifier/alias to invoke.',
+                        input_restrictions={' '}
+                    )
+                ),
+                (
+                    'assume_role_arn',
+                    OutputProperty(
+                        description='When provided, will use AssumeRole with this ARN',
+                        input_restrictions={' '}
+                    )
+                ),
+            ]
+        )
+
+    def _dispatch(self, alert, descriptor):
+        """Send alert to a Lambda function
+
+        The alert gets dumped to a JSON string to be sent to the Lambda function
+
+        Publishing:
+            By default this output sends the JSON-serialized alert record as the payload to the
+            lambda function. You can override this:
+
+            - @aws-lambda.alert_data (dict):
+                    Overrides the alert record. Will instead send this dict, JSON-serialized, to
+                    Lambda as the payload.
+
+        Args:
+            alert (Alert): Alert instance which triggered a rule
+            descriptor (str): Output descriptor
+
+        Returns:
+            bool: True if alert was sent successfully, False otherwise
+        """
+        creds = self._load_creds(descriptor)
+        if not creds:
+            LOGGER.error("No credentials found for descriptor: %s", descriptor)
+            return False
+
+        # Create the publication
+        publication = compose_alert(alert, self, descriptor)
+
+        # Defaults
+        default_alert_data = alert.record
+
+        # Override with publisher
+        alert_data = publication.get('@aws-lambda.alert_data', default_alert_data)
+        alert_string = json.dumps(alert_data, separators=(',', ':'))
+
+        client = self._build_client(creds)
+
+        function_name = creds['lambda_function_arn']
+        qualifier = creds.get('function_qualifier', False)
+
+        LOGGER.debug('Sending alert to Lambda function %s', function_name)
+        invocation_opts = {
+            'FunctionName': function_name,
+            'InvocationType': 'Event',
+            'Payload': alert_string,
+        }
+
+        # Use the qualifier if it's available. Passing an empty qualifier in
+        # with `Qualifier=''` or `Qualifier=None` does not work
+        if qualifier:
+            invocation_opts['Qualifier'] = qualifier
+
+        client.invoke(**invocation_opts)
+
+        return True
+
+    def _build_client(self, creds):
+        """
+        Generates a boto3 client for the current AWS Lambda invocation. Will perform AssumeRole
+        if an assume role is provided.
+
+        Params:
+            creds (dict): Result of _load_creds()
+
+        Returns:
+            boto3.session.Session.client
+        """
+        client_opts = {
+            'region_name': self.region
+        }
+
+        assume_role_arn = creds.get('assume_role_arn', False)
+        if assume_role_arn:
+            LOGGER.debug('Assuming role: %s', assume_role_arn)
+            sts_connection = boto3.client('sts')
+            acct_b = sts_connection.assume_role(
+                RoleArn=assume_role_arn,
+                RoleSessionName="streamalert_alert_processor"
+            )
+
+            client_opts['aws_access_key_id'] = acct_b['Credentials']['AccessKeyId']
+            client_opts['aws_secret_access_key'] = acct_b['Credentials']['SecretAccessKey']
+            client_opts['aws_session_token'] = acct_b['Credentials']['SessionToken']
+
+        return boto3.client(
+            'lambda',
+            **client_opts
+        )
+
+
 @StreamAlertOutput
 class S3Output(AWSOutput):
     """S3Output handles all alert dispatching for AWS S3"""
@@ -481,6 +635,7 @@ def _dispatch(self, alert, descriptor):
 
         return True
 
+
 @StreamAlertOutput
 class SESOutput(OutputDispatcher):
     """Handle all alert dispatching for AWS SES"""

tests/unit/streamalert/alert_processor/outputs/test_aws.py

@@ -37,7 +37,8 @@
     SESOutput,
     SNSOutput,
     SQSOutput,
-    CloudwatchLogOutput
+    CloudwatchLogOutput,
+    LambdaOutputV2,
 )
 from tests.unit.streamalert.alert_processor import (
     CONFIG,
@@ -151,6 +152,44 @@ def test_dispatch_with_qualifier(self, log_mock):
                                     self.SERVICE, alt_descriptor)
 
 
+@patch.object(aws_outputs, 'boto3', MagicMock())
+class TestLambdaV2Output:
+    """Test class for LambdaOutput"""
+    DESCRIPTOR = 'unit_test_lambda'
+    SERVICE = 'aws-lambda-v2'
+    OUTPUT = ':'.join([SERVICE, DESCRIPTOR])
+    CREDS = {
+        'lambda_function_arn': 'arn:aws:lambda:us-east-1:11111111:function:my_func',
+        'function_qualifier': 'production',
+        'assume_role_arn': 'arn:aws:iam::11111111:role/my_path/my_role',
+    }
+
+    @patch('streamalert.alert_processor.outputs.output_base.OutputCredentialsProvider')
+    def setup(self, provider_constructor):
+        """Setup before each method"""
+        provider = MagicMock()
+        provider_constructor.return_value = provider
+        provider.load_credentials = Mock(
+            side_effect=lambda x: self.CREDS if x == self.DESCRIPTOR else None
+        )
+
+        self._provider = provider
+        self._dispatcher = LambdaOutputV2(None)
+
+    def test_locals(self):
+        """LambdaOutput local variables"""
+        assert_equal(self._dispatcher.__class__.__name__, 'LambdaOutputV2')
+        assert_equal(self._dispatcher.__service__, self.SERVICE)
+
+    @patch('logging.Logger.info')
+    def test_dispatch(self, log_mock):
+        """LambdaOutput dispatch"""
+        assert_true(self._dispatcher.dispatch(get_alert(), self.OUTPUT))
+
+        log_mock.assert_called_with('Successfully sent alert to %s:%s',
+                                    self.SERVICE, self.DESCRIPTOR)
+
+
 @mock_s3
 class TestS3Output:
     """Test class for S3Output"""

tests/unit/streamalert/alert_processor/outputs/test_output_base.py

@@ -91,6 +91,7 @@ def test_output_loading():
     expected_outputs = {
         'aws-firehose',
         'aws-lambda',
+        'aws-lambda-v2',
         'aws-s3',
         'aws-ses',
         'aws-sns',