#!/usr/bin/perl -w
#
# Postgresql wrapper script
# 
# (c)2006 Max Baker
# See LICENSE in netdisco script
#
# $Id: pg,v 1.7 2009/01/31 02:06:37 jeneric Exp $

use FindBin;
use Compress::Zlib;

$|=1;
$EXIT=0;

BEGIN {
    # Use this begin block to allow dynamic lib directory
    #   to be specified as an argument.
    use Getopt::Long;
    GetOptions(\%ARGS,'c|configfile=s',
                      'b|batch',
                      'u|dbuser=s',
                      'd|homedir=s',
                      'p|psql=s',
                      'i|init',
                      'back',
                      'h|help',
              );

    $DefaultDir = "/usr/local/netdisco";
    # add more possible locations here
    foreach $poss ($FindBin::Bin . "/..", "..") {
	if (-f $poss . "/netdisco.pm") {
	    $DefaultDir = $poss;
	}
    }
    $Dir        = $ARGS{d} || $DefaultDir;
    $ConfigFile = $ARGS{c} || "$Dir/netdisco.conf";
    $Psql       = $ARGS{p} || 'psql';
    eval "use lib '$Dir';";
    die "use lib '$Dir': $@." if ($@);
}

my $OS = $^O;

if ($OS =~ /^MSWin/i) {
    $OS = 'WINDOWS';
}
else {
    $OS = 'UNIX';
}

eval "use Win32::NetAdmin qw(GetUsers GroupIsMember
                             UserGetAttributes UserSetAttributes
                             FILTER_NORMAL_ACCOUNT)";

my $split = ($OS eq 'WINDOWS') ? ';' : ':';
my @dirList = split /$split/, $ENV{PATH};
my $Path;
foreach $dir (@dirList) {
  if ( -e "$dir/psql"  or -e "$dir/psql.exe") {
    if (($OS eq 'WINDOWS') and ($dir !~ /\\$/)) {
        $Path = "$dir" . "\\";
    }
    elsif (($OS eq 'UNIX') and ($dir !~ /\/$/)) {
        $Path = "$dir" . "/";
    }
    else {
     $Path = $dir;
    }
    last;
  }
}

unless ($Path or $Psql ne 'psql') {
    print "Unable to determine psql cli executable, define with -p option\n\n";
    die &usage
}

if ($Psql =~ /(.+)psql(\.exe)?$/){
    $Path=$1;
}

die &usage if defined $ARGS{h};

print "$0 - Netdisco Postgres Wrapper.\n";

# Load netdisco library
eval "use netdisco;";
if ($@) {
    if (-d $Dir && ! -x _) {
        warn "$Dir exists but is not readable by you.\n";
    }
    if (-f $Dir . "/netdisco.pm" && ! -r _) {
        warn "$Dir/netdisco.pm exists but is not readable by you.\n";
    }
    die "Could not find netdisco.pm: $@.  Please set option -d to point to
    your netdisco directory.\n";
}

# Load Config File
unless (-r $ConfigFile) { 
    die "The netdisco config file specified, $ConfigFile, does not exist or
    is not readable.\n";
}
netdisco::config($ConfigFile);

# Get Database Configuration
$DB='netdisco';
if ($CONFIG{db_Pg} =~ /dbname=([^:]+)/) {
	$DB = $1;
}
$DB_USER = $CONFIG{db_Pg_user};
$DB_PW   = $CONFIG{db_Pg_pw};

unless (defined $DB_USER) {
    die "Please set a database user in netdisco.conf under db_Pg_user.\n";
    }
unless (defined $DB_PW)   {
    die "Please set a database password in netdisco.conf under db_Pg_pw.\n";
    }

# Init all tables, run a command, or launch the shell
$File = shift @ARGV || '';

if ($ARGS{i}) {
    init();
} elsif ($ARGS{back}) {
    back();
} elsif (-r $File) {
    run($File);
} else {
    interactive();
}

exit($EXIT);

sub run {
    my $file = shift;

    print "Loading $file into $Psql\n";

    # Prompt for destructoness
    unless ($file =~ /^upgrade/ or defined $ARGS{b}) {
        if ($file =~ /^db_back-(\d\d\d\d)(\d\d)(\d\d)/) {
            $msg = "This will restore the database backup taken on $2-$3-$1";
        } else {
            $msg = "This will nuke all the data in the table in $file";
        }
        print "$msg.  You Sure (y/n)?";
        my $reply = <STDIN>;
        chomp($reply);
        unless ($reply =~ /^y/i) {
            print "Never Mind.\n";
            exit(1);
        }
    }
    print "If prompted enter '$DB_PW' as password.\n";
    psql("-e -f $file -U $DB_USER $DB");

}

sub interactive {
    print "If prompted enter '$DB_PW' as password.\n";
    psql("-e -U $DB_USER $DB");
}

