#!/usr/bin/env python

# Copyright (c) 2011. All Right Reserved, http://chart.io/
#
# THIS CODE AND INFORMATION ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY
# KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
# PARTICULAR PURPOSE.

import ConfigParser
import getpass
import optparse
import os
import pprint
import random
import signal
import string
import subprocess
import sys
import time
import urllib
import urllib2

try:
    import json
except ImportError:
    try:
        import simplejson as json
    except ImportError:
        print 'Please install simplejson module'
        sys.exit(1)

KEY = 'FEUUB1JB4MHNHZ474R14VW3K62XGNP466GPRMD1N3WF5ER047DTOXUO190KXR2VFBO31XTDOODU2H7XNDRL6EA8D5F5HUC52LMHW'
BASE_URL = os.environ.get('CHARTIO_BASE_URL', 'https://chart.io')
VERSION = 1
# Global configuration data
CONFIG_DATA = None
# Default install location
PREFIX_DEFAULT = '~/.chartio.d'
# Distributed version
VERSION = '1.1.7'


class ConfigData(object):
    '''Yuck

    Double Yuck. Needs to be kept in sync with chartio_connect to limit
    PYTHONPATH dependencies.

    '''
    def __init__(self, prefix=None):
        # Directories
        self.PREFIX = prefix or os.path.expanduser(PREFIX_DEFAULT)
        self.LOG_DIRECTORY = os.path.join(self.PREFIX, 'logs')
        self.RUN_DIRECTORY = os.path.join(self.PREFIX, 'run')
        self.SSH_DIRECTORY = os.path.join(self.PREFIX, 'sshkey')
        # Files
        self.CONFIG_FILE = os.path.join(self.PREFIX, 'chartio.cfg')
        self.SSH_KEY = os.path.join(self.SSH_DIRECTORY, 'id_rsa')
        # Config file sections
        self.SSHTUNNEL_SECTION = 'SSHTunnel'
        # Ensure the directories exist. Exit if they do not (or cannot be created).
        self.directory_create(self.PREFIX, 0755)
        self.directory_create(self.LOG_DIRECTORY, 0755)
        self.directory_create(self.RUN_DIRECTORY, 0755)
        self.directory_create(self.SSH_DIRECTORY, 0700)

    def directory_create(self, path, mode):
        '''Create a directory if it does not exist.

        Exit if it exists and is not a directory (or symlink) or unable to create it.

        Arguments
            path -- directory path
            mode -- mode to set on directory

        '''
        if os.path.exists(path):
            if os.path.isdir(path):
                try:
                    os.chmod(path, mode)
                except Exception, exc:
                    TermColor.print_error('Failed to change mode of %r. Exiting.' % (path))
                    TermColor.print_error(str(exc))
                    sys.exit(1)
            else:
                TermColor.print_error('The path %r is not a directory. Exiting.' % (path))
                sys.exit(1)
        else:
            try:
                os.makedirs(path, mode)
            except Exception, exc:
                TermColor.print_error('Failed to create %r. Exiting.' % (path))
                TermColor.print_error(str(exc))
                sys.exit(1)


def config_data_set(prefix):
    '''Update the global config data.

    Exits when unable to create configuration directories.

    Arguments
        prefix -- The install prefix

    '''
    global CONFIG_DATA
    CONFIG_DATA = ConfigData(prefix)
    print 'Installing into %r' % (CONFIG_DATA.PREFIX)


class TermColor(object):
    '''Print colored text on a neutral background to the terminal'''
    CLRS = {
        'white': '\033[37m',
        'green': '\033[92m',
        'red': '\033[91m',
        'yellow': '\033[33m',
        'bg': '\033[40m\033m'
    }

    END = '\033[0m'

    @classmethod
    def print_clr(cls, color, txt, newline=True):
        sys.stdout.write(cls.CLRS.get(color, '')
                         + cls.CLRS['bg']
                         + txt
                         + cls.END)
        if newline:
            sys.stdout.write('\n')


    @classmethod
    def print_cmd(cls, txt, newline=True):
        cls.print_clr('yellow', txt, newline)

    @classmethod
    def print_header(cls, txt, newline=True):
        cls.print_clr('white', txt, newline)

    @classmethod
    def print_ok(cls, txt, newline=True):
        cls.print_clr('green', txt, newline)

    @classmethod
    def print_error(cls, txt, newline=True):
        cls.print_clr('red', 'Error: ' + txt, newline)

    @classmethod
    def print_delay(cls, txt, newline=True):
        cls.print_clr('red', '==> ', False)
        cls.print_ok(txt, newline)


