Merge pull request #2365 from spulec/elb-cognito

Add cognito support to Elb

moto/elbv2/models.py

@@ -2,9 +2,11 @@ from __future__ import unicode_literals
 
 import datetime
 import re
+from jinja2 import Template
 from moto.compat import OrderedDict
 from moto.core.exceptions import RESTError
 from moto.core import BaseBackend, BaseModel
+from moto.core.utils import camelcase_to_underscores
 from moto.ec2.models import ec2_backends
 from moto.acm.models import acm_backends
 from .utils import make_arn_for_target_group
@@ -213,13 +215,12 @@ class FakeListener(BaseModel):
                 action_type = action['Type']
                 if action_type == 'forward':
                     default_actions.append({'type': action_type, 'target_group_arn': action['TargetGroupArn']})
-                elif action_type == 'redirect':
+                elif action_type in ['redirect', 'authenticate-cognito']:
-                    redirect_action = {'type': action_type, }
+                    redirect_action = {'type': action_type}
-                    for redirect_config_key, redirect_config_value in action['RedirectConfig'].items():
+                    key = 'RedirectConfig' if action_type == 'redirect' else 'AuthenticateCognitoConfig'
+                    for redirect_config_key, redirect_config_value in action[key].items():
                         # need to match the output of _get_list_prefix
-                        if redirect_config_key == 'StatusCode':
+                        redirect_action[camelcase_to_underscores(key) + '._' + camelcase_to_underscores(redirect_config_key)] = redirect_config_value
-                            redirect_config_key = 'status_code'
-                        redirect_action['redirect_config._' + redirect_config_key.lower()] = redirect_config_value
                     default_actions.append(redirect_action)
                 else:
                     raise InvalidActionTypeError(action_type, i + 1)
@@ -231,6 +232,32 @@ class FakeListener(BaseModel):
         return listener
 
 
+class FakeAction(BaseModel):
+    def __init__(self, data):
+        self.data = data
+        self.type = data.get("type")
+
+    def to_xml(self):
+        template = Template("""<Type>{{ action.type }}</Type>
+            {% if action.type == "forward" %}
+            <TargetGroupArn>{{ action.data["target_group_arn"] }}</TargetGroupArn>
+            {% elif action.type == "redirect" %}
+            <RedirectConfig>
+                <Protocol>{{ action.data["redirect_config._protocol"] }}</Protocol>
+                <Port>{{ action.data["redirect_config._port"] }}</Port>
+                <StatusCode>{{ action.data["redirect_config._status_code"] }}</StatusCode>
+            </RedirectConfig>
+            {% elif action.type == "authenticate-cognito" %}
+            <AuthenticateCognitoConfig>
+                <UserPoolArn>{{ action.data["authenticate_cognito_config._user_pool_arn"] }}</UserPoolArn>
+                <UserPoolClientId>{{ action.data["authenticate_cognito_config._user_pool_client_id"] }}</UserPoolClientId>
+                <UserPoolDomain>{{ action.data["authenticate_cognito_config._user_pool_domain"] }}</UserPoolDomain>
+            </AuthenticateCognitoConfig>
+            {% endif %}
+            """)
+        return template.render(action=self)
+
+
 class FakeRule(BaseModel):
 
     def __init__(self, listener_arn, conditions, priority, actions, is_default):
@@ -402,6 +429,7 @@ class ELBv2Backend(BaseBackend):
         return new_load_balancer
 
     def create_rule(self, listener_arn, conditions, priority, actions):
+        actions = [FakeAction(action) for action in actions]
         listeners = self.describe_listeners(None, [listener_arn])
         if not listeners:
             raise ListenerNotFoundError()
@@ -429,20 +457,7 @@ class ELBv2Backend(BaseBackend):
             if rule.priority == priority:
                 raise PriorityInUseError()
 
-        # validate Actions
+        self._validate_actions(actions)
-        target_group_arns = [target_group.arn for target_group in self.target_groups.values()]
-        for i, action in enumerate(actions):
-            index = i + 1
-            action_type = action['type']
-            if action_type == 'forward':
-                action_target_group_arn = action['target_group_arn']
-                if action_target_group_arn not in target_group_arns:
-                    raise ActionTargetGroupNotFoundError(action_target_group_arn)
-            elif action_type == 'redirect':
-                # nothing to do
-                pass
-            else:
-                raise InvalidActionTypeError(action_type, index)
 
         # TODO: check for error 'TooManyRegistrationsForTargetId'
         # TODO: check for error 'TooManyRules'
@@ -452,6 +467,21 @@ class ELBv2Backend(BaseBackend):
         listener.register(rule)
         return [rule]
 