sub init {
    $CREATEUSER = "\"$Path"."psql\" -a template1 --command \\\"CREATE USER $DB_USER PASSWORD '$DB_PW' NOCREATEDB NOCREATEUSER\\\"";
    $CREATEDB   = "\"$Path"."createdb\" $DB";

    open (ALL,">pg_all.input") or die "Can't create pg_all.input for output. $!\n";

    # Cat all tables into one big file
    opendir(SQLDIR, "$Dir/sql") or die "Can't open $Dir/sql: $!";
    while( defined (my $f = readdir SQLDIR) ) {
        next if $f =~ /^\.\.?$/ or $f =~ /^upgrade-.*\.sql$/;
	next unless $f =~ /\.sql$/;
        print "[ $f ]\n";
        open (F, "<$Dir/sql/$f") or die "Can't read $f. $!\n";
        while (<F>) {
            print ALL;
        }
        close F;
    }

    close ALL or die "Can't write pg_all.input. $!\n";

    # Check if Root
    unless ($> == 0){ warn "You probably should be running this as root.\n"; }
    print "\n";

    # Determine the System Database User

    $SYS_DB_USER=$ARGS{u};

    # Windows first
    if ($OS eq 'WINDOWS') {

        my %hash;
        GetUsers("", FILTER_NORMAL_ACCOUNT , \%hash)
            or die "GetUsers() failed: $^E";

        foreach (keys %hash) {
            if ($_ =~ /^(postgres|pgsql)/i){
                $SYS_DB_USER = $_;
                last;
            }

            my ($password, $passwordAge, $privilege,
                $homeDir, $comment, $flags, $scriptPath);
            UserGetAttributes("", $_, $password, $passwordAge, $privilege,
                              $homeDir, $comment, $flags, $scriptPath)
                or die "UserGetAttributes() failed: $^E";

            if ($comment =~ /^(postgres|pgsql)/i){
                $SYS_DB_USER = $_;
                last;
            }
            next;
        }
    }

    # UNIX
    if ($OS eq 'UNIX') {
        foreach $user qw(postgres pgsql) {
            if (($SYS_DB_UID = getpwnam($user))) {
                $SYS_DB_USER=$user;
                last;
            }
        }
    }

    # Wrap calls in 'su' if we have sys db user
    if (defined $SYS_DB_USER and $OS eq 'UNIX' and $> ne $SYS_DB_UID){
        print "Running database commands as $SYS_DB_USER user.\n";
        $CREATEUSER = "su $SYS_DB_USER -c \"$CREATEUSER\"";
        $CREATEDB   = "su $SYS_DB_USER -c \"$CREATEDB\"";
    } elsif (defined $SYS_DB_USER and $OS eq 'WINDOWS') {
        print "Running database commands as $SYS_DB_USER user.\n";
	print "When prompted enter the $SYS_DB_USER user password.\n";
        $CREATEUSER = "\"$Path"."psql\" -U $SYS_DB_USER -W -a template1 --command \"CREATE USER $DB PASSWORD '$DB_PW' NOCREATEDB NOCREATEUSER\"";
        $CREATEDB   = "\"$Path"."createdb\" -U $SYS_DB_USER -W $DB";	
    } else {
        print "Could not determine which user Postgresql is running under. Will try as current user.\n";
        print "Use option -u to specify the unix user the database runs under.\n";
    }

    print "Creating User $DB_USER.\n";
    print "Enter the $SYS_DB_USER user password.\n" if ($OS eq 'WINDOWS') ;
    system($CREATEUSER) and bug("createuser did not return success.\n"); 

    print "Creating Database $DB.\n";
    print "Enter the $SYS_DB_USER user password.\n" if ($OS eq 'WINDOWS');
    system($CREATEDB) and bug("createdb did not return success.\n"); 

    print "Creating tables and indicies.\n";
    print "If prompted enter '$DB_PW' as password.\n";

    psql("-e -f pg_all.input -U $DB_USER $DB");
}

sub back {
    # TODO - Change this to pureperl
    my ($sec,$min,$hour,$mday,$mon,$year,$wday,$yday,$isdst) = localtime();
    my $DATE = sprintf("%d%02d%02d",$year+1900,$mon+1,$mday);
    my $FILE="db_back-${DATE}.sql";
    
    print "Backing up $DB to $FILE\n";
    $EXIT=system("\"$Path"."pg_dump\" -a -d -D -O -U ${DB_USER} ${DB} > ${FILE}");
    open(OLD, "< $FILE") or die "Can't open $FILE: $!\n";
       my $gz = gzopen("$FILE.gz", "wb") or die "Can't open $FILE.gz: $gzerrno\n";
       while (<OLD>) {
           $gz->gzwrite($_) 
              or die "Error writing: $gzerrno\n" ;
       }
       $gz->gzclose ;
       close(OLD)      or die "Can't close $FILE: $!\n";
       unlink($FILE)   or die "Can't delete $FILE: $!\n";    
    chmod (0600, "$FILE.gz");
}

sub psql {
    my $cmd = shift;
    $EXIT = system("\"$Psql\" $cmd");
}


sub bug {
    my $b = shift;
    print $b if defined $b;
    $EXIT++;
}

sub usage {
    return <<"end_usage";
pg - Netdisco Postgres shell

USAGE
    $0 [options] [file.sql]

Specify a SQL file to be run batch mode, or 
Don't give any arguments to launch the psql shell

OPTIONS
   -b                        -- Batch mode, no override protection
   -d /usr/local/netdisco    -- Path to netdisco.pm
   -c /path/to/netdisco.conf -- Netdisco Config file to use
   -u pgsql                  -- Database UNIX user
   -p /path/to/psql          -- psql cli executable

   -i                        -- Init all tables (Will delete all existing data!)
   --back                    -- Create Backup of database

end_usage
}
