gen_certs.py

#!/usr/bin/env python2

import argparse
import subprocess
import logging
import distutils.util
from shutil import copyfile

fmt_string = 'certs %(levelname)s:%(asctime)s line:%(lineno)d] %(message)s'
logging.basicConfig(format=fmt_string)
logger = logging.getLogger("certs")
logger.setLevel(logging.INFO)

DEFAULT_CERTIFICATE_AUTHORITY = "root_certificate_authority"


def run_subprocess(cmd):
    logger.info("Running command: exec bash -c '%s'" % cmd)
    proc = subprocess.Popen(
        "exec bash -c '%s'" % cmd,
        stdout=sys.stdout,
        stderr=sys.stderr,
        shell=True)
    return_code = 0
    try:
        return_code = proc.wait()
        sys.stdout.flush()
        sys.stderr.flush()
    except Exception as e:
        proc.kill()
        raise
    if return_code != 0:
        raise subprocess.CalledProcessError(return_code, cmd)


def generate_options():
    parser = argparse.ArgumentParser(
        description='Generate certificates for kafka, mesos, zk & concord.')
    parser.add_argument(
        '--log',
        type=str,
        default='INFO',
        help='info,debug, type log levels. i.e: --log=debug')
    parser.add_argument(
        '--kafka_certs',
        type=distutils.util.strtobool,
        default=True,
        help='Kafka certs')
    parser.add_argument(
        '--mesos_certs',
        type=distutils.util.strtobool,
        default=True,
        help='Mesos certs')
    parser.add_argument(
        '--ca_cert',
        type=distutils.util.strtobool,
        default=True,
        help='CA cert')
    parser.add_argument(
        '--zookeeper_certs',
        type=distutils.util.strtobool,
        default=True,
        help='zookeeper certs')
    parser.add_argument(
        '--password',
        type=str,
        default="abcdefgh",
        help='Password for kafka cert')

    return parser


def gen_key(name):
    logger.info("Generating key: %s", name)
    cmd = "openssl genrsa -out %s.key 2048" % name
    run_subprocess(cmd)


def gen_csr(name):
    logger.info("Generating csr for: %s", name)
    cmd = "openssl req -new -key %s.key -out %s.csr" % (name, name)
    run_subprocess(cmd)


def gen_self_sign_cert(ca, name):
    logger.info("Self signing cert for ca: %s, and csr: %s", ca, name)
    cmd = "openssl x509 -req -in %s.csr -CA %s.pem -CAKey" \
        " %s.key -CAcreateserial -out %s.crt -sha256" % (name, ca, ca, name)
    run_subprocess(cmd)


def gen_pem(name):
    logger.info("Generating pem")
    cmd ="openssl req -x509 -new -nodes -key %s.key"\
        "sha256 -out %s.pem" % (name,name)
    run_subprocess(cmd)
    logger.info("For JVM %s.pem == %s.chain_cert" % (name, name))
    copyfile("%s.pem" % name, "%s.chain_cert" % name)


def gen_ca_cert(name):
    if os.path.exists("%s.key" % name): return
    logger.info('Generating ca:%s', name)
    gen_key(name)
    gen_pem(name)


def gen_certs(name, ca):
    gen_ca_cert(ca)
    gen_key(name)
    gen_csr(name)
    gen_self_sign_cert(ca, name)


def gen_mesos_certs():
    gen_certs("mesos_master", DEFAULT_CERTIFICATE_AUTHORITY)
    gen_certs("mesos_agent", DEFAULT_CERTIFICATE_AUTHORITY)
    gen_certs("concord_scheduler", DEFAULT_CERTIFICATE_AUTHORITY)
    gen_certs("concord_executor", DEFAULT_CERTIFICATE_AUTHORITY)


def gen_keystore(name, passwd):
    logger.info("Generating keystore for %s with password %s", name, passwd)
    cmd = "keytool -storepass %s -keypass %s -keystore %s.keystore.jks" \
        "-alias localhost -genkey" %(passwd,passwd,name)
    run_subprocess(cmd)


def gen_truststore(name, passwd, ca):
    logger.info("Generating truststore for %s with passwd %s and ca: %s", name,
                passwd, ca)
    cmd = "keytool -storepass %s -keystore %s.truststore.jks "\
        "-alias %s -import -file %s.pem" %(passwd,name,ca,ca)
    run_subprocess(cmd)


def gen_keystore_cert(name, passwd):
    logger.info("Generating keystore cert for %s with passwd: %s", name,
                passwd)
    cmd = "keytool -storepass %s -keypass %s -keystore %s.keystore.jks -alias" \
    " localhost -certreq -file %s.cert" %(passwd, passwd, name, name)
    run_subprocess(cmd)


def gen_self_signed_keystore_cert(name, passwd, ca):
    logger.info("Generating self-signed keystore cert for: %s with passwd %s" \
                " and ca: %s", name, passwd, ca)
    cmd = "openssl x509 -req -CA %s.pem -CAkey %s.key -in %s.cert" \
        " -out %s.cert.signed -CAcreateserial -passin pass:%s" %(ca, ca, name, name, passwd)
    run_subprocess(cmd)


def gen_import_certificate_authority_to_truststore(name, passwd, ca):
    logger.info("Importing ca: %s into the %s.keystore.jks with passwd %s", ca,
                name, passwd)
    cmd = "keytool -storepass %s -keypass %s -keystore %s.keystore.jks" \
        " -alias %s -import -file %s.pem" % (passwd, passwd, name,ca, ca)
    run_subprocess(cmd)


def gen_import_self_signed_cert_to_truststore(name, passwd):
    logger.info("Importing self signed cert: %s.cert.signed" \
                "to %s.keystore.jks with passwd: %s", name, name, passwd)
    cmd = "keytool -storepass %s -keypass %s -keystore %s.keystore.jks" \
        " -alias localhost -import -file %s.cet.signed" %(passwd, passwd, name,
                                                          name)
    run_subprocess(cmd)


def gen_certs_with_jvm(name, password, ca):
    gen_keystore(name, password)
    gen_truststore(name, password, ca)
    gen_keystore_cert(name, password, ca)
    gen_self_signed_keystore_cert(name, password, ca)
    gen_import_certificate_authority_to_truststore(name, passwd, ca)
    gen_import_self_signed_cert_to_truststore(name, password)
    gen_certs("%s.client" % name)


def gen_kafka_certs(password):
    gen_certs_with_jvm("kafka", password, DEFAULT_CERTIFICATE_AUTHORITY)


def gen_zookeeper_certs(password):
    gen_certs_with_jvm("zookeeper", password, DEFAULT_CERTIFICATE_AUTHORITY)


def main():
    parser = generate_options()
    options, program_options = parser.parse_known_args()
    logger.info(options)
    logger.info("Common Name points to the ip address i.e.: 10.0.0.1")

    try:
        if options.ca_cert:
            gen_ca_cert()
        if options.kafka_certs:
            gen_kafka_certs(options.password)
        if options.mesos_certs:
            gen_mesos_certs()
        if options.zookeeper_certs:
            gen_zookeeper_certs(options.password)
    except Exception as e:
        log.error(str(e))


if __name__ == '__main__':
    main()