bootstrap.py

#!/usr/bin/env python

"""
This module will bootstrap a machine using chef.  The purpose of this
script is actually to work with AWS Auto Scaling Groups.  The user data
for the Launch Configuration is set to download this script and then 
run it.  This is also stores the results in a private gist and sends
a message to logstash with the results.
"""

import argparse
import datetime
import json
import os
import os.path
import re
import socket
import sys
import urllib2

import boto
from boto.s3.key import Key

import requests

import sh

INSTANCE_ID_URL = 'http://169.254.169.254/latest/meta-data/instance-id'
GIST_POST_URL = 'https://api.github.com/gists'

CLIENT_RB = """
log_level        :info
log_location     STDOUT
verbose_logging 

chef_server_url  "{0}"
validation_client_name "chef-validator"

file_backup_path   "/var/lib/chef"
file_cache_path    "/var/cache/chef"

pid_file           "/var/run/chef/client.pid"

Ohai::Config[:plugin_path] << "/etc/chef/ohai_plugins"

node_name          "{1}"
"""


def configure_chef(server_details, chef_server, bucket):
    """
    Get all of the chef configuration files from S3 and store them
    in the /etc/chef/ directory.

    :type server_details: dict
    :param server_details: Collection of properties for server

    :type bucket: boto.s3.Bucket
    :param bucket: The S3 bucket containing config.
    """
    for filename in ['validation.pem', 'encrypted_data_bag_secret']:
        target = os.path.join('/etc/chef', filename)
        key = Key(bucket)
        key.key = filename
        key.get_contents_to_filename(target)
        os.chmod(target, 0600)
    with open('/etc/chef/client.rb', 'w') as handle:
        handle.write(CLIENT_RB.format(chef_server, server_details['name']))

def tag_server(server_details, ec2):
    """
    Apply ec2 tags to this instance.

    :type server_details: dict
    :param server_details: Collection of properties for server

    :type ec2: boto.EC2Connection
    :param ec2: A connection to EC2
    """
    tags = {
            "Name": server_details["name"],
            "Environment": server_details["environment"],
            "ServerClass": server_details["class"]
            }
    role_regex = re.compile(r"role\[(.*)\]")
    for item in server_details['run_list']:
        match = role_regex.search(item)
        if match is not None:
            tags[match.group(1)] = "True"
    ec2.create_tags([server_details['instance_id']], tags)


def generate_server_details(server_class, environment, bucket):
    """
    Generate a server name based on the server class and environment

    :type redis_host: string
    :param redis_host: The redis server to contact to increment node_id

    :type server_class: String
    :param server_class: What "type" of server is this

    :type environment: dict
    :param environment: Chef environment for this node

    :type s3: boto.S3Connection
    :param s3: A connection to S3
    """
    response = urllib2.urlopen(INSTANCE_ID_URL)
    instance_id = response.read()
    server_name = "i_{0}_{1}_{2}".format(environment, server_class, instance_id)
    key = Key(bucket)
    key.key = 'first-boot-scripts/{0}/{1}.json'.format(environment, server_class)
    key.get_contents_to_filename('/etc/chef/first-boot.json')

    with open('/etc/chef/first-boot.json', 'r') as handle:
        run_list = json.loads(handle.read())['run_list']

    return {
            'name': server_name,
            'instance_id': instance_id,
            'environment': environment,
            'class': server_class,
            'run_list': run_list
            }