def get_choice(question, choices, default=None):
    '''Prompt for a response from a selection of choices.

    Return
        The selected item from choices (value, not an index)

    Raises
        ValueError -- if default is not in choices

    Arguments
        question -- the prompt
        choices -- possible answers
        default -- optional default value

    Example
        choice = get_choice('What fruit do you want?', ['apples', 'oranges'], 'apples')
        print 'choice', choice

    '''
    if default is None:
        default_idx = None
    else:
        try:
            default_idx = choices.index(default)
        except ValueError:
            raise ValueError('Default choice %r is not in choice collection' % (default))

    enum_choices = list(enumerate(choices))
    prompt = ((default_idx is not None and ('[%d]: ' % (default_idx + 1)))
              or ': ')
    while True:
        TermColor.print_header(question)
        for idx, item in enum_choices:
            TermColor.print_ok('    %d.' % ((idx + 1)), newline=False)
            print ' %s' % (item)
        input_raw = raw_input(prompt)
        if is_integer(input_raw):
            choice_idx = int(input_raw) - 1
        elif not input_raw.strip() and default_idx is not None:
            choice_idx = default_idx
        else:
            choice_idx = None

        if choice_idx is None:
            TermColor.print_error('invalid choice value: %s' % (input_raw))
        elif choice_idx < 0 or (len(choices) <= choice_idx):
            TermColor.print_error('choice out of range: %s' % (choice_idx))
        else:
            # !!! Exit loop
            break
    return choices[choice_idx]


def get_value(name, default=None, validate=None, validate_explanation=None,
              is_password=False):
    '''Prompt for and read a value from the terminal.

    Return
        string -- the read value

    Arguments
        name -- value name
        default -- [optional] default value
        validate -- [optional] callable to validate the input. Defaults to
            ensuring something was entered.
        validate_explanation -- [optional] error message if validate() fails
        is_password -- [optional] if True, do not echo the response. Defaults to False.

    '''
    prompt_default = (default and ' [%s]' % (default)) or ''
    prompt = '%s%s: ' % (name, prompt_default)
    if is_password:
        input_fn = lambda: getpass.getpass('')
    else:
        input_fn = lambda: raw_input()
    validate_fn = validate or (lambda x: bool(x))
    error_msg = validate_explanation or ('Invalid Input.  Please try again.')
    while True:
        TermColor.print_header(prompt, newline=False)
        input_raw = input_fn().strip()
        input = input_raw or default or ''
        if validate_fn(input):
            # !!! Loop exit
            break
        else:
            TermColor.print_error(error_msg)
    return input


def is_integer(value):
    '''Determine whether value is convertible to an integer'''
    try:
        int(value)
    except (TypeError, ValueError), e:
        rc = False
    else:
        rc = True
    return rc


def name_generate(db_name, max_length):
    '''Generate a database user name which does not exceed a maximum length

    If the name would exceed the maximum length, the name is shortened
    and some random digits appended.

    Return
        string -- the generated name

    Raises
        RuntimeError -- if database name and maximum length values do not
            permit generation of a useful name.

    Arguments
        db_name -- database name
        max_length -- the inclusive size limit of the generated name

    '''
    PREFIX = 'chartio_'
    RANDOM_SUFFIX_LENGTH = 3
    MAX_REQUIRED = len(PREFIX) + RANDOM_SUFFIX_LENGTH
    if max_length < MAX_REQUIRED:
        raise RuntimeError('Unable to generate a name with fewer than %d characters (%d specified).'
                           % (MAX_REQUIRED, max_length))
    name_full = ('%s%s%s' % (PREFIX, db_name.strip(), RANDOM_SUFFIX_LENGTH * 'X'))[:max_length]
    chars = string.digits
    name = (name_full[:-RANDOM_SUFFIX_LENGTH]
            + ''.join([random.choice(chars) for i in range(RANDOM_SUFFIX_LENGTH)]))
    return name


class ResultTrue(object):
    def __init__(self, info):
        self.info = info

    def __nonzero__(self):
        return True

    def __repr__(self):
        return 'ResultTrue(%r)' % (self.info)


class ResultFalse(object):
    def __init__(self, info):
        self.info = info

    def __nonzero__(self):
        return False

    def __repr__(self):
        return 'ResultFalse(%r)' % (self.info)


def empty_is_ok(_value):
    '''Validator for accepting an empty string'''
    return True


def is_yes_no(value):
    '''Determine whether a string begins with y/n'''
    lowered = value.strip().lower()
    retval = lowered.startswith('y') or lowered.startswith('n')
    return retval


def bool_from_yes_no(value):
    '''Convert a y/n string to a bool'''
    lowered = value.strip().lower()
    retval = lowered.startswith('y')
    return retval


def random_password_generate(length=24):
    '''Generate a random password of a specified length

    '''
    valid_chars = string.letters + string.digits
    password = ''.join(map(lambda x: random.choice(valid_chars), range(length)))
    return password


def database_choices(chartio_api):
    '''XXXwalrus stub routine to avoid dealing with HTTP for now'''
    db_map = {'MySQL': {'name': 'MySQL',
                        'default_port': 3306,
                        'user_name_limit': 16},
              'PostgreSQL': {'name': 'PostgreSQL',
                             'default_port': 5432,
                             'user_name_limit': 63}}
    return db_map


