update-tg-config.py

#!/usr/bin/python

import boto3
import requests
import sys
import re
import tempfile
import os
import difflib
import subprocess
import stat
import hashlib

def main():
  if len(sys.argv) != 4:
    print 'Usage: ' + sys.argv[0] + ' <service-name> <tg-name> <config-file>'
    return -1

  service_name = sys.argv[1]
  tg_name = sys.argv[2]
  config_file = sys.argv[3]

  if service_name not in ['wsrep', 'garb', 'drupal_varnish', 'drupal_memcache', 'varnish', 'efs', 'lsyncd', 'haproxy']:
    sys.stderr.write('Unknown service: ' + service_name + '\n')
    return -1

  prefix = get_prefix()
  tg_ips = get_tg_ips(prefix + tg_name)
  
  if len(tg_ips) == 0:
    print 'No IP addresses found in target group, aborting!'
    return -1
  
  with open(config_file) as file:
    config = file.read()
    file.close()
  new_config = globals()['generate_' + service_name + '_config'](config, tg_ips)
      
  #No update needed, so we can stop
  if new_config == config:
    return 0
  
  print 'Making the following change to ' + config_file + ':'
  for line in difflib.unified_diff(config.split("\n"), new_config.split("\n")):
    print line.strip()
  
  with open(config_file, 'w') as file:
    file.write(new_config)

  if 'load_' + service_name + '_config' in globals():
    globals()['load_' + service_name + '_config']()

  return 0

# Returns the region the instance is currently running in (eg: 'us-east')
def get_region():
  session_token = requests.put('http://169.254.169.254/latest/api/token', headers={'X-aws-ec2-metadata-token-ttl-seconds': 60}).text
  instance_info = requests.get('http://169.254.169.254/2018-09-24/dynamic/instance-identity/document/', headers={'X-aws-ec2-metadata-token': session_token}).json()
  return instance_info['region']

# Returns the instance_id of this VM.
def get_instance_id():
  with open('/var/lib/cloud/data/instance-id') as file:
    instance_id = file.read().strip()
  return instance_id

# Returns the value of the 'unlcms:prefix' tag on this VM.
def get_prefix():
  ec2 = boto3.resource('ec2', region_name=get_region())
  
  dept = ''
  campus = ''
  service = ''
  env = ''
  for tag in ec2.Instance(get_instance_id()).tags:
    if tag['Key'] == 'Department':
        dept = tag['Value']
    elif tag['Key'] == 'Environment':
        env = tag['Value']
    elif tag['Key'] == 'Campus':
        campus = tag['Value']
    elif tag['Key'] == 'Service':
        service = tag['Value']
  return "{}-{}-{}-{}-".format(
    dept,
    campus,
    service,
    env,
  )

# Returns a list of IP addresses registered to the named Target Group
def get_tg_ips(tg_name):
  region = get_region()
  elb = boto3.client('elbv2', region_name=region)
  ec2 = boto3.resource('ec2', region_name=region)
  
  ips = []

  tg_arn = elb.describe_target_groups(Names=[tg_name])['TargetGroups'][0]['TargetGroupArn']
  for instance in elb.describe_target_health(TargetGroupArn=tg_arn)['TargetHealthDescriptions']:
    ips.append(ec2.Instance(instance['Target']['Id']).private_ip_address)
  return ips

# Returns the updated contents of the wsrep.conf config file
def generate_wsrep_config(config, tg_ips):
  galera_nodes  = 'wsrep_cluster_address=gcomm://' + ','.join(tg_ips)
  
  pattern = re.compile(r'^wsrep_cluster_address=gcomm:\/\/.*$', re.MULTILINE)
  config = re.sub(pattern, galera_nodes, config)
  return config

# Updates GALERA_NODES with the list of IP Addresses
def generate_garb_config(config, tg_ips):
  galera_nodes  = 'GALERA_NODES="'
  galera_nodes += ','.join([ip + ':4567' for ip in tg_ips])
  galera_nodes += '"'
  
  pattern = re.compile(r'^GALERA_NODES="[^"]*"$', re.MULTILINE)
  config = re.sub(pattern, galera_nodes, config)
  return config

# Updates a file included by settings.php to configure varnish servers
def generate_drupal_varnish_config(config, tg_ips):
  config  = "<?php\n"
  config += "$conf['varnish_control_terminal'] = '"
  config += ' '.join([ip + ":6082" for ip in tg_ips])
  config += "';\n"
  return config

def generate_drupal_memcache_config(config, tg_ips):
  config  = "<?php\n"
  config += "$conf['memcache_servers'] = ["
  config += ','.join(["'" + ip + ":11211' => 'default'" for ip in tg_ips])
  config += '];\n'
  return config

def generate_varnish_config(config, tg_ips):
  NODE_BLOCK = (
    'backend node%s {\n'
    '  .host = "%s";\n'
    '}\n'
  )
  
  config = 'vcl 4.1;\n\n'
  for i, ip in enumerate(tg_ips):
    config += NODE_BLOCK % (i, ip)
  config += (
    'probe default {\n'
    '  .url = "/status.php";\n'
    '  .timeout = 30s;\n'
    '  .window = 6;\n'
    '  .threshold = 5;\n'
    '}\n'
    '\n'
    'import directors;\n'
    'sub vcl_init {\n'
    '  new vdir = directors.round_robin();\n'
  )
  for i, ip in enumerate(tg_ips):
    config += '  vdir.add_backend(node' + str(i) + ');\n'
  config += (
    '}\n'
    '\n'
    'sub vcl_recv {\n'
    '  set req.backend_hint = vdir.backend();\n'
    '}\n'
  )
  
  try:
    tmp_vcl = tempfile.NamedTemporaryFile(delete=False)
    tmp_vcl.write(config)
    tmp_vcl.close()
    os.chmod(tmp_vcl.name, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH)
    subprocess.check_call(['/sbin/varnishd', '-f', tmp_vcl.name, '-C'], stderr=open('/dev/null', 'w'))
  except:
    raise Exception('Error validating new VCL!')
  finally:
    os.unlink(tmp_vcl.name)
  
  return config

def generate_lsyncd_config(config, tg_ips):
  drupal_nodes  = "tg_member_ips = {'" + "','".join(tg_ips) + "'}"
  pattern = re.compile(r'^tg_member_ips = {[^}]*}$', re.MULTILINE)
  config = re.sub(pattern, drupal_nodes, config)
  return config

def generate_haproxy_config(config, tg_ips):
  server_list = (
    '  #BEGIN SERVER LIST\n'
    '  #DO NOT EDIT!\n'
  )
  for i, ip in enumerate(tg_ips):
    server_list += '  server %s %s:80\n' % (hashlib.sha256(ip).hexdigest()[:8], ip)
  server_list += '  #END SERVER LIST\n'
  pattern = re.compile(r'^\s*#BEGIN SERVER LIST\n  #DO NOT EDIT!\n(.*\n)*  #END SERVER LIST\n', re.MULTILINE)
  config = re.sub(pattern, server_list, config)
  return config
  

def load_varnish_config():
  subprocess.call(['/sbin/varnishreload', '-m', '0'])

def load_lsyncd_config():
  subprocess.call(['/bin/systemctl', 'restart', 'lsyncd'])
  
def load_haproxy_config():
  subprocess.call(['/bin/systemctl', 'reload', 'haproxy'])

sys.exit(main())