#!/usr/bin/python
# Copyright 2013 Hewlett-Packard Development Company, L.P.
# All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
#

import base64
import json
import logging
import os
import sys
import subprocess

logging.basicConfig(level='INFO',
   format='[%(asctime)s] (%(name)s) [%(levelname)s] %(message)s')
LOG = logging.getLogger(os.path.basename(sys.argv[0]))

USERS_FILE='/etc/rabbitmq/heat_users.json'
HANDLE_FILE = '/etc/rabbitmq/heat_password_handle'

# USERS_FILE can exist without any content, in which case it will just be an
# empty line.
if not os.path.exists(USERS_FILE) or os.path.getsize(USERS_FILE) < 2:
    sys.exit(0)

PASSWORD_HANDLE = None
if os.path.exists(HANDLE_FILE):
    with open(HANDLE_FILE) as hf:
        PASSWORD_HANDLE = hf.read().rstrip()

def get_existing_users():
    list_users = subprocess.check_output(['rabbitmqctl', 'list_users'],
                                         stderr=subprocess.STDOUT)
    users = {}
    for line in list_users.splitlines():
        if line.startswith('Listing users'):
            continue
        if line.startswith('...'):
            continue
        try:
            user, tags = line.split(None, 1)
            tags = tags.split('[')[1].split(']')[0].split(',')
        except ValueError as e:
            LOG.warn('{%s} did not have expected format.' % (line))
            LOG.exception(e)
            continue
        users[user] = {'username': user, 'tags': tags}
    return users

with open(USERS_FILE) as uf:
    user_map = json.loads(uf.read())

have_map = get_existing_users()

want = set(user_map.keys())
have = set(have_map.keys())
need = want - have

LOG.info("want = %s" % want)
LOG.info("have = %s" % have)
LOG.info("need = %s" % need)

for need_user in need:
    detail = user_map[need_user]
    username = detail['username']
    tags = detail['tags']
    if username not in have:
        if 'password' in detail:
            password = detail['password']
        else:
            password = base64.b64encode(os.urandom(40))
        subprocess.check_call(['rabbitmqctl', 'add_user', username, password],
                              stdout=sys.stderr)
        if PASSWORD_HANDLE:
            subprocess.check_call(['cfn-signal',
                                   '-s', 'true',
                                   '-i', username,
                                   '--data', password,
                                   PASSWORD_HANDLE])
        else:
            print '%s:%s' % (username, password)
    args = ['rabbitmqctl', 'set_user_tags', username]
    args.extend(tags)
    subprocess.check_call(args, stdout=sys.stderr)

have = set(get_existing_users().keys())
if want - have:
    LOG.error('Desired users missing: want=%s have=%s', (want, have))
    sys.exit(1)
