~launchpad-pqm/launchpad/devel

#!/usr/bin/python -S

#

# Copyright 2009-2010 Canonical Ltd.  This software is licensed under the

# GNU Affero General Public License version 3 (see the file LICENSE).

__metaclass__ = type

# pylint: disable-msg=W0403

import _pythonpath

from collections import defaultdict

from ConfigParser import SafeConfigParser

from itertools import chain

from optparse import OptionParser

import os

import sys

import psycopg2

from canonical.database.sqlbase import connect

from canonical.launchpad.scripts import logger_options, logger, db_options

from fti import quote_identifier

import replication.helpers

# The 'read' group does not get given select permission on the following

# tables. This is to stop the ro user being given access to secrurity

# sensitive information that interactive sessions don't need.

SECURE_TABLES = [

    'public.accountpassword',

    'public.oauthnonce',

    'public.openidnonce',

    'public.openidconsumernonce',

    ]

class DbObject(object):

    def __init__(

            self, schema, name, type_, owner, arguments=None, language=None):

        self.schema = schema

        self.name = name

        self.type = type_

        self.owner = owner

        self.arguments = arguments

        self.language = language

    def __eq__(self, other):

        return self.schema == other.schema and self.name == other.name

    @property

    def fullname(self):

        fn = "%s.%s" % (self.schema, self.name)

        if self.type == 'function':

            fn = "%s(%s)" % (fn, self.arguments)

        return fn

    @property

    def seqname(self):

        if self.type != 'table':

            return ''

        return "%s.%s" % (self.schema, self.name + '_id_seq')

class DbSchema(dict):

    groups = None  # List of groups defined in the db

    users = None  # List of users defined in the db

    def __init__(self, con):

        super(DbSchema, self).__init__()

        cur = con.cursor()

        cur.execute('''

            SELECT

                n.nspname as "Schema",

                c.relname as "Name",

                CASE c.relkind

                    WHEN 'r' THEN 'table'

                    WHEN 'v' THEN 'view'

                    WHEN 'i' THEN 'index'

                    WHEN 'S' THEN 'sequence'

                    WHEN 's' THEN 'special'

                END as "Type",

                u.usename as "Owner"

            FROM pg_catalog.pg_class c

                LEFT JOIN pg_catalog.pg_user u ON u.usesysid = c.relowner

                LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace

            WHERE c.relkind IN ('r','v','S','')

                AND n.nspname NOT IN ('pg_catalog', 'pg_toast')

                AND pg_catalog.pg_table_is_visible(c.oid)

            ORDER BY 1,2

            ''')

        for schema, name, type_, owner in cur.fetchall():

            key = '%s.%s' % (schema, name)

            self[key] = DbObject(schema, name, type_, owner)

        cur.execute(r"""

            SELECT

                n.nspname as "schema",

                p.proname as "name",

                pg_catalog.oidvectortypes(p.proargtypes) as "Argument types",

                u.usename as "owner",

                l.lanname as "language"

            FROM pg_catalog.pg_proc p

                LEFT JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace

                LEFT JOIN pg_catalog.pg_language l ON l.oid = p.prolang

                LEFT JOIN pg_catalog.pg_user u ON u.usesysid = p.proowner

                LEFT JOIN pg_catalog.pg_type r ON r.oid = p.prorettype

            WHERE

                r.typname NOT IN ('trigger', 'language_handler')

                AND pg_catalog.pg_function_is_visible(p.oid)

                AND n.nspname <> 'pg_catalog'

                """)

        for schema, name, arguments, owner, language in cur.fetchall():

            self['%s.%s(%s)' % (schema, name, arguments)] = DbObject(

                    schema, name, 'function', owner, arguments, language)

        # Pull a list of groups

        cur.execute("SELECT groname FROM pg_group")

        self.groups = [r[0] for r in cur.fetchall()]

        # Pull a list of users

        cur.execute("SELECT usename FROM pg_user")

        self.users = [r[0] for r in cur.fetchall()]

    @property

    def principals(self):

        return chain(self.groups, self.users)

class CursorWrapper(object):

    def __init__(self, cursor):

        self.__dict__['_cursor'] = cursor

    def execute(self, cmd, params=None):

        cmd = cmd.encode('utf8')

        if params is None:

            log.debug2('%s' % (cmd, ))

            return self.__dict__['_cursor'].execute(cmd)

        else:

            log.debug2('%s [%r]' % (cmd, params))

            return self.__dict__['_cursor'].execute(cmd, params)

    def __getattr__(self, key):

        return getattr(self.__dict__['_cursor'], key)

    def __setattr__(self, key, value):

        return setattr(self.__dict__['_cursor'], key, value)

