#!/bin/env python

import sys
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT

def inputStuff(title, default = None, empty = True):
    if default is not None:
        uinput = raw_input("%s [DEFAULT %s]: " %(title, default))
    else:
        uinput = raw_input("%s: " %(title))
    if default:
        if uinput == '':
            return default
    if uinput == '' and not empty:
        print "Error: Zero length input not allowed!"
        return inputStuff(title, default, empty)
    return uinput

def main():
    if len(sys.argv) == 6:
        thiscript, hostnameip, dspamuser, dspampass, dspamdb, createscript = sys.argv
    elif len(sys.argv) == 7:
        thiscript, hostnameip, dspamuser, dspampass, dspamdb, createscript, vuserscript = sys.argv
    else:
        print "Usage: %s [database host] [dspam user] [dspam passwd] [dspam dbname] [set-up DB script] [OPTIONAL: set-up virutal-user script]" %(sys.argv[0])
        sys.exit(1)
    con = connect(hostnameip)
    con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT) # CREATE DATABASE needs this
    cur = con.cursor()
    createDbAndUser(cur, dspamuser, dspampass, dspamdb)
    con.close()
    executeScript(dspamuser, dspampass, dspamdb, hostnameip, createscript)
    if len(sys.argv) == 7:
        executeScript(dspamuser, dspampass, dspamdb, hostnameip, vuserscript)

def createDbAndUser(cur, dspamuser, dspampass, dspamdb):
    print "Creating DSPAM PostgreSQL database %s and user %s" %(dspamdb, dspamuser)
    try:    
        cur.execute("""CREATE USER %s WITH PASSWORD '%s' NOCREATEDB NOCREATEUSER; CREATE DATABASE %s; GRANT ALL PRIVILEGES ON DATABASE %s TO %s; GRANT ALL PRIVILEGES ON SCHEMA public TO %s; UPDATE pg_database SET datdba=(SELECT usesysid FROM pg_shadow WHERE usename='%s') WHERE datname='%s';""" %(dspamuser, dspampass, dspamdb, dspamdb, dspamuser, dspamuser, dspamuser, dspamdb))
    except Exception, e:
        print "ERROR:", e
        sys.exit(2)

def executeScript(dspamuser, dspampass, dspamdb, hostnameip, script):
    print "Executing script %s for %s (u: %s p: %s)" %(script, dspamdb, dspamuser, dspampass)
    try:
        con = psycopg2.connect(database=dspamdb, user=dspamuser, password=dspampass, host=hostnameip)
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT) # Needed for plpgsql
        cur = con.cursor()
    except Exception, e:
        print "ERROR: ",e
        sys.exit(3)
    try:
        f = open(script, 'r')
        script = f.read()
        f.close()
    except Exception, e:
        print "ERROR: ",e
        sys.exit(4)
    try:
        cur.execute(script)
    except Exception, e:
        print "ERROR: ",e
        sys.exit(5)
    con.close()
    
def connect(hostnameip):
    connectuser = 'postgres'
    connectpass = ''
    notconnected = 1
    while(notconnected):
        try:    
            con = psycopg2.connect(database='template1', user=connectuser, password=connectpass, host=hostnameip)
            notconnected = 0
        except psycopg2.OperationalError, e:
            if 'no password' in e[0]:
                print "ERROR: Could not connect to database with default user 'postgres' an no password"
            else:
                print e
            connectuser = inputStuff('Connecting DB-User (default: postgres)', 'postgres', empty = False)
            connectpass = inputStuff('DB-User Password')
        except Exception, e:
            print e
            sys.exit(2)
    return con

if __name__ == '__main__':
    main()