def database_choices(chartio_api):
    db_settings = {'MySQL': {'default_port': 3306, 'user_name_limit': 16},
                   'PostgreSQL': {'default_port': 5432, 'user_name_limit': 63}}
    json_response = chartio_api.post('/connectionclient/databasetypes/')
    response = json.loads(json_response)
    databases = response.get('databasetypes', [])
    db_map = {}
    for db in databases:
        name = db['name']
        db_map[name] = {'name': name, 'id': db['id']}
        settings = db_settings.get(name)
        if settings:
            db_map[name].update(settings)
    return db_map


class Settings(object):
    '''Structure class'''
    def __init__(self):
        self.database_id = None
        self.database_name = None
        self.database_port = None
        self.readonly_user = None
        self.readonly_password = None
        # For dev
        self.db = None


class Poster(object):
    '''Class to POST information to Chartio.'''

    def __init__(self):
        self.opener = urllib2.build_opener(urllib2.HTTPCookieProcessor())
        urllib2.install_opener(self.opener)

    def post(self, url, data_param=None):
        '''A simple POST request wrapper'''
        if data_param is None:
            data = {}
        else:
            data = data_param
        if not url.startswith('http'):
            url = BASE_URL + url
        data['key'] = KEY
        encoded_args = urllib.urlencode(data)
        try:
            response = self.opener.open(url, encoded_args)
        except urllib2.URLError, exc:
            TermColor.print_error('Problem connecting to the Chartio service.'
                                  ' Please try again later.')
            f = open('/tmp/chartio-error.html', 'w')
            if hasattr(exc, 'read'):
                f.write(exc.read())
            f.close()
            sys.exit(1)
        return response.read()


def admin_user_password_dec(wrapped_fn):
    '''Decorator handling prompting of admin user/password as needed

    Requirements
        self.settings.admin_user data attribute

    References
        self.have_admin to determine whether admin is already configured
        self.admin_user_password() to set/validate/check rights

    '''
    def fn(self, *args, **kwargs):
        if not self.have_admin:
            TermColor.print_header('\nThis step requires administrator access.\n'
                                   'The credentials will be used during setup only for\n'
                                   ' - determining the names of local databases [optional]\n'
                                   ' - creating read-only user account [optional]\n'
                                   ' - granting read-only database access to the user account [optional]\n'
                                   'You may instead choose to do any or all of the above steps prior to\n'
                                   ' running chartio_setup. The script will prompt you to enter the\n'
                                   ' relevant information directly as needed.\n')
            admin_user = get_value('Database administrator name')
            TermColor.print_header('Please enter the password for the database'
                                   ' administrator (leave empty for none)')
            # Validate any value to permit a blank password
            admin_password = get_value('Database administrator password',
                                       is_password=True,
                                       validate=empty_is_ok)
            run_wrapped_fn = self.admin_user_password(admin_user, admin_password)
        else:
            run_wrapped_fn = True
        if run_wrapped_fn:
            retval = wrapped_fn(self, *args, **kwargs)
        else:
            retval = None
        return retval
    return fn


class DbAdmin(object):
    def __init__(self, chartio_poster, db_accessors):
        self.db_accessor = None
        self.db_accessors = db_accessors
        self.db_choice = None
        self.have_admin = False
        self.chartio_poster = chartio_poster
        self.settings = Settings()

    def admin_user_password(self, user, password):
        '''Means of validating and storing admin user and password'''
        self.have_admin = self.db_accessor.admin_user_password(user, password)
        return self.have_admin

    def readonly_user_password(self, database, user, password):
        '''Means of validating and storing readonly user and password'''
        rc = self.db_accessor.readonly_user_password(database, user, password)
        if rc:
            self.settings.readonly_user = user
            self.settings.readonly_password = password
        return rc

    def start_step(self):
        '''Beginning of interaction state machine'''
        return self.database_type_step

    def database_type_step(self):
        '''Get database type'''
        db_types = database_choices(self.chartio_poster)
        db_type_name = get_choice('What type of database are you connecting?',
                                  sorted(db_types))
        self.db_choice = db_types[db_type_name]
        self.settings.database_id = db_types[db_type_name]['id']
        self.db_accessor = self.db_accessors[db_type_name.lower()]()
        return self.database_listen_port_step

    def database_listen_port_step(self):
        '''Database listen port'''
        TermColor.print_header('Enter database local listen port')
        default = self.db_choice['default_port']
        port = get_value('Database listen port',
                         default=default,
                         validate=is_integer)
        self.db_accessor.port = port
        self.settings.database_port = port
        return self.database_name_step

    def database_name_step(self):
        '''Get database name'''
        TermColor.print_header('Enter the database name to connect to Chartio'
                               '\n(leave blank to list of available databases)')
        database = (get_value('Database name', validate=empty_is_ok)
                    or self.database_name_from_choice())
        if database is None:
            retval = False
        else:
            self.settings.database_name = database
            retval = self.readonly_user_step
        return retval

    @admin_user_password_dec
    def database_name_from_choice(self):
        '''Select database from automatically generated list (requires admin)'''
        databases = self.db_accessor.databases_get()
        if databases is None:
            retval = False
        else:
            TermColor.print_header('\nSelect which database to connect:')
            retval = get_choice('Database name', sorted(databases))
        return retval

    def readonly_user_step(self):
        '''Get/create read-only role and password and grant database access'''
        database = self.settings.database_name
        TermColor.print_header('Enter an existing read-only role for Chartio to use'
                               '\nor leave blank to create a new one automatically')
        user = get_value('Read-only role name', validate=empty_is_ok)
        if user:
            user_password = self.readonly_user_password_get(database, user=user)
            while not user_password:
                user_password = self.readonly_user_password_get(database, user=None)
        else:
            user_password = self.readonly_user_create_and_grant(database)
        retval = (user_password and self.datasource_register_step) or False
        return retval

    def readonly_user_password_get(self, database, user=None):
        '''Prompt for existing read-only database role

        Return
        (user, password) | False -- a tuple of the user and password iff apparently successful

        '''
        if user is None:
            user = get_value('Read-only role name')
        password = get_value('Read-only role password', is_password=True)
        if self.readonly_user_password(database, user, password):
            retval = (user, password)
        else:
            retval = False
        return retval

    @admin_user_password_dec
    def readonly_user_create_and_grant(self, database):
        '''Create a read-only user

        Return
        (user, password) | False -- a tuple of the user and password iff apparently successful

        '''
        user = name_generate(database.replace(' ', ''),
                             self.db_choice['user_name_limit'])
        password = random_password_generate()
        TermColor.print_delay('Creating read-only user %r' % (user))
        if (self.db_accessor.readonly_user_create_and_grant(database, user, password)
            and self.readonly_user_password(database, user, password)):
            retval = (user, password)
        else:
            retval = False
        return retval

    def datasource_register_step(self):
        return None


