path: root/rp/config/rpki-sql-setup

#!/usr/bin/env python

# $Id$
#
# Copyright (C) 2015-2016  Parsons Government Services ("PARSONS")
# Portions 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 PARSONS, DRL, AND ISC DISCLAIM
# ALL WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS.  IN NO EVENT SHALL
# PARSONS, 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 textwrap
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 db_accessible(self, udb):
        try:
            self._db_accessible_test(udb)
        except:
            return False
        else:
            return True

    def fetchone(self):
        return self._cur.fetchone()

    def fetchall(self):
        return self._cur.fetchall()

    def close(self):
        self._cur.close()
        self._db.close()

    def log(self, msg):
        if self.args.verbose:
            sys.stderr.write(msg + "\n")


class MySQL_Driver(Abstract_Driver):

    _initialized = False

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

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

    def db_exists(self, udb):
        self.execute("SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = '{0.database}'".format(udb))
        return bool(self.fetchone()[0])

    def execute(self, query, args = None):
        try:
            self._cur
        except AttributeError:
            self.log("MySQL driver initializing root connection")
            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._db.autocommit(True)
            self._cur = self._db.cursor()
        self.log("MySQL driver executing {!r} {!r}".format(query, args))
        return self._cur.execute(query, args)

    def create(self, udb):
        self.execute("CREATE DATABASE IF NOT EXISTS {0.database}".format(udb))
        self.fix_grants(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("GRANT ALL ON {0.database}.* TO {0.username}@localhost IDENTIFIED BY %s".format(udb),
                     (udb.password,))

class SQLite3_Driver(Abstract_Driver):

    def __init__(self, args):
        try:
            self.driver
        except AttributeError:
            import sqlite3
            self.driver = sqlite3
            self.args = args

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

    def db_exists(self, udb):
        return os.path.exists(udb.database)

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

    def create(self, udb):
        self._db_accessible_test(udb.database)
        self._grant(udb)

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

    def script_drop(self, udb):
        self.args.script_output.write("rm {}\n".format(udb.database))

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


class PostgreSQL_Driver(Abstract_Driver):

    def __init__(self, args):
        try:
            self.driver
        except AttributeError:
            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
            self.log("Initialized PostgreSQL driver, pw {!r}".format(self._pw))

    def _seteuid(self, new_uid):
        old_uid = os.geteuid()
        if new_uid != old_uid:
            self.log("PostgreSQL driver changing EUID from {} to {}".format(old_uid, new_uid))
            os.seteuid(new_uid)
        return old_uid

    def execute(self, query, args = None):
        try:
            self._cur
        except AttributeError:
            self.log("PostgreSQL driver opening connection to database {}".format(self.args.postgresql_root_database))
            if self._pw is not None:
                euid = self._seteuid(self._pw.pw_uid)
            try:
                self._db = self.driver.connect(database = self.args.postgresql_root_database)
                self._db.autocommit = True
                self._cur = self._db.cursor()
            finally:
                if self._pw is not None:
                    self._seteuid(euid)
        self.log("PostgreSQL driver executing {!r} {!r}".format(query, args))
        return self._cur.execute(query, args)

    def _db_accessible_test(self, udb):
        pw = pwd.getpwnam(udb.username)
        uid = self._seteuid(pw.pw_uid)
        try:
            self.driver.connect(database = udb.database, user = udb.username , password = udb.password).close()
        finally:
            self._seteuid(uid)

    def db_exists(self, udb):
        self.execute("SELECT COUNT(*) FROM pg_database WHERE datname = '{0.database}'".format(udb))
        return bool(self.fetchone()[0])

    def role_in_use(self, udb):
        self.execute(textwrap.dedent('''\
            SELECT COUNT(*) FROM pg_database
            JOIN pg_roles ON pg_database.datdba = pg_roles.oid
            WHERE pg_roles.rolname = '{0.username}'
            '''.format(udb)))
        return bool(self.fetchone()[0])

    def create(self, udb):
        if not self.role_in_use(udb):
            self.execute("CREATE ROLE {0.username} LOGIN PASSWORD '{0.password}'".format(udb))
        if not self.db_exists(udb):
            self.execute("CREATE DATABASE {0.database} OWNER {0.username}".format(udb))

    def drop(self, udb):
        self.execute("DROP DATABASE IF EXISTS {0.database}".format(udb))
        if not self.role_in_use(udb):
            self.execute("DROP ROLE IF EXISTS {0.username}".format(udb))

    def script_drop(self, udb):
        self.args.script_output.write(textwrap.dedent('''\
            DROP DATABASE IF EXISTS {0.database};
            DO $$ BEGIN
                IF NOT EXISTS (SELECT * FROM pg_database JOIN pg_roles
                               ON pg_database.datdba = pg_roles.oid
                               WHERE pg_roles.rolname = '{0.username}')
                THEN
                    DROP ROLE IF EXISTS {0.username};
                END IF;
            END $$;
            '''.format(udb)))

    def fix_grants(self, udb):
        self.execute("ALTER DATABASE {0.database} OWNER TO {0.username}".format(udb))
        self.execute("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.db_accessible(self):
            self.driver.drop(self)

    def create(self):
        if self.args.force or not self.driver.db_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.db_accessible(self):
            self.driver.drop(self)
        self.driver.create(self)

    def fix_grants(self):
        if self.args.force or not self.driver.db_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)]
    names.append("rcynic")

    # For now, we quietly ignore missing sections rather than throwing an exception.
    # I could make a case either way for this, but ignoring missing sections is a
    # lot easier to clean up while debugging the installation scripts.

    for name in names:
        if cfg.has_section(name):
            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))