+    def _validate_actions(self, actions):
+        # validate Actions
+        target_group_arns = [target_group.arn for target_group in self.target_groups.values()]
+        for i, action in enumerate(actions):
+            index = i + 1
+            action_type = action.type
+            if action_type == 'forward':
+                action_target_group_arn = action.data['target_group_arn']
+                if action_target_group_arn not in target_group_arns:
+                    raise ActionTargetGroupNotFoundError(action_target_group_arn)
+            elif action_type in ['redirect', 'authenticate-cognito']:
+                pass
+            else:
+                raise InvalidActionTypeError(action_type, index)
+
     def create_target_group(self, name, **kwargs):
         if len(name) > 32:
             raise InvalidTargetGroupNameError(
@@ -495,26 +525,22 @@ class ELBv2Backend(BaseBackend):
         return target_group
 
     def create_listener(self, load_balancer_arn, protocol, port, ssl_policy, certificate, default_actions):
+        default_actions = [FakeAction(action) for action in default_actions]
         balancer = self.load_balancers.get(load_balancer_arn)
         if balancer is None:
             raise LoadBalancerNotFoundError()
         if port in balancer.listeners:
             raise DuplicateListenerError()
 
+        self._validate_actions(default_actions)
+
         arn = load_balancer_arn.replace(':loadbalancer/', ':listener/') + "/%s%s" % (port, id(self))
         listener = FakeListener(load_balancer_arn, arn, protocol, port, ssl_policy, certificate, default_actions)
         balancer.listeners[listener.arn] = listener
-        for i, action in enumerate(default_actions):
+        for action in default_actions:
-            action_type = action['type']
+            if action.type == 'forward':
-            if action_type == 'forward':
+                target_group = self.target_groups[action.data['target_group_arn']]
-                if action['target_group_arn'] in self.target_groups.keys():
+                target_group.load_balancer_arns.append(load_balancer_arn)
-                    target_group = self.target_groups[action['target_group_arn']]
-                    target_group.load_balancer_arns.append(load_balancer_arn)
-            elif action_type == 'redirect':
-                # nothing to do
-                pass
-            else:
-                raise InvalidActionTypeError(action_type, i + 1)
 
         return listener
 
@@ -648,6 +674,7 @@ class ELBv2Backend(BaseBackend):
         raise ListenerNotFoundError()
 
     def modify_rule(self, rule_arn, conditions, actions):
+        actions = [FakeAction(action) for action in actions]
         # if conditions or actions is empty list, do not update the attributes
         if not conditions and not actions:
             raise InvalidModifyRuleArgumentsError()
@@ -673,20 +700,7 @@ class ELBv2Backend(BaseBackend):
                 # TODO: check pattern of value for 'path-pattern'
 
         # validate Actions
-        target_group_arns = [target_group.arn for target_group in self.target_groups.values()]
+        self._validate_actions(actions)
-        if actions:
-            for i, action in enumerate(actions):
-                index = i + 1
-                action_type = action['type']
-                if action_type == 'forward':
-                    action_target_group_arn = action['target_group_arn']
-                    if action_target_group_arn not in target_group_arns:
-                        raise ActionTargetGroupNotFoundError(action_target_group_arn)
-                elif action_type == 'redirect':
-                    # nothing to do
-                    pass
-                else:
-                    raise InvalidActionTypeError(action_type, index)
 
         # TODO: check for error 'TooManyRegistrationsForTargetId'
         # TODO: check for error 'TooManyRules'
@@ -851,6 +865,7 @@ class ELBv2Backend(BaseBackend):
         return target_group
 
     def modify_listener(self, arn, port=None, protocol=None, ssl_policy=None, certificates=None, default_actions=None):
+        default_actions = [FakeAction(action) for action in default_actions]
         for load_balancer in self.load_balancers.values():
             if arn in load_balancer.listeners:
                 break
@@ -917,7 +932,7 @@ class ELBv2Backend(BaseBackend):
             for listener in load_balancer.listeners.values():
                 for rule in listener.rules:
                     for action in rule.actions:
-                        if action.get('target_group_arn') == target_group_arn:
+                        if action.data.get('target_group_arn') == target_group_arn:
                             return True
         return False
 

moto/elbv2/responses.py

@@ -775,16 +775,7 @@ CREATE_LISTENER_TEMPLATE = """<CreateListenerResponse xmlns="http://elasticloadb
         <DefaultActions>
           {% for action in listener.default_actions %}
           <member>
-            <Type>{{ action.type }}</Type>
+            {{ action.to_xml() }}
-            {% if action["type"] == "forward" %}
-            <TargetGroupArn>{{ action["target_group_arn"] }}</TargetGroupArn>
-            {% elif action["type"] == "redirect" %}
-            <RedirectConfig>
-                <Protocol>{{ action["redirect_config._protocol"] }}</Protocol>
-                <Port>{{ action["redirect_config._port"] }}</Port>
-                <StatusCode>{{ action["redirect_config._status_code"] }}</StatusCode>
-            </RedirectConfig>
-            {% endif %}
           </member>
           {% endfor %}
         </DefaultActions>
@@ -888,16 +879,7 @@ DESCRIBE_RULES_TEMPLATE = """<DescribeRulesResponse xmlns="http://elasticloadbal
         <Actions>
           {% for action in rule.actions %}
           <member>
-            <Type>{{ action["type"] }}</Type>
+            {{ action.to_xml() }}
-            {% if action["type"] == "forward" %}
-            <TargetGroupArn>{{ action["target_group_arn"] }}</TargetGroupArn>
-            {% elif action["type"] == "redirect" %}
-            <RedirectConfig>
-                <Protocol>{{ action["redirect_config._protocol"] }}</Protocol>
-                <Port>{{ action["redirect_config._port"] }}</Port>
-                <StatusCode>{{ action["redirect_config._status_code"] }}</StatusCode>
-            </RedirectConfig>
-            {% endif %}
           </member>
           {% endfor %}
         </Actions>
@@ -989,16 +971,7 @@ DESCRIBE_LISTENERS_TEMPLATE = """<DescribeLoadBalancersResponse xmlns="http://el
         <DefaultActions>
           {% for action in listener.default_actions %}
           <member>
-            <Type>{{ action.type }}</Type>
+            {{ action.to_xml() }}
-            {% if action["type"] == "forward" %}
-            <TargetGroupArn>{{ action["target_group_arn"] }}</TargetGroupArn>m
-            {% elif action["type"] == "redirect" %}
-            <RedirectConfig>
-                <Protocol>{{ action["redirect_config._protocol"] }}</Protocol>
-                <Port>{{ action["redirect_config._port"] }}</Port>
-                <StatusCode>{{ action["redirect_config._status_code"] }}</StatusCode>
-            </RedirectConfig>
-            {% endif %}
           </member>
           {% endfor %}
         </DefaultActions>
@@ -1048,8 +1021,7 @@ MODIFY_RULE_TEMPLATE = """<ModifyRuleResponse xmlns="http://elasticloadbalancing
         <Actions>
           {% for action in rule.actions %}
           <member>
-            <Type>{{ action["type"] }}</Type>
+            {{ action.to_xml() }}
-            <TargetGroupArn>{{ action["target_group_arn"] }}</TargetGroupArn>
           </member>
           {% endfor %}
         </Actions>
@@ -1432,16 +1404,7 @@ MODIFY_LISTENER_TEMPLATE = """<ModifyListenerResponse xmlns="http://elasticloadb
         <DefaultActions>
           {% for action in listener.default_actions %}
           <member>
-            <Type>{{ action.type }}</Type>
+            {{ action.to_xml() }}
-            {% if action["type"] == "forward" %}
-            <TargetGroupArn>{{ action["target_group_arn"] }}</TargetGroupArn>
-            {% elif action["type"] == "redirect" %}
-            <RedirectConfig>
-                <Protocol>{{ action["redirect_config._protocol"] }}</Protocol>
-                <Port>{{ action["redirect_config._port"] }}</Port>
-                <StatusCode>{{ action["redirect_config._status_code"] }}</StatusCode>
-            </RedirectConfig>
-            {% endif %}
           </member>
           {% endfor %}
         </DefaultActions>

tests/test_elbv2/test_elbv2.py

@@ -1811,3 +1811,132 @@ def test_redirect_action_listener_rule_cloudformation():
             'Port': '443', 'Protocol': 'HTTPS', 'StatusCode': 'HTTP_301',
         }
     },])
+
+
+@mock_elbv2
+@mock_ec2
+def test_cognito_action_listener_rule():
+    conn = boto3.client('elbv2', region_name='us-east-1')
+    ec2 = boto3.resource('ec2', region_name='us-east-1')
+
+    security_group = ec2.create_security_group(
+        GroupName='a-security-group', Description='First One')
+    vpc = ec2.create_vpc(CidrBlock='172.28.7.0/24', InstanceTenancy='default')
+    subnet1 = ec2.create_subnet(
+        VpcId=vpc.id,
+        CidrBlock='172.28.7.192/26',
+        AvailabilityZone='us-east-1a')
+    subnet2 = ec2.create_subnet(
+        VpcId=vpc.id,
+        CidrBlock='172.28.7.128/26',
+        AvailabilityZone='us-east-1b')
+
+    response = conn.create_load_balancer(
+        Name='my-lb',
+        Subnets=[subnet1.id, subnet2.id],
+        SecurityGroups=[security_group.id],
+        Scheme='internal',
+        Tags=[{'Key': 'key_name', 'Value': 'a_value'}])
+    load_balancer_arn = response.get('LoadBalancers')[0].get('LoadBalancerArn')
+
+    action = {
+        'Type': 'authenticate-cognito',
+        'AuthenticateCognitoConfig': {
+             'UserPoolArn': 'arn:aws:cognito-idp:us-east-1:123456789012:userpool/us-east-1_ABCD1234',
+             'UserPoolClientId': 'abcd1234abcd',
+             'UserPoolDomain': 'testpool',
+         }
+    }
+    response = conn.create_listener(LoadBalancerArn=load_balancer_arn,
+                                    Protocol='HTTP',
+                                    Port=80,
+                                    DefaultActions=[action])
+
+    listener = response.get('Listeners')[0]
+    listener.get('DefaultActions')[0].should.equal(action)
+    listener_arn = listener.get('ListenerArn')
+
+    describe_rules_response = conn.describe_rules(ListenerArn=listener_arn)
+    describe_rules_response['Rules'][0]['Actions'][0].should.equal(action)
+
+    describe_listener_response = conn.describe_listeners(ListenerArns=[listener_arn, ])
+    describe_listener_actions = describe_listener_response['Listeners'][0]['DefaultActions'][0]
+    describe_listener_actions.should.equal(action)
+
+
+@mock_elbv2
+@mock_cloudformation
+def test_cognito_action_listener_rule_cloudformation():
+    cnf_conn = boto3.client('cloudformation', region_name='us-east-1')
+    elbv2_client = boto3.client('elbv2', region_name='us-east-1')
+
+    template = {
+        "AWSTemplateFormatVersion": "2010-09-09",
+        "Description": "ECS Cluster Test CloudFormation",
+        "Resources": {
+            "testVPC": {
+                "Type": "AWS::EC2::VPC",
+                "Properties": {
+                    "CidrBlock": "10.0.0.0/16",
+                },
+            },
+            "subnet1": {
+                "Type": "AWS::EC2::Subnet",
+                "Properties": {
+                    "CidrBlock": "10.0.0.0/24",
+                    "VpcId": {"Ref": "testVPC"},
+                    "AvalabilityZone": "us-east-1b",
+                },
+            },
+            "subnet2": {
+                "Type": "AWS::EC2::Subnet",
+                "Properties": {
+                    "CidrBlock": "10.0.1.0/24",
+                    "VpcId": {"Ref": "testVPC"},
+                    "AvalabilityZone": "us-east-1b",
+                },
+            },
+            "testLb": {
+                "Type": "AWS::ElasticLoadBalancingV2::LoadBalancer",
+                "Properties": {
+                    "Name": "my-lb",
+                    "Subnets": [{"Ref": "subnet1"}, {"Ref": "subnet2"}],
+                    "Type": "application",
+                    "SecurityGroups": [],
+                }
+            },
+            "testListener": {
+                "Type": "AWS::ElasticLoadBalancingV2::Listener",
+                "Properties": {
+                    "LoadBalancerArn": {"Ref": "testLb"},
+                    "Port": 80,
+                    "Protocol": "HTTP",
+                    "DefaultActions": [{
+                        "Type": "authenticate-cognito",
+                        "AuthenticateCognitoConfig": {
+                            'UserPoolArn': 'arn:aws:cognito-idp:us-east-1:123456789012:userpool/us-east-1_ABCD1234',
+                            'UserPoolClientId': 'abcd1234abcd',
+                            'UserPoolDomain': 'testpool',
+                        }
+                    }]
+                }
+
+            }
+        }
+    }
+    template_json = json.dumps(template)
+    cnf_conn.create_stack(StackName="test-stack", TemplateBody=template_json)
+
+    describe_load_balancers_response = elbv2_client.describe_load_balancers(Names=['my-lb',])
+    load_balancer_arn = describe_load_balancers_response['LoadBalancers'][0]['LoadBalancerArn']
+    describe_listeners_response = elbv2_client.describe_listeners(LoadBalancerArn=load_balancer_arn)
+
+    describe_listeners_response['Listeners'].should.have.length_of(1)
+    describe_listeners_response['Listeners'][0]['DefaultActions'].should.equal([{
+        'Type': 'authenticate-cognito',
+        "AuthenticateCognitoConfig": {
+            'UserPoolArn': 'arn:aws:cognito-idp:us-east-1:123456789012:userpool/us-east-1_ABCD1234',
+            'UserPoolClientId': 'abcd1234abcd',
+            'UserPoolDomain': 'testpool',
+        }
+    },])