#!/usr/bin/env python

# $Id$
#
# Copyright (C) 2014  Dragon Research Labs ("DRL")
# Portions copyright (C) 2009-2013  Internet Systems Consortium ("ISC")
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notices and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND DRL AND ISC DISCLAIM ALL
# WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS.  IN NO EVENT SHALL DRL OR
# ISC BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL
# DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA
# OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
# TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
# PERFORMANCE OF THIS SOFTWARE.

"""
Automated setup of SQL stuff used by the RPKI tools.  Pulls
configuration from rpki.conf, prompts for SQL password when needed.
"""

import os
import pwd
import sys
import getpass
import argparse
import rpki.config


class Abstract_Driver(object):

    # Kludge to make classes derived from this into singletons.  Net
    # of a Million Lies says this is Not Pythonic, but it seems to
    # work, so long as one doesn't attempt to subclass the resulting
    # driver classes.  For our purposes, it will do.

    __instance = None

    def __new__(cls, *args, **kwargs):
        if cls.__instance is None:
            cls.__instance = object.__new__(cls, *args, **kwargs)
        return cls.__instance

    def accessible(self, udb):
        try:
            self._accessible_test(udb)
        except:
            return False
        else:
            return True

class MySQL_Driver(Abstract_Driver):

    _initialized = False

    def __init__(self, args):
        from rpki.mysql_import import MySQLdb
        self.driver = MySQLdb
        self.args = args

    def _initialize(self):
        if not self._initialized:
            if self.args.verbose:
                print "Initializing MySQL driver"
            if self.args.mysql_defaults:
                mysql_cfg = rpki.config.parser(set_filename = self.args.mysql_defaults, section = "client")
                self._db = self.driver.connect(db     = "mysql",
                                               user   = mysql_cfg.get("user"),
                                               passwd = mysql_cfg.get("password"))
            else:
                self._db = self.driver.connect(db     = "mysql",
                                               user   = "root",
                                               passwd = getpass.getpass("Please enter your MySQL root password: "))
            self._cur = self._db.cursor()
            self._initialized = True

    def _accessible_test(self, udb):
        self.driver.connect(db = udb.database, user = udb.username, passwd = udb.password).close()

    def _grant(self, udb):
        self._cur.execute("GRANT ALL ON {0.database}.* TO {0.username}@localhost IDENTIFIED BY %s".format(udb),
                          (udb.password,))

    def create(self, udb):
        self._initialize()
        self._cur.execute("CREATE DATABASE IF NOT EXISTS {0.database}".format(udb))
        self._grant(udb)
        self._db.commit()

    def drop(self, udb):
        self._initialize()
        self._cur.execute("DROP DATABASE IF EXISTS {0.database}".format(udb))
        self._db.commit()

    def script_drop(self, udb):
        self.args.script_output.write("DROP DATABASE IF EXISTS {};\n".format(udb.database))

    def fix_grants(self, udb):
        self._grant(udb)
        self._db.commit()


class SQLite3_Driver(Abstract_Driver):

    def __init__(self, args):
        import sqlite3
        self.driver = sqlite3
        self.args = args
        self.can_chown = os.getuid() == 0 or os.geteuid() == 0

    def _accessible_test(self, udb):
        self.driver.connect(udb.database).close()

    def _grant(self, udb):
        if self.can_chown and udb.username:
            pw = pwd.getpwnam(udb.username)
            os.chown(udb.database, pw.pw_uid, pw.pw_gid)

    def create(self, udb):
        self.driver.connect(udb.database).close()
        self._grant(udb)

    def drop(self, udb):
        os.unlink(udb.database)

    def script_drop(self, udb):
        pass

    def fix_grants(self, udb):
        self._grant(udb)