CONFIG_DEFAULTS = {

    'groups': '',

    }

def main(options):

    # Load the config file

    config = SafeConfigParser(CONFIG_DEFAULTS)

    configfile_name = os.path.join(os.path.dirname(__file__), 'security.cfg')

    config.read([configfile_name])

    con = connect(options.dbuser)

    if options.cluster:

        nodes = replication.helpers.get_nodes(con, 1)

        if nodes:

            # If we have a replicated environment, reset permissions on all

            # Nodes.

            con.close()

            for node in nodes:

                log.info("Resetting permissions on %s (%s)" % (

                    node.nickname, node.connection_string))

                reset_permissions(

                    psycopg2.connect(node.connection_string), config, options)

            return 0

        log.warning("--cluster requested, but not a Slony-I cluster.")

    log.info("Resetting permissions on single database")

    reset_permissions(con, config, options)

    return 0

def list_identifiers(identifiers):

    """List all of `identifiers` as SQL, quoted and separated by commas.

    :param identifiers: A sequence of SQL identifiers.

    :return: A comma-separated SQL string consisting of all identifiers

        passed in.  Each will be quoted for use in SQL.

    """

    return ', '.join([

        quote_identifier(identifier) for identifier in identifiers])

class PermissionGatherer:

    """Gather permissions for bulk granting or revocation.

    Processing such statements in bulk (with multiple users, tables,

    or permissions in one statement) is faster than issuing very large

    numbers of individual statements.

    """

    def __init__(self, entity_keyword):

        """Gather for SQL entities of one kind (TABLE, FUNCTION, SEQUENCE).

        :param entity_keyword: The SQL keyword for the kind of entity

            that permissions will be gathered for.

        """

        self.entity_keyword = entity_keyword

        self.permissions = defaultdict(dict)

    def add(self, permission, entity, principal, is_group=False):

        """Add a permission.

        Add all privileges you want to grant or revoke first, then use

        `grant` or `revoke` to process them in bulk.

        :param permission: A permission: SELECT, INSERT, EXECUTE, etc.

        :param entity: Table, function, or sequence on which to grant

            or revoke a privilege.

        :param principal: User or group to which the privilege should

            apply.

        :param is_group: Is `principal` a group?

        """

        if is_group:

            full_principal = "GROUP " + principal

        else:

            full_principal = principal

        self.permissions[permission].setdefault(entity, set()).add(

            full_principal)

    def tabulate(self):

        """Group privileges into single-statement work items.

        Each entry returned by this method represents a batch of

        privileges that can be granted or revoked in a single SQL

        statement.

        :return: A sequence of tuples of strings: permission(s) to

            grant/revoke, entity or entities to act on, and principal(s)

            to grant or revoke for.  Each is a string.

        """

        result = []

        for permission, parties in self.permissions.iteritems():

            for entity, principals in parties.iteritems():

                result.append(

                    (permission, entity, ", ".join(principals)))

        return result

    def countPermissions(self):

        """Count the number of different permissions."""

        return len(self.permissions)

    def countEntities(self):

        """Count the number of different entities."""

        return len(set(sum([

            entities.keys()

            for entities in self.permissions.itervalues()], [])))

    def countPrincipals(self):

        """Count the number of different principals."""

        principals = set()

        for entities_and_principals in self.permissions.itervalues():

            for extra_principals in entities_and_principals.itervalues():

                principals.update(extra_principals)

        return len(principals)

    def grant(self, cur):

        """Grant all gathered permissions.

        :param cur: A cursor to operate on.

        """

        log.debug(

            "Granting %d permission(s) on %d %s(s) for %d user(s)/group(s).",

            self.countPermissions(),

            self.countEntities(),

            self.entity_keyword,

            self.countPrincipals())

        grant_count = 0

        for permissions, entities, principals in self.tabulate():

            grant = "GRANT %s ON %s %s TO %s" % (

                permissions, self.entity_keyword, entities, principals)

            log.debug2(grant)

            cur.execute(grant)

            grant_count += 1

        log.debug("Issued %d GRANT statement(s).", grant_count)

    def revoke(self, cur):

        """Revoke all gathered permissions.

        :param cur: A cursor to operate on.

        """

        log.debug(

            "Revoking %d permission(s) on %d %s(s) for %d user(s)/group(s).",

            self.countPermissions(),

            self.countEntities(),

            self.entity_keyword,

            self.countPrincipals())

        revoke_count = 0

        for permissions, entities, principals in self.tabulate():

            revoke = "REVOKE %s ON %s %s FROM %s" % (

                permissions, self.entity_keyword, entities, principals)

            log.debug2(revoke)

            cur.execute(revoke)

            revoke_count += 1

        log.debug("Issued %d REVOKE statement(s).", revoke_count)