def argv_run(argv, input_str=None):
    '''Execute an argv with optional input
    
    Return
    ResultFalse|ResultTrue

    '''
    if input_str is None:
        stdin = None
    else:
        stdin = subprocess.PIPE
    try:
        proc = subprocess.Popen(argv,
                                stdin=stdin,
                                stderr=subprocess.STDOUT,
                                stdout=subprocess.PIPE)
        info, _unused = proc.communicate(input_str)
        proc_exited_cleanly = (0 == proc.returncode)
    except IOError:
        info = None
        proc_exited_cleanly = False
    if stdin is not None:
        proc.stdin.close()
    retval = {True: ResultTrue, False: ResultFalse}[proc_exited_cleanly](info)
    return retval


def argv_check(what, argv_result, check_fn):
    if argv_result:
        if isinstance(argv_result.info, basestring):
            info = argv_result.info.strip()
        else:
            info = argv_result.info
        rc = check_fn(info)
    else:
        rc = False
    if not rc:
        TermColor.print_error('Failed to ' + what + '.')
        if argv_result.info:
            TermColor.print_error(argv_result.info)
    return rc


def replacer_make(source, dest):
    '''A means of creating a replacement function which operates on strings.

    The use here is for escaping special characters from SQL arguments.

    '''
    def replacer(**kwargs):
        def item_replacer(item):
            key, value = item
            if isinstance(value, basestring):
                retval = (key, value.replace(source, dest))
            else:
                retval = item
            return retval
        retval = dict(map(item_replacer, kwargs.items()))
        return retval
    return replacer


