Adding ability to modify ELBv2 endpoints. (#624)

lemur/plugins/lemur_aws/elb.py

@@ -49,16 +49,6 @@ def is_valid(listener_tuple):
     return listener_tuple
 
 
-@sts_client('elb')
-@retry(retry_on_exception=retry_throttled, stop_max_attempt_number=7, wait_exponential_multiplier=1000)
-def get_elbs(**kwargs):
-    """
-    Fetches one page elb objects for a given account and region.
-    """
-    client = kwargs.pop('client')
-    return client.describe_load_balancers(**kwargs)
-
-
 def get_all_elbs(**kwargs):
     """
     Fetches all elbs for a given account/region
@@ -80,6 +70,80 @@ def get_all_elbs(**kwargs):
             kwargs.update(dict(marker=response['NextMarker']))
 
 
+def get_all_elbs_v2(**kwargs):
+    """
+    Fetches all elbs for a given account/region
+
+    :param kwargs:
+    :return:
+    """
+    elbs = []
+
+    while True:
+        response = get_elbs_v2(**kwargs)
+        elbs += response['LoadBalancers']
+
+        if not response.get('IsTruncated'):
+            return elbs
+
+        if response['NextMarker']:
+            kwargs.update(dict(marker=response['NextMarker']))
+
+
+@sts_client('elbv2')
+@retry(retry_on_exception=retry_throttled, stop_max_attempt_number=7, wait_exponential_multiplier=1000)
+def get_listener_arn_from_endpoint(endpoint_name, endpoint_port, **kwargs):
+    """
+    Get a listener ARN from a endpoint.
+    :param endpoint_name:
+    :param endpoint_port:
+    :return:
+    """
+    client = kwargs.pop('client')
+    elbs = client.describe_load_balancers(Names=[endpoint_name])
+    for elb in elbs['LoadBalancers']:
+        listeners = client.describe_listeners(LoadBalancerArn=elb['LoadBalancerArn'])
+        for listener in listeners['Listeners']:
+            if listener['Port'] == endpoint_port:
+                return listener['ListenerArn']
+
+
+@sts_client('elb')
+@retry(retry_on_exception=retry_throttled, stop_max_attempt_number=7, wait_exponential_multiplier=1000)
+def get_elbs(**kwargs):
+    """
+    Fetches one page elb objects for a given account and region.
+    """
+    client = kwargs.pop('client')
+    return client.describe_load_balancers(**kwargs)
+
+
+@sts_client('elbv2')
+@retry(retry_on_exception=retry_throttled, stop_max_attempt_number=7, wait_exponential_multiplier=1000)
+def get_elbs_v2(**kwargs):
+    """
+    Fetches one page of elb objects for a given account and region.
+
+    :param kwargs:
+    :return:
+    """
+    client = kwargs.pop('client')
+    return client.describe_load_balancers(**kwargs)
+
+
+@sts_client('elbv2')
+@retry(retry_on_exception=retry_throttled, stop_max_attempt_number=7, wait_exponential_multiplier=1000)
+def describe_listeners_v2(**kwargs):
+    """
+    Fetches one page of listener objects for a given elb arn.
+
+    :param kwargs:
+    :return:
+    """
+    client = kwargs.pop('client')
+    return client.describe_listeners(**kwargs)
+
+
 @sts_client('elb')
 def describe_load_balancer_policies(load_balancer_name, policy_names, **kwargs):
     """
@@ -91,6 +155,17 @@ def describe_load_balancer_policies(load_balancer_name, policy_names, **kwargs):
     return kwargs['client'].describe_load_balancer_policies(LoadBalancerName=load_balancer_name, PolicyNames=policy_names)
 
 
+@sts_client('elbv2')
+def describe_ssl_policies_v2(policy_names, **kwargs):
+    """
+    Fetching all policies currently associated with an ELB.
+
+    :param policy_names:
+    :return:
+    """
+    return kwargs['client'].describe_ssl_policies(Names=policy_names)
+
+
 @sts_client('elb')
 def describe_load_balancer_types(policies, **kwargs):
     """
@@ -120,3 +195,23 @@ def attach_certificate(name, port, certificate_id, **kwargs):
             current_app.logger.warning("Loadbalancer does not exist.")
         else:
             raise e
+
+
+@sts_client('elbv2')
+@retry(retry_on_exception=retry_throttled, stop_max_attempt_number=7, wait_exponential_multiplier=1000)
+def attach_certificate_v2(listener_arn, port, certificates, **kwargs):
+    """
+    Attaches a certificate to a listener, throws exception
+    if certificate specified does not exist in a particular account.
+
+    :param listener_arn:
+    :param port:
+    :param certificates:
+    """
+    try:
+        return kwargs['client'].modify_listener(ListenerArn=listener_arn, Port=port, Certificates=certificates)
+    except botocore.exceptions.ClientError as e:
+        if e.response['Error']['Code'] == 'LoadBalancerNotFound':
+            current_app.logger.warning("Loadbalancer does not exist.")
+        else:
+            raise e

lemur/plugins/lemur_aws/plugin.py

@@ -44,6 +44,108 @@ def get_region_from_dns(dns):
     return dns.split('.')[-4]
 
 
+def format_elb_cipher_policy_v2(policy):
+    """
+    Attempts to format cipher policy information for elbv2 into a common format.
+    :param policy:
+    :return:
+    """
+    ciphers = []
+    name = None
+
+    for descr in policy['SslPolicies']:
+        name = descr['Name']
+        for cipher in descr['Ciphers']:
+            ciphers.append(cipher['Name'])
+
+    return dict(name=name, ciphers=ciphers)
+
+
+def format_elb_cipher_policy(policy):
+    """
+    Attempts to format cipher policy information into a common format.
+    :param policy:
+    :return:
+    """
+    ciphers = []
+    name = None
+    for descr in policy['PolicyDescriptions']:
+        for attr in descr['PolicyAttributeDescriptions']:
+            if attr['AttributeName'] == 'Reference-Security-Policy':
+                name = attr['AttributeValue']
+                continue
+
+            if attr['AttributeValue'] == 'true':
+                ciphers.append(attr['AttributeName'])
+
+    return dict(name=name, ciphers=ciphers)
+
+
+def get_elb_endpoints(account_number, region, elb_dict):
+    """
+    Retrieves endpoint information from elb response data.
+    :param account_number:
+    :param region:
+    :param elb_dict:
+    :return:
+    """
+    endpoints = []
+    for listener in elb_dict['ListenerDescriptions']:
+        if not listener['Listener'].get('SSLCertificateId'):
+            continue
+
+        if listener['Listener']['SSLCertificateId'] == 'Invalid-Certificate':
+            continue
+
+        endpoint = dict(
+            name=elb_dict['LoadBalancerName'],
+            dnsname=elb_dict['DNSName'],
+            type='elb',
+            port=listener['Listener']['LoadBalancerPort'],
+            certificate_name=iam.get_name_from_arn(listener['Listener']['SSLCertificateId'])
+        )
+
+        if listener['PolicyNames']:
+            policy = elb.describe_load_balancer_policies(elb_dict['LoadBalancerName'], listener['PolicyNames'], account_number=account_number, region=region)
+            endpoint['policy'] = format_elb_cipher_policy(policy)
+
+        endpoints.append(endpoint)
+
+    return endpoints
+
+
+def get_elb_endpoints_v2(account_number, region, elb_dict):
+    """
+    Retrieves endpoint information from elbv2 response data.
+    :param account_number:
+    :param region:
+    :param elb_dict:
+    :return:
+    """
+    endpoints = []
+    listeners = elb.describe_listeners_v2(account_number=account_number, region=region, LoadBalancerArn=elb_dict['LoadBalancerArn'])
+    for listener in listeners['Listeners']:
+        if not listener['Certificates']:
+            continue
+
+        for certificate in listener['Certificates']:
+            endpoint = dict(
+                name=elb_dict['LoadBalancerName'],
+                dnsname=elb_dict['DNSName'],
+                type='elbv2',
+                port=listener['Port'],
+                certificate_name=iam.get_name_from_arn(certificate['CertificateArn'])
+            )
+
+        if listener['SslPolicy']:
+            policy = elb.describe_ssl_policies_v2([listener['SslPolicy']], account_number=account_number, region=region)
+            endpoint['policy'] = format_elb_cipher_policy_v2(policy)
+
+        endpoints.append(endpoint)
+
+    return endpoints
+
+
 class AWSDestinationPlugin(DestinationPlugin):
     title = 'AWS'
     slug = 'aws-destination'
@@ -77,10 +179,6 @@ class AWSDestinationPlugin(DestinationPlugin):
             if e.error_code != 'EntityAlreadyExists':
                 raise Exception(e)
 
-        e = self.get_option('elb', options)
-        if e:
-            iam.attach_certificate(kwargs['accountNumber'], ['region'], e['name'], e['port'], e['certificateId'])
-
     def deploy(self, elb_name, account, region, certificate):
         pass
 
@@ -135,28 +233,17 @@ class AWSSourcePlugin(SourcePlugin):
 
         for region in regions:
             elbs = elb.get_all_elbs(account_number=account_number, region=region)
-            current_app.logger.info("Describing load balancers in {0}-{1}".format(account_number, region))
+            current_app.logger.info("Describing classic load balancers in {0}-{1}".format(account_number, region))
+
             for e in elbs:
-                for listener in e['ListenerDescriptions']:
+                endpoints.extend(get_elb_endpoints(account_number, region, e))
-                    if not listener['Listener'].get('SSLCertificateId'):
-                        continue
 
-                    if listener['Listener']['SSLCertificateId'] == 'Invalid-Certificate':
+            # fetch advanced ELBs
-                        continue
+            elbs_v2 = elb.get_all_elbs_v2(account_number=account_number, region=region)
+            current_app.logger.info("Describing advanced load balancers in {0}-{1}".format(account_number, region))
 
-                    endpoint = dict(
+            for e in elbs_v2:
-                        name=e['LoadBalancerName'],
+                endpoints.extend(get_elb_endpoints_v2(account_number, region, e))
-                        dnsname=e['DNSName'],
-                        type='elb',
-                        port=listener['Listener']['LoadBalancerPort'],
-                        certificate_name=iam.get_name_from_arn(listener['Listener']['SSLCertificateId'])
-                    )
-
-                    if listener['PolicyNames']:
-                        policy = elb.describe_load_balancer_policies(e['LoadBalancerName'], listener['PolicyNames'], account_number=account_number, region=region)
-                        endpoint['policy'] = format_elb_cipher_policy(policy)
-
-                    endpoints.append(endpoint)
 
         return endpoints
 
@@ -167,6 +254,11 @@ class AWSSourcePlugin(SourcePlugin):
         # relies on the fact that region is included in DNS name
         region = get_region_from_dns(endpoint.dnsname)
         arn = iam.create_arn_from_cert(account_number, region, certificate.name)
+
+        if endpoint.type == 'elbv2':
+            listener_arn = elb.get_listener_arn_from_endpoint(endpoint.name, endpoint.port, account_number=account_number, region=region)
+            elb.attach_certificate_v2(listener_arn, endpoint.port, [{'CertificateArn': arn}], account_number=account_number, region=region)
+        else:
             elb.attach_certificate(endpoint.name, endpoint.port, arn, account_number=account_number, region=region)
 
     def clean(self, options, **kwargs):
@@ -186,26 +278,6 @@ class AWSSourcePlugin(SourcePlugin):
         return orphaned
 
 
-def format_elb_cipher_policy(policy):
-    """
-    Attempts to format cipher policy information into a common format.
-    :param policy:
-    :return:
-    """
-    ciphers = []
-    name = None
-    for descr in policy['PolicyDescriptions']:
-        for attr in descr['PolicyAttributeDescriptions']:
-            if attr['AttributeName'] == 'Reference-Security-Policy':
-                name = attr['AttributeValue']
-                continue
-
-            if attr['AttributeValue'] == 'true':
-                ciphers.append(attr['AttributeName'])
-
-    return dict(name=name, ciphers=ciphers)
-
-
 class S3DestinationPlugin(DestinationPlugin):
     title = 'AWS-S3'
     slug = 'aws-s3'

lemur/sources/service.py

@@ -169,13 +169,13 @@ def sync_certificates(source, user):
 
 
 def sync(source, user):
-    new, updated = sync_certificates(source, user)
+    new_certs, updated_certs = sync_certificates(source, user)
-    new, updated = sync_endpoints(source)
+    new_endpoints, updated_endpoints = sync_endpoints(source)
 
     source.last_run = arrow.utcnow()
     database.update(source)
 
-    return {'endpoints': (new, updated), 'certificates': (new, updated)}
+    return {'endpoints': (new_endpoints, updated_endpoints), 'certificates': (new_certs, updated_certs)}
 
 
 def clean(source):