def run_chef(server_details, chef_server, bucket, github_token, logstash_host, logstash_port):
    """
    First, install chef on the node.  Then, run the chef-client using the
    configuration which was created earlier
    """
    response = urllib2.urlopen('http://www.opscode.com/chef/install.sh')
    with open('/tmp/chef-install.sh', 'w') as handle:
        handle.write(response.read())
    chef_install_log = sh.bash('/tmp/chef-install.sh')
    configure_chef(server_details, chef_server, bucket)
    chef_client = sh.Command('/opt/chef/bin/chef-client')
    chef_run_log = chef_client(config='/etc/chef/client.rb', environment=server_details['environment'], j='/etc/chef/first-boot.json')
    gist_payload = {
        "description": "Boostrap for {0}".format(server_details['name']),
        "public": False,
        "files": {
            "chef-install.log": {
                "content": str(chef_install_log)
            },
            "chef-run.log": {
                "content": str(chef_run_log)
            }
        }
    }
    r = requests.post(GIST_POST_URL, params={"access_token": github_token}, data=json.dumps(gist_payload))
    gist_url = json.loads(r.text)['html_url']
    log_message = {
        '@timestamp': datetime.datetime.utcnow().isoformat('T') + 'Z',
        '@tags': ['bootstrap', 'chef', 'ASG'],
        '@type': 'bootstrap',
        '@source': server_details['name'],
        '@fields': {'gist_url': gist_url},
        '@message': "Instance bootstrapped automaticatlly.  Report: {0}".format(gist_url)
    }
    sys.stdout.write(json.dumps(log_message, indent=4) + '\n')
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    except socket.error, msg:
        sys.stderr.write("[ERROR] Failed to write to logstash: {0}".format(msg))
        sys.exit(1)
    try:
        sock.connect((logstash_host, int(logstash_port)))
    except socket.error, msg:
        sys.stderr.write("[ERROR] Failed to write to logstash: {0}".format(msg))
        sys.exit(2)

    sock.send(json.dumps(log_message) + '\n')
    sock.close()
      

def main():
    """docstring for main"""
    parser = argparse.ArgumentParser("Bootstrap an EC2 box for chef")
    parser.add_argument('-k', '--key', help='AWS Access Key ID')
    parser.add_argument('-s', '--secret', help='AWS Secret Access Key')
    parser.add_argument('-b', '--bucket',
        help='S3 Bucket which contains config files')
    parser.add_argument('-t', '--server_class', help='Server class (type)')
    parser.add_argument('-e', '--environment', help='The chef environment')
    parser.add_argument('-c', '--chefserver', help='URL of the chef server')
    parser.add_argument('-g', '--githubtoken', help='Github OAuth Token to use to post the gist')
    parser.add_argument('-l', '--logstashhost', help='Host of the logstash server')
    parser.add_argument('-L', '--logstashport', help='Port of the logstash server')
    args = parser.parse_args()
    sh.mkdir('/etc/chef', parents=True)
    s3 = boto.connect_s3(aws_access_key_id=args.key, aws_secret_access_key=args.secret)
    bucket = s3.get_bucket(args.bucket)
    ec2 = boto.connect_ec2(aws_access_key_id=args.key, aws_secret_access_key=args.secret)
    server_details = generate_server_details(args.server_class, args.environment, bucket)
    tag_server(server_details, ec2)
    run_chef(server_details, args.chefserver, bucket, args.githubtoken, args.logstashhost, args.logstashport)
    s3.close()
    ec2.close()

if __name__ == '__main__':
    main()

user_data.sh

# This is a sample of what you would put into your user-data

#/bin/bash

export AWS_ACCESS_KEY_ID=********************
export AWS_SECRET_ACCESS_KEY_ID=***************************
export LOGSTASH_HOST=*****************
export LOGSTASH_PORT=*****************
export GITHUB_TOKEN=***********************

curl -o /tmp/virtualenv.py https://raw.github.com/pypa/virtualenv/master/virtualenv.py 
python /tmp/virtualenv.py /tmp/install 
/tmp/install/bin/pip install boto
/tmp/install/bin/pip install sh
/tmp/install/bin/pip install requests

curl -o /tmp/bootstrap.tar.gz https://gist.github.com/gists/4218533/download
mkdir /tmp/bootstrap
cd /bootstrap
tar -xzvf bootstrap.tar.gz
mv */bootstrap.py /tmp

/tmp/install/bin/python /tmp/bootstrap.py -b your-s3-bucket -t server-type -e stg -c http://your.chef.url -k $AWS_ACCESS_KEY_ID -s $AWS_SECRET_ACCESS_KEY_ID -l $LOGSTASH_HOST -L $LOGSTASH_PORT -g $GITHUB_TOKEN > /tmp/bootstrap.log 2> /tmp/bootstrap.error

rm -rf /tmp/install /tmp/bootstrap*