#!/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} || '/usr/lib/netdisco';
    $ConfigFile = $ARGS{c} || "/etc/netdisco/netdisco.conf";
    $Psql       = $ARGS{p} || 'psql';
}

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/lib/netdisco      -- Path to netdisco "Home Directory" (can be ignored)
   -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
}

__END__

=head1 NAME npg

npg - NetDisco Postgres shell

=head1 SYNOPSIS

 npg [options] [file.sql]

=head1 DESCRIPTION

Specify an SQL file to be run in batch mode, or don't give any arguments to
launch the psql shell.

=head1 OPTIONS

 -b                        -- Batch mode, no override protection
 -d /usr/lib/netdisco      -- Path to netdisco "Home Directory" (can be ignored)
 -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

=head1 AUTHOR

Max Baker

=head1 COPYRIGHT & LICENSE

Copyright (c) 2006 Max Baker.

Redistribution and use in source and binary forms, with or without 
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright notice,
      this list of conditions and the following disclaimer in the documentation
      and/or other materials provided with the distribution.
    * Neither the name of the University of California, Santa Cruz nor the 
      names of its contributors may be used to endorse or promote products 
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

=cut