class PostgreSQL_Driver(Abstract_Driver):

    def __init__(self, args):
        import psycopg2
        self.driver = psycopg2
        self.args = args
        if args.postgresql_root_username and (os.getuid() == 0 or os.geteuid() == 0):
            self._pw = pwd.getpwnam(args.postgresql_root_username)
        else:
            self._pw = None
        if self.args.verbose:
            print "Initialized PostgreSQL driver, pw {!r}".format(pw)

    def _execute(self, *sql_commands):
        try:
            pid = None if self._pw is None else os.fork()
            if self.args.verbose:
                print "PostgreSQL driver fork {}".format(pid)
            if pid == 0:
                os.setgid(self._pw.pw_gid)
                os.setuid(self._pw.pw_uid)
            if not pid:
                with self.driver.connect(database = self.args.postgresql_root_database) as db:
                    with db.cursor() as cur:
                        for sql_command in sql_commands:
                            if self.args.verbose:
                                print "PostgreSQL driver executing {!r}".format(sql_command)
                            cur.execute(command)
            if pid == 0:
                os._exit(0)
            elif pid:
                os.waitpid(pid, 0)
        except Exception as e:
            if self.args.verbose:
                print "PostgreSQL driver exception {!s}".format(e)
            if pid == 0:
                os._exit(1)

    def _accessible_test(self, udb):
        self.driver.connect(database = udb.database, user = udb.username , password = usb.password).close()

    # At some point we'll have to do something about DROP ROLE [IF EXISTS],
    # but it's a bit complicated because we need to defer dropping the role until
    # after we've dropped all associated databases, which gets messy when
    # interleaved with all the other things we're doing, and may require
    # restructuring all of the drivers to maintain a queue of actions to be taken
    # so that we can make sure that all DROP ROLEs go after all DROP DATABASEs.
    #
    # Punt on this for now, but will need to come back to it, particularly if we're
    # serious about using PostgreSQL on Debian and cleaning up after ourselves.

    def create(self, udb):
        #
        # CREATE ROLE doesn't take a IF NOT EXISTS modifier, but we can fake it using plpgsql.
        # http://stackoverflow.com/questions/8092086/create-postgresql-role-user-if-it-doesnt-exist
        #
        self._execute('''
                        DO $$ BEGIN
                          IF NOT EXISTS (SELECT * FROM pg_catalog.pg_user WHERE usename = '{0.username}') THEN
                            CREATE ROLE {0.username} LOGIN PASSWORD '{0.password}';
                          END IF;
                        END $$
                      '''.format(udb),
                      "CREATE DATABASE IF NOT EXISTS {0.database} OWNER {0.username}".format(udb))

    def drop(self, udb):
        self._execute("DROP DATABASE IF EXISTS {0.database}".format(udb))

    def script_drop(self, udb):
        self.args.script_output.write("DROP DATABASE IF EXISTS {};\n".format(udb.database))

    def fix_grants(self, udb):
        self._execute("ALTER DATABASE {0.database} OWNER TO {0.username}".format(udb),
                      "ALTER ROLE {0.username} WITH PASSWORD '{0.password}".format(udb))


class UserDB(object):
    """
    Class to wrap access parameters for a particular database.
    """

    drivers = dict(sqlite3    =    SQLite3_Driver,
                   mysql      =      MySQL_Driver,
                   postgresql = PostgreSQL_Driver)

    def __init__(self, args, name):
        self.database = cfg.get("sql-database", section = name)
        self.username = cfg.get("sql-username", section = name)
        self.password = cfg.get("sql-password", section = name)
        self.engine   = cfg.get("sql-engine",   section = name)
        self.driver   = self.drivers[self.engine](args)
        self.args     = args

    def drop(self):
        if self.args.force or self.driver.accessible(self):
            self.driver.drop(self)

    def create(self):
        if self.args.force or not self.driver.accessible(self):
            self.driver.create(self)

    def script_drop(self):
        self.driver.script_drop(self)

    def drop_and_create(self):
        if self.args.force or self.driver.accessible(self):
            self.driver.drop(self)
        self.driver.create(self)

    def fix_grants(self):
        if self.args.force or not self.driver.accessible(self):
            self.driver.fix_grants(self)


parser = argparse.ArgumentParser(description = __doc__)
parser.add_argument("-c", "--config",
                    help = "specify alternate location for rpki.conf")
parser.add_argument("-d", "--debug", action = "store_true",
                    help = "enable debugging (eg, Python backtraces)")
parser.add_argument("-v", "--verbose", action = "store_true",
                    help = "whistle while you work")
parser.add_argument("-f", "--force", action = "store_true",
                    help = "force database create, drop, or grant regardless of current state")

parser.add_argument("--mysql-defaults",
                    help = "specify MySQL root access credentials via a configuration file")


parser.add_argument("--postgresql-root-database", default = "postgres",
                    help = "name of PostgreSQL control database")
parser.add_argument("--postgresql-root-username",
                    help = "username of PostgreSQL control role")

subparsers = parser.add_subparsers(title = "Commands", metavar = "", dest = "dispatch")

subparsers.add_parser("create",
                      help = "create databases and load schemas")

subparsers.add_parser("drop",
                      help = "drop databases")

subparser = subparsers.add_parser("script-drop",
                                  help = "show SQL commands to drop databases")
subparser.add_argument("script_output",
                       nargs = "?", type = argparse.FileType("w"), default = "-",
                       help = "destination for drop script")

subparsers.add_parser("drop-and-create",
                      help = "drop databases then recreate them and load schemas")

subparsers.add_parser("fix-grants",
                      help = "whack database to match configuration file")

args = parser.parse_args()

try:

    cfg  = rpki.config.parser(set_filename = args.config, section = "myrpki")

    names = [name for name in ("irdbd", "rpkid", "pubd")
             if cfg.getboolean("start_" + name, False)]

    # In the long run this should probably become mandatory, but for
    # the moment I'm getting installation errors because we don't yet
    # have an [rcynic] section in rpki.conf, and I don't want to shave
    # that yak today.

    if cfg.has_section("rcynic"):
        names.append("rcynic")

    for name in names:
        udb = UserDB(args = args, name = name)
        method = args.dispatch.replace("-", "_")
        getattr(udb, method)()

except Exception, e:
    if args.debug:
        raise
    else:
        sys.exit(str(e))