Merge pull request #804 from edx/e0d/dns-refactor

Refactoring

util/vpc-tools/vpc_dns.py

@@ -26,61 +26,84 @@ import argparse
 import boto
 import datetime
 from vpcutil import vpc_for_stack_name
+import xml.dom.minidom
+import re
 
 r53 = boto.connect_route53()
 
 
-def add_or_update_record(zone, record_name, record_type,
+extra_play_dns = {"edxapp":["courses","studio"]}
+
+class DNSRecord():
+
+    def __init__(self, zone, record_name, record_type,
                          record_ttl, record_values):
+        self.zone = zone
+        self.record_name = record_name
+        self.record_type = record_type
+        self.record_ttl = record_ttl
+        self.record_values = record_values
+
+def add_or_update_record(dns_records):
     """
     Creates or updates a DNS record in a hosted route53
     zone
     """
+    change_set = boto.route53.record.ResourceRecordSets()
 
-    status_msg = """
-    record_name:   {}
-    record_type:   {}
-    record_ttl:    {}
-    record_values: {}
-             """.format(record_name, record_type,
-                        record_ttl, record_values)
+    for record in dns_records:
 
-    if args.noop:
-        print("Would have updated DNS record:\n{}".format(status_msg))
-        return
+        status_msg = """
+        record_name:   {}
+        record_type:   {}
+        record_ttl:    {}
+        record_values: {}
+                 """.format(record.record_name, record.record_type,
+                            record.record_ttl, record.record_values)
 
-    zone_id = zone.Id.replace("/hostedzone/", "")
+        if args.noop:
+            print("Would have updated DNS record:\n{}".format(status_msg))
 
-    records = r53.get_all_rrsets(zone_id)
+        zone_id = record.zone.Id.replace("/hostedzone/", "")
 
-    old_records = {r.name[:-1]: r for r in records}
+        records = r53.get_all_rrsets(zone_id)
 
-    change_set = boto.route53.record.ResourceRecordSets()
+        old_records = {r.name[:-1]: r for r in records}
 
-    # If the record name already points to something.
-    # Delete the existing connection.
-    if record_name in old_records.keys():
-        print("Deleting record:\n{}".format(status_msg))
-        change = change_set.add_change(
-            'DELETE',
-            record_name,
-            record_type,
-            record_ttl)
+        # If the record name already points to something.
+        # Delete the existing connection.
+        if record.record_name in old_records.keys():
+            if args.force:
+                print("Deleting record:\n{}".format(status_msg))
+                change = change_set.add_change(
+                    'DELETE',
+                    record.record_name,
+                    record.record_type,
+                    record.record_ttl)
+            else:
+                raise RuntimeError(
+                    "DNS record exists for {} and force was not specified.".
+                    format(record.record_name))
 
-        for value in old_records[record_name].resource_records:
-            change.add_value(value)
+            for value in old_records[record.record_name].resource_records:
+                change.add_value(value)
 
-    change = change_set.add_change(
-        'CREATE',
-        record_name,
-        record_type,
-        record_ttl)
+        change = change_set.add_change(
+            'CREATE',
+            record.record_name,
+            record.record_type,
+            record.record_ttl)
 
-    for value in record_values:
-        change.add_value(value)
+        for value in record.record_values:
+            change.add_value(value)
 
-    r53.change_rrsets(zone_id, change_set.to_xml())
-    print("Updated DNS record:\n{}".format(status_msg))
+    if args.noop:
+        print("Would have submitted the following change set:\n")
+        xml_doc = xml.dom.minidom.parseString(change_set.to_xml())
+        print xml_doc.toprettyxml()
+    else:
+        r53.change_rrsets(zone_id, change_set.to_xml())
+        print("Updated DNS record:\n{}".format(status_msg))
 
 
 def get_or_create_hosted_zone(zone_name):
@@ -112,12 +135,42 @@ def get_or_create_hosted_zone(zone_name):
 
     if parent_zone:
         print("Updating parent zone {}".format(parent_zone_name))
-        add_or_update_record(parent_zone,
-                             zone_name, 'NS', 900,
-                             zone.NameServers)
+
+        dns_records = set()
+        dns_records.add(DNSRecord(parent_zone,zone_name,'NS',900,zone.NameServers))
+        add_or_update_record(dns_records)
 
     return zone
 
+def get_security_group_dns(group_name):
+    # stage-edx-RabbitMQELBSecurityGroup-YB8ZKIZYN1EN
+    environment,deployment,sec_group,salt = group_name.split('-')
+    play = sec_group.replace("ELBSecurityGroup","").lower()
+    return environment, deployment, play
+
+def get_dns_from_instances(elb):
+
+    ec2_con = boto.connect_ec2()
+
+    for inst in elb.instances:
+        instance = ec2_con.get_all_instances(
+                instance_ids=[inst.id])[0].instances[0]
+        try:
+            env_tag = instance.tags['environment']
+            if 'play' in instance.tags:
+                play_tag = instance.tags['play']
+            else:
+                # deprecated, for backwards compatibility
+                play_tag = instance.tags['role']
+
+            break  # only need the first instance for tag info
+        except KeyError:
+            print("Instance {}, attached to elb {} does not "
+                  "have tags for environment and play".format(elb, inst))
+            raise
+
+    return env_tag, play_tag
+
 
 def update_elb_rds_dns(zone):
     """
@@ -127,9 +180,11 @@ def update_elb_rds_dns(zone):
     to the ELBs to create the dns name
     """
 
+    dns_records = set()
+
     elb_con = boto.connect_elb()
-    ec2_con = boto.connect_ec2()
     rds_con = boto.connect_rds()
+
     vpc_id = vpc_for_stack_name(args.stack_name)
 
     if not zone and args.noop:
@@ -139,41 +194,38 @@ def update_elb_rds_dns(zone):
     else:
         zone_name = zone.Name[:-1]
 
+    stack_elbs = [elb for elb in elb_con.get_all_load_balancers()
+                  if elb.vpc_id == vpc_id]
+
+    for elb in stack_elbs:
+
+        if "RabbitMQ" in elb.source_security_group.name or "ElasticSearch" in elb.source_security_group.name:
+            env_tag,deployment,play_tag = get_security_group_dns(elb.source_security_group.name)
+            fqdn = "{}-{}.{}".format(env_tag, play_tag, zone_name)
+            dns_records.add(DNSRecord(zone,fqdn,'CNAME',600,[elb.dns_name]))
+        else:
+            env_tag,play_tag = get_dns_from_instances(elb)
+            fqdn = "{}-{}.{}".format(env_tag, play_tag, zone_name)
+            dns_records.add(DNSRecord(zone,fqdn,'CNAME',600,[elb.dns_name]))
+
+        if extra_play_dns.has_key(play_tag):
+            for name in extra_play_dns.get(play_tag):
+                fqdn = "{}-{}.{}".format(env_tag, name, zone_name)
+                dns_records.add(DNSRecord(zone,fqdn,'CNAME',600,[elb.dns_name]))
+
+
     stack_rdss = [rds for rds in rds_con.get_all_dbinstances()
                   if hasattr(rds.subnet_group, 'vpc_id') and
                   rds.subnet_group.vpc_id == vpc_id]
-    for rds in stack_rdss:
-        fqdn = "{}.{}".format('rds', zone_name)
-        add_or_update_record(zone, fqdn, 'CNAME', 600,
-                             [stack_rdss[0].endpoint[0]])
 
-    stack_elbs = [elb for elb in elb_con.get_all_load_balancers()
-                  if elb.vpc_id == vpc_id]
+    # TODO the current version of the RDS API doesn't support
+    # looking up RDS instance tags.  Hence, we are using the 
+    # env_tag that was set via the loop over instances above.
+    for rds in stack_rdss:
+        fqdn = "{}-{}.{}".format(env_tag,'rds', zone_name)
+        dns_records.add(DNSRecord(zone,fqdn,'CNAME',600,[stack_rdss[0].endpoint[0]]))
 
-    for elb in stack_elbs:
-        for inst in elb.instances:
-            instance = ec2_con.get_all_instances(
-                instance_ids=[inst.id])[0].instances[0]
-            try:
-                env_tag = instance.tags['environment']
-                if 'play' in instance.tags:
-                    play_tag = instance.tags['play']
-                else:
-                    # deprecated, for backwards compatibility
-                    play_tag = instance.tags['role']
-                fqdn = "{}-{}.{}".format(env_tag, play_tag, zone_name)
-                add_or_update_record(zone, fqdn, 'CNAME', 600, [elb.dns_name])
-                if play_tag == 'edxapp':
-                    # create courses and studio CNAME records for edxapp
-                    for name in ['courses', 'studio']:
-                        fqdn = "{}-{}.{}".format(env_tag, name, zone_name)
-                        add_or_update_record(zone, fqdn, 'CNAME',
-                                             600, [elb.dns_name])
-                break  # only need the first instance for tag info
-            except KeyError:
-                print("Instance {}, attached to elb {} does not "
-                      "have tags for environment and play".format(elb, inst))
-                raise
+    add_or_update_record(dns_records)
 
 if __name__ == "__main__":
     description = "Give a cloudformation stack name, for an edx stack, setup \
@@ -188,7 +240,11 @@ if __name__ == "__main__":
     parser.add_argument('-z', '--zone-name', default="vpc.edx.org",
                         help="The name of the zone under which to "
                              "create the dns entries.")
+    parser.add_argument('-f', '--force',
+                        help="Force reuse of an existing name in a zone",
+                        action="store_true",default=False)
 
     args = parser.parse_args()
     zone = get_or_create_hosted_zone(args.zone_name)
     update_elb_rds_dns(zone)
+