class MysqlAccessor(object):
    '''Intermediary for accessing a MySQL database'''

    squote_escape = staticmethod(replacer_make("'", "''"))
    backtick_escape = staticmethod(replacer_make("`", "``"))

    def __init__(self):
        self.port = None
        self.admin_user = None
        self.admin_password = None

    def _sql_cmd_argv(self, user, password, database=None):
        argv = ['mysql',
                '--silent',
                '-u%s' % (user)]
        if self.port is not None:
            argv.extend(['--host=127.0.0.1',
                         '--port=%s' % (self.port),
                         '--protocol=tcp'])
        if password:
            argv.extend(['-p%s' % (password)])
        if database is not None:
            argv.append(database)
        return argv

    def sql_commands_print(self):
        cmds = ('''SHOW_DATABASES;''',
                '''SELECT 'SUCCESS';''',
                '''SELECT 'SUCCESS' FROM mysql.user WHERE user = '<USER>' AND host='127.0.0.1';''',
                ('''GRANT SELECT, SHOW VIEW ON <DATABASE>.* TO `<USER>`@`127.0.0.1`'''
                 ''' IDENTIFIED BY '<PASSWORD>';'''))
        for cmd in cmds:
            TermColor.print_cmd(cmd)

    def admin_user_password(self, user, password):
        '''Confirm and set admin user/password values'''
        rc = self.user_password_is_valid(user, password, database='mysql')
        if rc:
            self.admin_user = user
            self.admin_password = password
        return rc

    def readonly_user_password(self, database, user, password):
        '''Confirm and set readonly user/password values'''
        rc = self.user_password_is_valid(user, password, database)
        if rc:
            self.ro_user = user
            self.ro_password = password
        return rc

    def databases_get(self):
        '''Fetch a sequence of databases'''
        argv = self._sql_cmd_argv(self.admin_user, self.admin_password, database=None)
        sql = 'SHOW DATABASES;'
        sql_retval = argv_run(argv, sql)
        sql_rc = argv_check('list available databases', sql_retval, empty_is_ok)
        if sql_rc:
            databases = [db.strip() for db in sql_retval.info.split()]
        else:
            databases = None
        return databases

    def user_password_is_valid(self, user, password, database):
        '''Determine whether a user/password combination is valid for access to a given database'''
        argv = self._sql_cmd_argv(user, password, database)
        sql = '''SELECT 'SUCCESS';'''
        sql_retval = argv_run(argv, sql)
        rc = argv_check('validate user/password for %r' % (user),
                        sql_retval,
                        lambda x: 'SUCCESS' == x)
        return rc

    def user_has_role_create_access(self, user, password):
        argv = self._sql_cmd_argv(user, password, database='mysql')
        sql = ('''SELECT COUNT(1) FROM user WHERE User='%(user)s' AND grant_priv='y';'''
               % squote_escape(user=user))
        sql_retval = argv_run(argv, sql)
        rc = argv_check('verify user %r has role create access' % (user),
                        sql_retval,
                        lambda x: 'SUCCESS' == x)
        return rc

    def readonly_user_create_and_grant(self, database, ro_user, ro_password):
        admin_user, admin_password = self.admin_user, self.admin_password
        user_exists = self.readonly_user_exists(admin_user, admin_password, database, ro_user)
        if user_exists or (user_exists is None):
            rc = False
        else:
            rc = self.readonly_user_access_grant(admin_user, admin_password, database, ro_user, ro_password)
        return rc

    def readonly_user_exists(self, user, password, database, ro_user):
        argv = self._sql_cmd_argv(user, password, database)
        sql = ('''SELECT 'SUCCESS' FROM mysql.user WHERE User='%(user)s' AND Host='127.0.0.1';'''
               % self.squote_escape(user=ro_user))
        sql_retval = argv_run(argv, sql)
        if sql_retval:
            rc = ('SUCCESS' == sql_retval.info.strip())
            if rc:
                TermColor.print_error('Role %r already exists.' % (ro_user))
                TermColor.print_header('''You may reset the password with\n'''
                                       '''    SET PASSWORD FOR '%s'@'127.0.0.1' = 'New-Password';\n''' % (ro_user))
        else:
            TermColor.print_error('Error checking whether role %r exists.' % (ro_user))
            TermColor.print_error(sql_retval.info)
            rc = None
        return rc

    def readonly_user_access_grant(self, user, password, database, ro_user, ro_password):
        argv = self._sql_cmd_argv(user, password, database)
        params = self.squote_escape(role=ro_user, password=ro_password)
        params.update(self.backtick_escape(database=database))
        sql = ('''GRANT SELECT, SHOW VIEW ON `%(database)s`.* TO '%(role)s'@`127.0.0.1`'''
               ''' IDENTIFIED BY "%(password)s";'''
               % params)
        sql_retval = argv_run(argv, sql)
        rc = argv_check('create/grant access to database %r for role %r' % (database, user),
                        sql_retval,
                        lambda x: '' == x)
        return rc


