write_secrets

#!/usr/bin/env python
# The contents of this file are subject to the Common Public Attribution
# License Version 1.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://code.reddit.com/LICENSE. The License is based on the Mozilla Public
# License Version 1.1, but Sections 14 and 15 have been added to cover use of
# software over a computer network and provide for limited attribution for the
# Original Developer. In addition, Exhibit A has been modified to be consistent
# with Exhibit B.
#
# Software distributed under the License is distributed on an "AS IS" basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License for
# the specific language governing rights and limitations under the License.
#
# The Original Code is reddit.
#
# The Original Developer is the Initial Developer.  The Initial Developer of
# the Original Code is reddit Inc.
#
# All portions of the code written by reddit are Copyright (c) 2006-2015 reddit
# Inc. All Rights Reserved.
###############################################################################

import base64
import ConfigParser
import fileinput
import getpass
import json
import os
import sys

import kazoo

from kazoo.security import make_digest_acl

from r2.lib.utils import parse_ini_file
from r2.lib.zookeeper import connect_to_zookeeper
from r2.lib.app_globals import SECRETS_NODE, extract_secrets


USERNAME = "live-config"


def _encode_secrets(secrets):
    return json.dumps({key: base64.b64encode(secret)
                       for key, secret in secrets.iteritems()})


def write_secrets_to_zookeeper(reddit_config, username, password, secrets):
    # read the zk configuration from the app's config
    zk_hostlist = reddit_config.get("DEFAULT", "zookeeper_connection_string")
    app_username = reddit_config.get("DEFAULT", "zookeeper_username")
    app_password = reddit_config.get("DEFAULT", "zookeeper_password")

    # connect to zk!
    client = connect_to_zookeeper(zk_hostlist, (username, password))

    # we're going to assume that any parent parts of the node path were
    # already created by write_live_config.
    json_data = _encode_secrets(secrets)
    try:
        client.create(SECRETS_NODE, json_data, acl=[
            make_digest_acl(username, password, read=True, write=True),
            make_digest_acl(app_username, app_password, read=True),
        ])
    except kazoo.exceptions.NodeExistsException:
        client.set(SECRETS_NODE, json_data)


def main():
    progname = os.path.basename(sys.argv[0])

    input = fileinput.input()
    try:
        config = parse_ini_file(input)
    except (IOError, ConfigParser.Error), e:
        print >> sys.stderr, "%s: %s" % (progname, e)
        return 1

    secrets = extract_secrets(config)
    password = getpass.getpass("ZooKeeper Password: ")

    try:
        write_secrets_to_zookeeper(config, USERNAME, password, secrets)
    except kazoo.exceptions.NoAuthException:
        print >> sys.stderr, "%s: incorrect password" % progname
        return 1
    except Exception as e:
        print >> sys.stderr, "%s: %s" % (progname, e)
        return 1


if __name__ == "__main__":
    sys.exit(main())