ansible/postgresql_user
Pepe Barbe 511ab8697c Query for all active privileges instead
Use a different method to query for current 
privileges at the table and database level. 
This method is more robust if newer privileges 
are added in future versions and also supports the 
ALL wildcard.
2012-08-22 10:04:58 -05:00

289 lines
9.3 KiB
Python
Executable file

#!/usr/bin/python
# -*- coding: utf-8 -*-
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible. If not, see <http://www.gnu.org/licenses/>.
import re
try:
import psycopg2
except ImportError:
postgresqldb_found = False
else:
postgresqldb_found = True
# ===========================================
# PostgreSQL module specific support methods.
#
def user_exists(cursor, user):
query = "SELECT rolname FROM pg_roles WHERE rolname=%(user)s"
cursor.execute(query, {'user': user})
return cursor.rowcount > 0
def user_add(cursor, user, password):
"""Create a new user with write access to the database"""
query = "CREATE USER %(user)s with PASSWORD '%(password)s'"
cursor.execute(query % {"user": user, "password": password})
return True
def user_chpass(cursor, user, password):
"""Change user password"""
changed = False
# Handle passwords.
if password is not None:
select = "SELECT rolpassword FROM pg_authid where rolname=%(user)s"
cursor.execute(select, {"user": user})
current_pass_hash = cursor.fetchone()[0]
# Not sure how to hash the new password, so we just initiate the
# change and check if the hash changed
alter = "ALTER USER %(user)s WITH PASSWORD '%(password)s'"
cursor.execute(alter % {"user": user, "password": password})
cursor.execute(select, {"user": user})
new_pass_hash = cursor.fetchone()[0]
if current_pass_hash != new_pass_hash:
changed = True
return changed
def user_delete(cursor, user):
"""Try to remove a user. Returns True if successful otherwise False"""
cursor.execute("SAVEPOINT ansible_pgsql_user_delete")
try:
cursor.execute("DROP USER %s" % user)
except:
cursor.execute("ROLLBACK TO SAVEPOINT ansible_pgsql_user_delete")
cursor.execute("RELEASE SAVEPOINT ansible_pgsql_user_delete")
return False
cursor.execute("RELEASE SAVEPOINT ansible_pgsql_user_delete")
return True
def has_table_privilege(cursor, user, table, priv):
query = 'SELECT has_table_privilege(%s, %s, %s)'
cursor.execute(query, (user, table, priv))
return cursor.fetchone()[0]
def get_table_privileges(cursor, user, table):
if '.' in table:
schema, table = table.split('.', 1)
else:
schema = 'public'
query = '''SELECT privilege_type FROM information_schema.role_table_grants
WHERE grantee=%s AND table_name=%s AND table_schema=%s'''
cursor.execute(query, (user, table, schema))
return set([x[0] for x in cursor.fetchall()])
def grant_table_privilege(cursor, user, table, priv):
prev_priv = get_table_privileges(cursor, user, table)
query = 'GRANT %s ON TABLE %s TO %s' % (priv, table, user)
cursor.execute(query)
curr_priv = get_table_privileges(cursor, user, table)
return len(curr_priv) > len(prev_priv)
def revoke_table_privilege(cursor, user, table, priv):
prev_priv = get_table_privileges(cursor, user, table)
query = 'REVOKE %s ON TABLE %s FROM %s' % (priv, table, user)
cursor.execute(query)
curr_priv = get_table_privileges(cursor, user, table)
return len(curr_priv) < len(prev_priv)
def get_database_privileges(cursor, user, db):
priv_map = {
'C':'CREATE',
'T':'TEMPORARY',
'c':'CONNECT',
}
query = 'SELECT datacl FROM pg_database WHERE datname = %s'
cursor.execute(query, (db,))
datacl = cursor.fetchone()[0]
r = re.search('%s=(C?T?c?)/[a-z]+\,?' % user, datacl)
if r is None:
return []
o = []
for v in r.group(1):
o.append(priv_map[v])
return o
def has_database_privilege(cursor, user, db, priv):
query = 'SELECT has_database_privilege(%s, %s, %s)'
cursor.execute(query, (user, db, priv))
return cursor.fetchone()[0]
def grant_database_privilege(cursor, user, db, priv):
prev_priv = get_database_privileges(cursor, user, db)
query = 'GRANT %s ON DATABASE %s TO %s' % (priv, db, user)
cursor.execute(query)
curr_priv = get_database_privileges(cursor, user, db)
return len(curr_priv) > len(prev_priv)
def revoke_database_privilege(cursor, user, db, priv):
prev_priv = get_database_privileges(cursor, user, db)
query = 'REVOKE %s ON DATABASE %s FROM %s' % (priv, db, user)
cursor.execute(query)
curr_priv = get_database_privileges(cursor, user, db)
return len(curr_priv) < len(prev_priv)
def revoke_privileges(cursor, user, privs):
if privs is None:
return False
changed = False
for type_ in privs:
revoke_func = {
'table':revoke_table_privilege,
'database':revoke_database_privilege
}[type_]
for name, privileges in privs[type_].iteritems():
for privilege in privileges:
changed = revoke_func(cursor, user, name, privilege)\
or changed
return changed
def grant_privileges(cursor, user, privs):
if privs is None:
return False
changed = False
for type_ in privs:
grant_func = {
'table':grant_table_privilege,
'database':grant_database_privilege
}[type_]
for name, privileges in privs[type_].iteritems():
for privilege in privileges:
changed = grant_func(cursor, user, name, privilege)\
or changed
return changed
def parse_privs(privs, db):
"""
Parse privilege string to determine permissions for database db.
Format:
privileges[/privileges/...]
Where:
privileges := DATABASE_PRIVILEGES[,DATABASE_PRIVILEGES,...] |
TABLE_NAME:TABLE_PRIVILEGES[,TABLE_PRIVILEGES,...]
"""
if privs is None:
return privs
o_privs = {
'database':{},
'table':{}
}
for token in privs.split('/'):
if ':' not in token:
type_ = 'database'
name = db
priv_set = set(x.strip() for x in token.split(','))
else:
type_ = 'table'
name, privileges = token.split(':', 1)
priv_set = set(x.strip() for x in privileges.split(','))
o_privs[type_][name] = priv_set
return o_privs
# ===========================================
# Module execution.
#
def main():
module = AnsibleModule(
argument_spec=dict(
login_user=dict(default="postgres"),
login_password=dict(default=""),
login_host=dict(default=""),
user=dict(required=True, aliases=['name']),
password=dict(default=None),
state=dict(default="present", choices=["absent", "present"]),
priv=dict(default=None),
db=dict(default=''),
fail_on_user=dict(default='yes')
)
)
user = module.params["user"]
password = module.params["password"]
state = module.params["state"]
fail_on_user = module.params["fail_on_user"] == 'yes'
db = module.params["db"]
if db == '' and module.params["priv"] is not None:
module.fail_json(msg="privileges require a database to be specified")
privs = parse_privs(module.params["priv"], db)
if not postgresqldb_found:
module.fail_json(msg="the python psycopg2 module is required")
# To use defaults values, keyword arguments must be absent, so
# check which values are empty and don't include in the **kw
# dictionary
params_map = {
"login_host":"host",
"login_user":"user",
"login_password":"password",
"db":"database"
}
kw = dict( (params_map[k], v) for (k, v) in module.params.iteritems()
if k in params_map and v != "" )
try:
db_connection = psycopg2.connect(**kw)
cursor = db_connection.cursor()
except Exception, e:
module.fail_json(msg="unable to connect to database: %s" % e)
kw = dict(user=user)
changed = False
user_removed = False
if state == "present":
if user_exists(cursor, user):
changed = user_chpass(cursor, user, password)
else:
if password is None:
msg = "password parameter required when adding a user"
module.fail_json(msg=msg)
changed = user_add(cursor, user, password)
changed = grant_privileges(cursor, user, privs) or changed
else:
if user_exists(cursor, user):
changed = revoke_privileges(cursor, user, privs)
user_removed = user_delete(cursor, user)
changed = changed or user_removed
if fail_on_user and not user_removed:
msg = "unabel to remove user"
module.fail_json(msg=msg)
kw['user_removed'] = user_removed
if changed:
db_connection.commit()
kw['changed'] = changed
module.exit_json(**kw)
# this is magic, see lib/ansible/module_common.py
#<<INCLUDE_ANSIBLE_MODULE_COMMON>>
main()