Script for creating users in Consul in the format expected by nsscache.

create_consul_user.py

#!/usr/bin/env python

import argparse
import base64
import requests
import os


class ConsulUser(object):

  def __init__(self, server='localhost:8500', token=''):
    """initialize the consul object

    Args:
      server: The consul host:port pair to query.
      token: Authentication token.
    """

    self.server = server
    if not token:
      try:
        with open(os.path.join(os.environ['HOME'], '.consul'), 'r') as fd:
          token = fd.read().strip()
      except IOError:
        pass
    self.token = token
    self.admin_group = 'admins'

  def next_uid(self):
    """Claims the next highest available uid.

    Returns:
      The new UID.
    """

    url = 'http://{}/v1/kv/org/metadata/highestuid?token={}'.format(self.server, self.token)

    while True:
      resp = requests.get(url)
      if not resp.ok:
        raise Exception('{} returned {}: {}'.format(url, resp.status_code, resp.text))
      current = resp.json()[0]
      index = current['ModifyIndex']
      new_uid = int(base64.b64decode(current['Value'])) + 1

      resp = requests.put('{url}&cas={index}'.format(url=url, index=index),
                          data=str(new_uid))
      if not resp.ok:
        raise Exception('{}&cas={} returned {}: {}'.format(url, index, resp.status_code, resp.text))
      elif resp.text == 'true':
        return new_uid

  def create_key(self, key, value):
    """Creates and sets the inital value for a key. An
    exception will be thrown if the key already exists.

    Args:
      key: The key to create.
      value: The inital value of the key.
    """

    url = 'http://{}/v1/kv/{}?token={}&cas=0'.format(self.server, key, self.token)

    resp = requests.put(url, data=value)
    if not resp.ok:
      raise Exception('Creation for {} returned {}: {}'.format(key, resp.status_code, resp.text))
    elif resp.text != 'true':
      raise Exception('Unable to create key {} at {}'.format(key, url))

    return True

  def overwrite_key(self, key, value):
    """Sets the value of a key, ignoring any existing values.

    Args:
      key: The key to create.
      value: The inital value of the key.
    """

    url = 'http://{}/v1/kv/{}?token={}'.format(self.server, key, self.token)
    resp = requests.put(url, data=value)
    if not resp.ok:
      raise Exception('Creation for {} returned {}: {}'.format(key, resp.status_code, resp.text))
    elif resp.text != 'true':
      raise Exception('Unable to set key {} at {}'.format(key, url))

    return True

  def create_user(self, username, uid=None, gid=None, ssh_key=None, allowed_hosts=None, admin=False):
    """Creates a new user account in Consul.

    Args:
      username: Name of the user.
      uid: Global user ID. If not provided, the next available one will be assigned.
      gid: Global group ID. If not provided, the UID will be used.
      ssh_key: User's publish SSH key.
      allowed_hosts: Iterable of regexes that the user is allowed access to.
      admin: Boolean indicating whether the user has local admin access.
    """

    if not uid:
      uid = self.next_uid()

    if not gid:
      gid = uid
      self.create_group(username, gid)

    self.create_key('org/humans/{}/uid'.format(username), str(uid))
    self.create_key('org/humans/{}/gid'.format(username), str(gid))

    if ssh_key:
      self.add_ssh_key(username, ssh_key)

    if allowed_hosts:
      hosts = '\n'.join(allowed_hosts)
      self.create_key('org/humans/{}/allowed_hosts'.format(username), hosts)

    if admin:
      self.add_group_member(username, self.admin_group)

  def create_group(self, name, gid):
    key = 'org/factions/{}/gid'.format(name)
    self.create_key(key, str(gid))

  def add_ssh_key(self, username, ssh_key):
    url = 'http://{}/v1/kv/org/humans/{}/authorized_keys?token={}'.format(self.server, username, self.token)
    resp = requests.put('{}&cas=0'.format(url), data=ssh_key)

    while resp.text != 'true':
      resp = requests.get(url)
      if not resp.ok:
        raise Exception('Unable to retrieve existing ssh keys for {} - {}: {}'.format(
                        username, resp.status_code, resp.text))
      current = resp.json()[0]
      index = current['ModifyIndex']
      existing_keys = base64.b64decode(current['Value'])
      new_keys = '{}\n{}'.format(existing_keys, ssh_key)

      resp = requests.put('{url}&cas={index}'.format(url=url, index=index),
                          data=new_keys)
      if not resp.ok:
        raise Exception('PUT for {} ssh key with index {} returned {}: {}'.format(
                        username, index, resp.status_code, resp.text))

      return True

  def add_group_member(self, username, group):
    """Adds a user to a given group.

    Args:
      username: The user to add.
      group: The name of the group.
    """

    url = 'http://{}/v1/kv/org/factions/{}/members?token={}'.format(self.server, group, self.token)
    resp = requests.put('{}&cas=0'.format(url), data=username)

    while resp.text != 'true':
      resp = requests.get(url)
      if not resp.ok:
        raise Exception('Unable to retrieve existing group members for {} - {}: {}'.format(
                        group, resp.status_code, resp.text))
      current = resp.json()[0]
      index = current['ModifyIndex']
      members = base64.b64decode(current['Value']).split('\n')
      if username in members:
        return True

      members.append(username)
      resp = requests.put('{url}&cas={index}'.format(url=url, index=index),
                          data='\n'.join(members))
      if not resp.ok:
        raise Exception('PUT for {}/members with index {} returned {}: {}'.format(
                        group, index, resp.status_code, resp.text))

      return True

  def add_allowed_host(self, username, host_regex):
    url = 'http://{}/v1/kv/org/humans/{}/allowed_hosts?token={}'.format(self.server, username, self.token)
    resp = requests.put('{}&cas=0'.format(url), data=host_regex)

    while resp.text != 'true':
      resp = requests.get(url)
      if not resp.ok:
        raise Exception('Unable to retrieve existing allowed_hosts for {} - {}: {}'.format(
                        username, resp.status_code, resp.text))
      current = resp.json()[0]
      index = current['ModifyIndex']
      existing_hosts = base64.b64decode(current['Value'])
      new_hosts = '{}\n{}'.format(existing_hosts, host_regex)

      resp = requests.put('{url}&cas={index}'.format(url=url, index=index),
                          data=new_hosts)
      if not resp.ok:
        raise Exception('PUT for {}/allowed_hosts with index {} returned {}: {}'.format(
                        username, index, resp.status_code, resp.text))

      return True


if __name__ == '__main__':
  parser = argparse.ArgumentParser(description='Create a new user in Consul.')
  parser.add_argument('--token', help='ACL token. If not provided, this will be read from ~/.consul', default='')
  parser.add_argument('--server', default='localhost:8500',
                      help='Consul server.')
  parser.add_argument('--uid')
  parser.add_argument('--gid')
  parser.add_argument('--admin', action='store_true', help='Add user to the sudo group')
  parser.add_argument('--ssh-key', help='User\'s authorized ssh key')
  parser.add_argument('--allowed-hosts', default='.*-dev',
                      help='Comma separated list of host regexes the user has access to.')
  parser.add_argument('username', help='Name of the new user.')
  args = parser.parse_args()

  allowed_hosts = args.allowed_hosts.split(',')

  consul = ConsulUser(args.server, args.token)

  consul.create_user(args.username, uid=args.uid, gid=args.gid,
                     ssh_key=args.ssh_key, allowed_hosts=allowed_hosts,
                     admin=args.admin)