class PostgresqlAccessor(object):
    '''Intermediary for accessing a PostgreSQL database'''

    dquote_escape = staticmethod(replacer_make('"', '""'))
    squote_escape = staticmethod(replacer_make("'", "''"))

    def __init__(self):
        self.port = None
        self.database = None
        self.admin_user = None
        self.admin_password = None
        self.ro_user = None
        self.ro_password = None

    def _sql_cmd_argv(self, user, password, database=None):
        argv = []
        if password:
            argv.extend(['env', 'PGPASSWORD=%s' % (password)])
        argv.extend(['psql',
                     '-t',
                     '-U', user,
                     '-v', 'ON_ERROR_STOP=1'])
        if self.port is not None:
            argv.extend(['-h', '127.0.0.1',
                         '-p', str(self.port)])
        if database is not None:
            argv.append(database)
        return argv

    def sql_commands_print(self):
        cmds = ('''SELECT 'SUCCESS';''',
                ('''SELECT 'SUCCESS' FROM pg_roles WHERE rolname='<ADMIN-USER>' '''
                 '''AND rolcreaterole='t';'''),
                '''SELECT 'SUCCESS' FROM pg_roles WHERE rolname='<READONLY-USER>';''',
                ('''CREATE USER "<READONLY-USER>" PASSWORD '<PASSWORD>' '''
                 '''NOSUPERUSER NOCREATEDB NOCREATEROLE NOINHERIT;'''),
                '''SELECT datname FROM pg_database;''',
                ('''SELECT relname'''
                 ''' FROM pg_class JOIN pg_namespace ON pg_namespace.oid = pg_class.relnamespace'''
                 ''' WHERE nspname = 'public' AND relkind IN ('r','v')'''
                 ''' ORDER BY relname ASC;'''),
                '''GRANT SELECT ON TABLE "<TABLE>" TO "<READONLY-USER>";''')
        for cmd in cmds:
            TermColor.print_cmd(cmd)

    def admin_user_password(self, user, password):
        '''Confirm and set admin user/password values'''
        rc = (self.user_password_is_valid(user, password, database=None)
              and self.user_has_role_create_access(user, password))
        if rc:
            self.admin_user = user
            self.admin_password = password
        return rc

    def readonly_user_password(self, database, user, password):
        '''Confirm and set readonly user/password values'''
        rc = (self.user_password_is_valid(user, password, database)
              and self.user_has_readonly_access(user, password, database))
        if rc:
            self.ro_user = user
            self.ro_password = password
        return rc

    def user_password_is_valid(self, user, password, database):
        '''Return bool'''
        argv = self._sql_cmd_argv(user, password, database)
        sql = '''SELECT 'SUCCESS';'''
        sql_retval = argv_run(argv, sql)
        rc = argv_check('validate user/password for %r' % (user),
                        sql_retval,
                        lambda x: 'SUCCESS' == x)
        return rc

    def user_has_role_create_access(self, user, password):
        argv = self._sql_cmd_argv(user, password, database=None)
        sql = ('''SELECT 'SUCCESS' FROM pg_roles WHERE rolname='%(rolname)s' '''
               ''' AND rolcreaterole='t';'''
               % self.dquote_escape(rolname=user))
        sql_retval = argv_run(argv, sql)
        rc = argv_check('verify user %r has role create access' % (user),
                        sql_retval,
                        lambda x: 'SUCCESS' == x)
        return rc

    def user_has_readonly_access(self, user, password, database):
        argv = self._sql_cmd_argv(user, password, database)
        sql = '''SELECT 'SUCCESS';'''
        sql_retval = argv_run(argv, sql)
        rc = argv_check('verify role %r has readonly access to %r' % (user, database),
                        sql_retval,
                        lambda x: 'SUCCESS' == x)
        return rc

    def readonly_user_create_and_grant(self, database, ro_user, ro_password):
        admin_user, admin_password = self.admin_user, self.admin_password
        user_exists = self.readonly_user_exists(admin_user, admin_password, database, ro_user)
        if user_exists or (user_exists is None):
            rc = False
        else:
            rc = (self.readonly_user_create(admin_user, admin_password,
                                            database,
                                            ro_user, ro_password)
                  and self.readonly_user_access_grant(admin_user, admin_password, database, ro_user))
        return rc

    def readonly_user_exists(self, user, password, database, ro_user):
        argv = self._sql_cmd_argv(user, password, database)
        sql = ('''SELECT 'SUCCESS' FROM pg_roles WHERE rolname='%(rolname)s';'''
               % self.squote_escape(rolname=ro_user))
        sql_retval = argv_run(argv, sql)
        if sql_retval:
            rc = ('SUCCESS' == sql_retval.info.strip())
            if rc:
                TermColor.print_error('Role %r already exists.' % (ro_user))
                TermColor.print_header('''You may reset the password with\n'''
                                       '''    ALTER ROLE '%s' PASSWORD 'New-Password'\n''' % (ro_user))
        else:
            TermColor.print_error('Error checking whether role %r exists.' % (ro_user))
            TermColor.print_error(sql_retval.info)
            rc = None
        return rc

    def readonly_user_create(self, user, password, database, ro_user, ro_password):
        '''Create a database-specific Postgresql read-only user'''
        argv = self._sql_cmd_argv(user, password, database)
        sql = ('''CREATE USER "%(role)s" PASSWORD '%(password)s' '''
               ''' NOSUPERUSER NOCREATEDB NOCREATEROLE NOINHERIT;'''
               % self.dquote_escape(role=ro_user, password=ro_password))
        sql_retval = argv_run(argv, sql)
        argv_check('create read-only user %r' % (ro_user),
                   sql_retval,
                   lambda x: 'CREATE ROLE' == x)
        return sql_retval

    def databases_get(self):
        '''List available databases

        Return
        seq | None -- a sequence of database names iff successful; None otherwise.

        '''
        argv = self._sql_cmd_argv(self.admin_user, self.admin_password, database=None)
        sql = '''SELECT datname FROM pg_database ORDER BY datname ASC;'''
        sql_retval = argv_run(argv, sql)
        sql_rc = argv_check('list available databases', sql_retval, empty_is_ok)
        if sql_rc:
            databases = [db.strip() for db in sql_retval.info.split()]
        else:
            databases = None
        return databases

    def readonly_user_access_grant(self, user, password, database, ro_user):
        tables_retval = self.tables_get(user, password, database)
        if tables_retval:
            tables = tables_retval.info.strip().split()
            for table in tables:
                if not self.readonly_user_table_access_grant(user, password, database, ro_user, table):
                    rc = False
                    break
            else:
                rc = True
        else:
            rc = False
        return rc

    def tables_get(self, user, password, database):
        argv = self._sql_cmd_argv(user, password, database)
        sql = ('''SELECT relname'''
               ''' FROM pg_class JOIN pg_namespace ON pg_namespace.oid=pg_class.relnamespace'''
               ''' WHERE nspname='public' AND relkind IN ('r','v')'''
               ''' ORDER BY relname ASC;''')
        sql_retval = argv_run(argv, sql)
        argv_check('determine tables in %r' % self.database,
                   sql_retval,
                   lambda x: True)
        return sql_retval

    def readonly_user_table_access_grant(self, user, password, database, ro_user, table):
        argv = self._sql_cmd_argv(user, password, database)
        sql = ('''GRANT SELECT ON TABLE "%(table)s" TO "%(role)s";'''
               % self.dquote_escape(table=table, role=ro_user))
        TermColor.print_delay('Granting %r read-only access to table %r' % (ro_user, table))
        sql_retval = argv_run(argv, sql)
        argv_check('grant table select access on %r to "%r"' % (table, ro_user),
                   sql_retval,
                   lambda x: 'GRANT' == x)
        return sql_retval


