Simple script for generating Kerberos principals and keytabs from the Ambari host-keytab CSV file

setup-krb-users-and-keytabs.py

#!/usr/bin/env python
# File: setup-krb-users-and-keytabs.py
#
# This is a simple script that will handle the creation of principals on
# the KDC, the creation of keytabs with appropriate permissions, and
# the copying of keytabs to the appropriate locations. It's input is the
# host-principal-keytab-list file that Ambari generates.
#
# This script MUST be run as root on the KDC node.
#
#
# 'container' is a nested datastructure of the form:
#  {hostname : {keytabdir: {keytab: (perms) } } }
#
#

import sys
import os
import pwd
import grp
import socket
import tempfile
import shutil
from subprocess import call

runninghost = socket.gethostname()
tempktspath = tempfile.mkdtemp()


def usage():
    return "Usage: setup-krb-users-and-keytabs <host-principal-file>"


def get_filename(args):
    """Gets the file name from the list of arguments. For now, there's
            only one argument so we return that. """
    if len(args) < 1:
        print(usage())
        raise Exception("An error occurred")
    return os.path.abspath(args[0])


def make_keytab_dirs(container):
    """Create the directories for keytab files"""
    for host, val in container.iteritems():
        for dirs in val.keys():
            if host == runninghost:
                if not os.path.exists(dirs):
                    print('Creating {0}'.format(dirs))
                    os.makedirs(dirs, 0755)
            else:
                # create dir on remote host
                cmdstr = "ssh {0} 'mkdir {1} -m 0755'".format(host, dirs)
                print("Running command on remote host: {0}. Command: {1}".
                      format(host, cmdstr))
                try:
                    call(["ssh", host, "mkdir", dirs, "-m 0755"])
                except Exception, e:
                    print("Unable to run command on {}".format(host))
                    print(e)



def create_princs(principals):
    """Create principals in the KDC database using kadmin.local"""
    for p in principals:
        print("Running command: kadmin.local -q 'addprinc -randkey {0}'".format(p))
        call(['kadmin.local', '-q', 'addprinc -randkey ' + p])


def printmap(nestedmap):
    """Print out the nested map of form:
        { key1 : { key2 : object} }"""
    for x in nestedmap:
        print (x)
        for y in nestedmap[x]:
                print ("{0}: {1}".format(y, nestedmap[x][y]))
        print("\n")


def write_princs_to_keytabs(container):
    """Write the principals to the appropriate keytab files in temporary
        location"""
    # os.mkdir(tempktspath)
    for keytab, princs in container.iteritems():
        cmd = "ktadd -k {0}/{1} {2}".format(tempktspath, keytab, " ".join(princs))
        print(["kadmin.local", "-q", cmd])
        call(["kadmin.local", "-q", cmd])


def get_keytab_perms_map(lines):
    """Set permissions on the keytab files"""
    pathfileperms = {}
    for line in lines:
        keytabdir = line[4]
        keytab = line[3]
        perms = (line[5], line[6], line[7])
        if keytabdir not in pathfileperms:
            pathfileperms[keytabdir] = {}
        if keytab not in pathfileperms[keytabdir]:
            pathfileperms[keytabdir][keytab] = None
        pathfileperms[keytabdir][keytab] = perms
    return pathfileperms


def save_perms(path, username, grpname, permstr):
    try:
        uid = pwd.getpwnam(username).pw_uid
        gid = grp.getgrnam(grpname).gr_gid
    except Exception, e:
        print(e)
        return
    print(r'os.chown({0}, {1}, {2})'.format(path, uid, gid))
    os.chown(path, uid, gid)
    print(r'os.chmod({0}, {1})'.format(path, int(permstr, 8)))
    os.chmod("{0}".format(path), int(permstr, 8))


def copy_keytabs_set_perms(container):
    for host, val in container.iteritems():
        for keytabdir, val2 in val.iteritems():
            for keytab, perms in val2.iteritems():
                src = tempktspath + "/" + keytab
                if host == runninghost:
                    # copy file
                    print("Copying {0} to {1}".format(src, keytabdir))
                    shutil.copy(src, keytabdir)
                    print("Saving permissions")
                    save_perms(keytabdir + "/" + keytab, perms[0], perms[1],
                               perms[2])
                else:
                    # copy file
                    dsthost = host + ":" + keytabdir
                    print("Executing scp {0} {1}".format(src, dsthost))
                    ret = call(["scp", src, dsthost])
                    # set perms
                    if ret == 0:
                        ug = perms[0] + ":" + perms[1]
                        dst = keytabdir + "/" + keytab
                        perm = "0" + perms[2]
                        print("Executing ssh {0} chown {1} {2}".format(host, ug, dst))
                        call(["ssh", host, "chown", ug, dst])
                        print("Executing ssh {0} chmod {1} {2}".format(host, perm, dst))
                        call(["ssh", host, "chmod", perm, dst])
                    else:
                        print("scp failed with return code: {0}".format(ret))


def cleanup():
    shutil.rmtree(tempktspath)


def process(principals, keytabprincsmap, container):
    make_keytab_dirs(container)
    create_princs(principals)
    write_princs_to_keytabs(keytabprincsmap)
    copy_keytabs_set_perms(container)
    cleanup()


def main(args):
    """Main function which controls everything"""
    fname = get_filename(args)
    # container is: {hostname: {keytabdir: {keytab: (perms) } } }
    container = {}
    keytabprincsmap = {}
    principals = set()
    with open(fname, 'r') as f:
        # each line has the form:
        # <hostname>, <username>, <principal>, <keytab>, <keytabdir>,
        # <fowner>, <fgroup>, <fperms>
        for line in f:
            line = [r.rstrip() for r in line.split(",")]
            if len(line) == 1: continue
            principals.add(line[2])
            hostname, keytab, keytabdir = line[0], line[3], line[4]
            if hostname not in container:
                container[hostname] = {}
            if keytabdir not in container[hostname]:
                container[hostname][keytabdir] = {}
            if keytab not in container[hostname][keytabdir]:
                container[hostname][keytabdir][keytab] = (line[5], line[6], line[7])
            if keytab not in keytabprincsmap:
                keytabprincsmap[keytab] = set()
            keytabprincsmap[keytab].add(line[2])

    # printmap(container)
    # print(keytabprincsmap)
    process(principals, keytabprincsmap, container)


if __name__ == '__main__':
    # Before we go too far, check if the user is root
    if not os.geteuid() == 0:
        sys.exit("\nThis script must be run as root!\n")
    print("Running script on " + runninghost)
    main(sys.argv[1:])