def reset_permissions(con, config, options):

    schema = DbSchema(con)

    all_users = list_identifiers(schema.users)

    cur = CursorWrapper(con.cursor())

    # Add our two automatically maintained groups

    for group in ['read', 'admin']:

        if group in schema.principals:

            log.debug("Removing managed users from %s role" % group)

            cur.execute("ALTER GROUP %s DROP USER %s" % (

                    quote_identifier(group), all_users))

        else:

            log.debug("Creating %s role" % group)

            cur.execute("CREATE GROUP %s" % quote_identifier(group))

            schema.groups.append(group)

    # Create all required groups and users.

    for section_name in config.sections():

        if section_name.lower() == 'public':

            continue

        assert not section_name.endswith('_ro'), (

            '_ro namespace is reserved (%s)' % repr(section_name))

        type_ = config.get(section_name, 'type')

        assert type_ in ['user', 'group'], 'Unknown type %s' % type_

        role_options = [

            'NOCREATEDB', 'NOCREATEROLE', 'NOCREATEUSER', 'INHERIT']

        if type_ == 'user':

            role_options.append('LOGIN')

        else:

            role_options.append('NOLOGIN')

        for username in [section_name, '%s_ro' % section_name]:

            if username in schema.principals:

                if type_ == 'group':

                    if options.revoke:

                        log.debug("Revoking membership of %s role", username)

                        cur.execute("REVOKE %s FROM %s" % (

                            quote_identifier(username), all_users))

                else:

                    # Note - we don't drop the user because it might own

                    # objects in other databases. We need to ensure they are

                    # not superusers though!

                    log.debug("Resetting role options of %s role.", username)

                    cur.execute(

                        "ALTER ROLE %s WITH %s" % (

                            quote_identifier(username),

                            ' '.join(role_options)))

            else:

                log.debug("Creating %s role.", username)

                cur.execute(

                    "CREATE ROLE %s WITH %s"

                    % (quote_identifier(username), ' '.join(role_options)))

                schema.groups.append(username)

        # Set default read-only mode for our roles.

        cur.execute(

            'ALTER ROLE %s SET default_transaction_read_only TO FALSE'

            % quote_identifier(section_name))

        cur.execute(

            'ALTER ROLE %s SET default_transaction_read_only TO TRUE'

            % quote_identifier('%s_ro' % section_name))

    # Add users to groups

    for user in config.sections():

        if config.get(user, 'type') != 'user':

            continue

        groups = [

            g.strip() for g in config.get(user, 'groups', '').split(',')

            if g.strip()]

        # Read-Only users get added to Read-Only groups.

        if user.endswith('_ro'):

            groups = ['%s_ro' % group for group in groups]

        if groups:

            log.debug("Adding %s to %s roles", user, ', '.join(groups))

            for group in groups:

                cur.execute(r"""ALTER GROUP %s ADD USER %s""" % (

                    quote_identifier(group), quote_identifier(user)))

        else:

            log.debug("%s not in any roles", user)

    if options.revoke:

        # Change ownership of all objects to OWNER.

        # We skip this in --no-revoke mode as ownership changes may

        # block on a live system.

        for obj in schema.values():

            if obj.type in ("function", "sequence"):

                pass  # Can't change ownership of functions or sequences

            else:

                if obj.owner != options.owner:

                    log.info("Resetting ownership of %s", obj.fullname)

                    cur.execute("ALTER TABLE %s OWNER TO %s" % (

                        obj.fullname, quote_identifier(options.owner)))

        # Revoke all privs from known groups. Don't revoke anything for

        # users or groups not defined in our security.cfg.

        table_revocations = PermissionGatherer("TABLE")

        function_revocations = PermissionGatherer("FUNCTION")

        sequence_revocations = PermissionGatherer("SEQUENCE")

        # Gather all revocations.

        for section_name in config.sections():

            role = quote_identifier(section_name)

            if section_name == 'public':

                ro_role = None

            else:

                ro_role = quote_identifier(section_name + "_ro")

            for obj in schema.values():

                if obj.type == 'function':

                    gatherer = function_revocations

                else:

                    gatherer = table_revocations

                gatherer.add("ALL", obj.fullname, role)

                if obj.seqname in schema:

                    sequence_revocations.add("ALL", obj.seqname, role)

                    if ro_role is not None:

                        sequence_revocations.add("ALL", obj.seqname, ro_role)

        table_revocations.revoke(cur)

        function_revocations.revoke(cur)

        sequence_revocations.revoke(cur)

    else:

        log.info("Not resetting ownership of database objects")

        log.info("Not revoking permissions on database objects")

    # Set of all tables we have granted permissions on. After we have assigned

    # permissions, we can use this to determine what tables have been

    # forgotten about.

    found = set()

    # Set permissions as per config file

    table_permissions = PermissionGatherer("TABLE")

    function_permissions = PermissionGatherer("FUNCTION")

    sequence_permissions = PermissionGatherer("SEQUENCE")

    for username in config.sections():

        for obj_name, perm in config.items(username):

            if '.' not in obj_name:

                continue

            if obj_name not in schema.keys():

                log.warn('Bad object name %r', obj_name)

                continue

            obj = schema[obj_name]

            found.add(obj)

            perm = perm.strip()

            if not perm:

                # No perm means no rights. We can't grant no rights, so skip.

                continue

            who = quote_identifier(username)

            if username == 'public':

                who_ro = who

            else:

                who_ro = quote_identifier('%s_ro' % username)

            log.debug(

                "Granting %s on %s to %s", perm, obj.fullname, who)

            if obj.type == 'function':

                function_permissions.add(perm, obj.fullname, who)

                function_permissions.add("EXECUTE", obj.fullname, who_ro)

                function_permissions.add(

                    "EXECUTE", obj.fullname, "read", is_group=True)

                function_permissions.add(

                    "ALL", obj.fullname, "admin", is_group=True)

            else:

                table_permissions.add(

                    "ALL", obj.fullname, "admin", is_group=True)

                table_permissions.add(perm, obj.fullname, who)

                table_permissions.add("SELECT", obj.fullname, who_ro)

                is_secure = (obj.fullname in SECURE_TABLES)

                if not is_secure:

                    table_permissions.add(

                        "SELECT", obj.fullname, "read", is_group=True)

                if obj.seqname in schema:

                    if 'INSERT' in perm:

                        seqperm = 'USAGE'

                    elif 'SELECT' in perm:

                        seqperm = 'SELECT'

                    sequence_permissions.add(seqperm, obj.seqname, who)

                    if not is_secure:

                        sequence_permissions.add(

                            "SELECT", obj.seqname, "read", is_group=True)

                    sequence_permissions.add("SELECT", obj.seqname, who_ro)

                    sequence_permissions.add(

                        "ALL", obj.seqname, "admin", is_group=True)

    function_permissions.grant(cur)

    table_permissions.grant(cur)

    sequence_permissions.grant(cur)

    # Set permissions on public schemas

    public_schemas = [

        s.strip() for s in config.get('DEFAULT', 'public_schemas').split(',')

        if s.strip()]

    log.debug("Granting access to %d public schemas", len(public_schemas))

    for schema_name in public_schemas:

        cur.execute("GRANT USAGE ON SCHEMA %s TO PUBLIC" % (

            quote_identifier(schema_name),

            ))

    for obj in schema.values():

        if obj.schema not in public_schemas:

            continue

        found.add(obj)

        if obj.type == 'function':

            cur.execute('GRANT EXECUTE ON FUNCTION %s TO PUBLIC' %

                        obj.fullname)

        else:

            cur.execute('GRANT SELECT ON TABLE %s TO PUBLIC' % obj.fullname)

    # Raise an error if we have database objects lying around that have not

    # had permissions assigned.

    forgotten = set()

    for obj in schema.values():

        if obj not in found:

            forgotten.add(obj)

    forgotten = [obj.fullname for obj in forgotten

        if obj.type in ['table', 'function', 'view']]

    if forgotten:

        log.warn('No permissions specified for %r', forgotten)

    if options.dryrun:

        log.info("Dry run - rolling back changes")

        con.rollback()

    else:

        log.debug("Committing changes")

        con.commit()

if __name__ == '__main__':

    parser = OptionParser()

    parser.add_option(

        "-n", "--dry-run", dest="dryrun", default=False,

        action="store_true", help="Don't commit any changes")

    parser.add_option(

        "--revoke", dest="revoke", default=True, action="store_true",

        help="Revoke privileges as well as add them")

    parser.add_option(

        "--no-revoke", dest="revoke", default=True, action="store_false",

        help="Do not revoke any privileges. Just add.")

    parser.add_option(

        "-o", "--owner", dest="owner", default="postgres",

        help="Owner of PostgreSQL objects")

    parser.add_option(

        "-c", "--cluster", dest="cluster", default=False,

        action="store_true",

        help="Rebuild permissions on all nodes in the Slony-I cluster.")

    db_options(parser)

    logger_options(parser)

    (options, args) = parser.parse_args()

    log = logger(options)

    sys.exit(main(options))