def create_ssh_conf():
    '''Attempt to create the config file.

    Complains and exits if the attempt fails. Good luck testing.

    '''
    try:
        f = open(CONFIG_DATA.CONFIG_FILE, 'a')
    except IOError, exc:
        # !!! Early exit
        sys.stderr.write('Unable to write to config file %r\n' % (CONFIG_DATA.CONFIG_FILE))
        sys.stderr.write('    %s\n' % (exc))
        sys.stderr.write('Exiting.\n')
        sys.exit(1)
    else:
        f.close()


def write_ssh_conf(key, value):
    '''Write a config key and value to a config file SSHTunnel section.

    Argumnents
    key -- the storage key
    value -- the storage value

    '''
    conf = ConfigParser.ConfigParser()
    if os.path.exists(CONFIG_DATA.CONFIG_FILE):
        conf.read(CONFIG_DATA.CONFIG_FILE)
    section = CONFIG_DATA.SSHTUNNEL_SECTION
    if section not in conf.sections():
        conf.add_section(section)
    conf.set(section, key, value)
    try:
        f = open(CONFIG_DATA.CONFIG_FILE, 'w')
    except IOError, exc:
        # !!! Early exit
        sys.stderr.write('Unable to open config file for writing: %s\n' % (exc))
        sys.stderr.write('Exiting.\n')
        sys.exit(1)
    conf.write(f)
    f.close()


def get_ssh_conf_value(key):
    '''Retrieve the value of a key from a config file SSHTunnel section.

    Return
    string or None -- the associated value on success; None if the
        config file, section, or key was not found.

    Arguments
    key -- the lookup key

    '''
    conf = ConfigParser.ConfigParser()
    if os.path.exists(CONFIG_DATA.CONFIG_FILE):
        conf.read(CONFIG_DATA.CONFIG_FILE)
        try:
            retval = conf.get(CONFIG_DATA.SSHTUNNEL_SECTION, key)
        except (ConfigParser.NoOptionError, ConfigParser.NoSectionError), e:
            retval = None
    else:
        retval = None
    return retval

def chartio_connect_start():
    '''Start chartio_connect daemon process.

    Exits if the return code is non-zero.

    '''
    # Launching chartio connect
    TermColor.print_delay('Launching chartio_connect')
    retcode = subprocess.call(['chartio_connect',
                               '-d',
                               '--prefix=%s' % (CONFIG_DATA.PREFIX)])
    if 0 == retcode:
        # Wait for connection to establish
        time.sleep(5)
        TermColor.print_delay('chartio_connect running')
    else:
        TermColor.print_error('Failed to launch chartio_connect. Exiting.')
        sys.exit(1)


def _exit(*args):
    print ''
    TermColor.print_ok('Exiting')
    sys.exit(0)


def opt_args_gather():
    parser = optparse.OptionParser(version='%prog 1.1.7')
    parser.add_option('--prefix',
                      help=('installation prefix for configuration'
                            'and runtime information. Defaults to %r' % (PREFIX_DEFAULT)))
    opt_args = parser.parse_args()
    return opt_args


def main():
    # Handle control-c
    signal.signal(signal.SIGINT, _exit)

    print 'Welcome to the chart.io setup wizard.'

    (options, args) = opt_args_gather()

    # This exits on failure
    config_data_set(options.prefix)

    # This exits on failure
    create_ssh_conf()

    # Confirm things have been installed
    proc = subprocess.Popen(['which', 'chartio_connect'],
                            stderr=subprocess.STDOUT,
                            stdout=subprocess.PIPE)
    (conn_location, which_err) = proc.communicate()
    if 0 != proc.returncode:
        TermColor.print_error('Chartio does not appear installed. Please run\n'
                              '  easy_install chartio\n'
                              'or\n'
                              '  python setup.py install')
    conn_location = os.path.abspath(conn_location).strip()

    # Instantiate API poster
    chartio_api = Poster()

    LOGIN_ATTEMPT_LIMIT = 3
    for login_attempt in range(LOGIN_ATTEMPT_LIMIT):
        email = get_value('Enter the email address registered with chart.io',
                          validate = lambda x: 0 < x.find('@'),
                          validate_explanation = 'This is not a valid email')
        password = get_value('Enter your chart.io password', is_password=True)

        # Login user
        response = chartio_api.post('/connectionclient/login/',
                                    {'email': email,
                                     'password': password})

        if response != 'success':
            TermColor.print_error(response)
        else:
            TermColor.print_delay('Username and password confirmed')
            break
    else:
        TermColor.print_error('Login tries exceeded.')
        sys.exit(1)

    TermColor.print_delay('Checking for existing SSH keys')
    if os.path.exists(CONFIG_DATA.SSH_KEY):
        TermColor.print_delay('SSH key found. Using the existing SSH key.')
    else:
        TermColor.print_delay('Generating keys for SSH tunneling')
        ret = subprocess.call([
            'ssh-keygen',
            '-q', # shhh!
            '-N', '', # No passphrase
            '-C', 'chart.io ssh tunneling',
            '-t', 'rsa',
            '-f', CONFIG_DATA.SSH_KEY,
        ])

        if ret != 0:
            TermColor.print_error('Failed to generate SSH key. Please confirm you have'
                                  ' ssh-keygen installed.')
            sys.exit(1)
        TermColor.print_delay('Generated SSH keys.')

    if not get_ssh_conf_value('client_id'):
        TermColor.print_delay('''Creating tunnel account on chart.io's server.'''
                              ''' This will take a moment.''')
        ssh_key = open('%s.pub' % CONFIG_DATA.SSH_KEY).read()
        response = chartio_api.post('/connectionclient/create/',
                                    {'email': email,
                                     'password': password,
                                     'ssh_key': ssh_key,
                                     'version': VERSION
                                     })

        response = json.loads(response)

        write_ssh_conf('remotehost', response['connection']['server_hostname'])
        write_ssh_conf('remoteuser', response['connection']['server_username'])
        write_ssh_conf('remoteport', response['connection']['port'])
        write_ssh_conf('client_id', response['connection']['connectionclient_id'])
        TermColor.print_delay('Tunnel account created')
    else:
        TermColor.print_delay('Connection tunnel already set up')

    # Get the project
    projects = json.loads(chartio_api.post('/connectionclient/projects/')).get('projects')
    if not projects:
        print ('\nUnable to find projects for your account. You must define a project'
               ' through the Chart.io web interface before running this.')
        sys.exit(1)
    elif 1 == len(projects):
        project = projects[0]
    else:
        project_map = dict([(p['name'], p) for p in projects])
        project_name = get_choice('\nYou have multiple projects.'
                                  ' To which project would you like to attatch this database',
                                  sorted(project_map.keys()))
        project = project_map[project_name]

    # Run through database administration
    db_accessors = {'mysql': MysqlAccessor,
                    'postgresql': PostgresqlAccessor}
    db_admin = DbAdmin(chartio_api, db_accessors)
    db_step = db_admin.start_step()
    while db_step:
        db_step = db_step()
    if db_step is None:
        TermColor.print_delay('Finished configuring database information')
    elif not db_step:
        # !!! Early exit
        TermColor.print_error('Exiting.')
        sys.exit(1)

    register_args = {'project_id': project['id'],
                     'type': db_admin.settings.database_id,
                     'connectionclient_id': get_ssh_conf_value('client_id'),
                     'port': db_admin.settings.database_port,
                     'name': db_admin.settings.database_name,
                     'user': db_admin.settings.readonly_user,
                     'passwd': db_admin.settings.readonly_password}

    write_ssh_conf('localport', db_admin.settings.database_port)
    chartio_connect_start()

    TermColor.print_delay('Registering datasource with Chartio. This will take a moment.')
    reg_response = chartio_api.post('/connectionclient/register/', register_args)
    if reg_response == 'success':
        TermColor.print_delay('Datasource registered. chartio_connect is running.\n')
        TermColor.print_ok('To ensure chartio reconnects after a reboot, add it to your crontab by typing')
        TermColor.print_header('  crontab -e')
        TermColor.print_ok('and entering this as an entry:')
        TermColor.print_header('  @reboot %s -d --prefix=%s' % (conn_location, CONFIG_DATA.PREFIX))
        TermColor.print_ok('\nThen visit')
        TermColor.print_header('  https://%s.chart.io/' % (project['slug']))
        TermColor.print_ok('to explore your data.')
    else:
        TermColor.print_error('Problem setting up your datasource. If this'
                              ' continues, please contact support@chart.io')


if '__main__' == __name__:
    main()
