#!/usr/bin/perl -w

#  Update a PostgreSQL <pg_integrity.conf> file.
#
#  This program 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.

use strict;
use warnings;
use PgCommon;
use File::Temp qw(tempfile);
use Getopt::Long;
use POSIX qw(lchown);

# untaint environment
$ENV{'PATH'} = '/bin:/usr/bin';
delete @ENV{'IFS', 'CDPATH', 'ENV', 'BASH_ENV'};

# global variables
my ($version, $cluster, $command);
my (%info);

# do not trip over cwd not being accessible to postgres superuser
chdir '/';

# Update pg_integrity.conf whole section or only one line
# Arguments: <data> <target> <this is only one line>
sub update_integrity_conf
{

	my @data = @{$_[0]};
	my $target = $_[1];
	my $only_one_line = $_[2];
	my $filE = PgCommon::cluster_data_directory($version, $cluster)."/pg_integrity.conf";
	my $skip_section_copy = 0;
	my @sectionS = ('^\[FILE\]', '^\[PARAMETERS\]', '^\[SQL\]');

	# From file
	open O, $filE or error "open $filE: $!";
	# To file
	open N, ">$filE.new" or error "open $filE.new: $!";
	# Start copy
	while (my $line = <O>)
	{
		# Check target string
		if ($line =~ /$target/)
		{
			# One line or whole section
			unless ($only_one_line)
			{
				# Copy section name
				print N $line;
				$skip_section_copy = 1;
			}
			# Insert new line[s]
			foreach my $value ( @data )
			{
				 print N $value; 
			}
			# Blank line at the end of section
			unless ($only_one_line)
			{
				print N "\n";
			}
			next;
		}

		if ($skip_section_copy)
		{
			# Skip copying the old lines until the start of the new section

			if ($line =~ /$sectionS[0]|$sectionS[1]|$sectionS[2]/) {
				$skip_section_copy = 0;
			} else {
				next;
			}
		}

		# Just copy line from one file to another
		print N $line;
	}
	close O;
	close N;

	lchown $info{'owneruid'}, $info{'ownergid'}, "$filE.new";
	chmod 0640, "$filE.new";
	rename "$filE.new", $filE or error "rename: $!";
}


# Get [FILE] data
# Arguments: <psql> <socket> <port> <files>
sub get_files_data{
	my ($psql, $socket, $port, @files) = @_;
	my @argv = ($psql, '-h', $socket, ' -p', $port, ' -F\' : \'', '-d', 'template1');
	push @argv, (' -AXtc "');
	for my $file (@files)
	{
		push @argv, ('SELECT \''.$file.'\', \'\"\'||get_file_hash(\''.$file.'\')||\'\"\'');
		push @argv, (' UNION ') unless ($file eq $files[$#files]);
	}
	push @argv, ( ';"');
    return qx(@argv);
}

# Get [PARAMETERS] data
# Arguments: <psql> <socket> <port>
sub get_parameters_data{
	my ($psql, $socket, $port) = @_;
	my @argv = ($psql, '-h', $socket, ' -p', $port, ' -F\' : \'', '-d', 'template1');
	push @argv, (' -qc "set lc_messages = \'C\'"');
	push @argv, (' -AXtc "');
	push @argv, ('SELECT name, \'\"\'||setting||\'\"\' FROM pg_settings');
	push @argv, (' WHERE ');
	push @argv, ('category not in (
\'Client Connection Defaults \/ Statement Behavior\',
\'Client Connection Defaults \/ Locale and Formatting\',
\'Client Connection Defaults \/ Shared Library Preloading\',
\'Client Connection Defaults \/ Other Defaults\')');
	push @argv, (' AND ');
	push @argv, ('  name not in (
\'ac_capable_chmac\',
\'ac_capable_ignmaccat\',
\'ac_capable_ignmaclvl\',
\'ac_capable_mac_readsearch\',
\'ac_capable_setmac\') ');
	push @argv, ( ';"');
    return qx(@argv);
}

# Get [SQL] data
# Arguments: <database> <psql> <socket> <port> [<oid>]
sub get_procedure_data{
	my ($database, $psql, $socket, $port,  $oid) = @_;
	my @argv = ($psql, ' -h ', $socket, ' -p ', $port, ' -F\' : \'', ' -d ', $database);
	push @argv, (' -AXtc', ' "SELECT pr.proname,
\'\"\'||get_procedure_hash(pr.proname, nm.nspname, pg_get_function_arguments(pr.oid))||\'\"\',
\'\"\'||nm.nspname||\'\"\',
\'\"\'||regexp_replace(pg_get_function_arguments(pr.oid),\'\"\', \'\"\"\', \'gi\')||\'\"\',
\'\"\'||current_database()||\'\"\'
FROM pg_proc AS pr
JOIN pg_namespace AS nm ON pr.pronamespace=nm.oid');
	push @argv, (' WHERE pr.prokind != \'a\'');
	push @argv, (' AND pr.oid = ', $oid) if ($oid);
	push @argv, ( ';"');
    return qx(@argv);
}

# Arguments: <text array> <number of columns>
sub format_width{
	my @lines = @{$_[0]};
	my $n = $_[1];;

	my @colwidth =  (1) x $n;
	foreach my $line (@lines) {
	    for (my $i = 0; $i < @$line - 0; $i++) {
	        my $len = length @$line[$i];
	        $colwidth[$i] = $len+0 if ($len > $colwidth[$i]);
	    }
	}
	my $fmtstring = join ' ', map { "%-${_}s" } @colwidth;
    return $fmtstring;
}

# Print state
# Arguments: <text array> <format>
sub print_state{
	my @lines = @{$_[0]};
	my $fmtstring = $_[1];;

	my $fmtstringH = join ' ', map { "\033[4m${_}\033[0m" } split ' ', $fmtstring;
	my $head = 1;

	foreach my $line (@lines) {
		if ($head) {
			printf "$fmtstringH\n", @$line;
			$head = 0;
		} else {
			printf "$fmtstring\n", @$line;
		}
	}
}

#########################
#                       #
# Execution starts here #
#                       #
#########################


if ($#ARGV < 1) {
    print "Usage: $0 [OPTIONS] <version> <cluster name> <command>
Commands:
	check - integrity check
	refill { --file | --parameters | --sql } - rewrite section of integrity.conf file [<all sections> is set by default]
	clear | clean { --file | --parameters | --sql } - rewrite section of integrity.conf file [<all sections> is set by default]
	stop | off - stops integrity autochecking
	start | on { --level=| 'LOG' | 'LIMIT_CONN' | 'PANIC' --delay=<>} - starts integrity autochecking [--level='LOG' is set by default]
Examples:
	#pg_integrity 15 main check
	#pg_integrity 15 main refill
	#pg_integrity 15 main refill --file --sql
	#pg_integrity 15 main on --delay=30
	#pg_integrity 15 main clear --parameters
\n";
    exit 1;
}

($version) = $ARGV[0] =~ /^(\d+\.?\d+)$/;
($cluster) = $ARGV[1] =~ /^([-.\w]+)$/;
if ($#ARGV == 1) {
	($command) = 'state';
}else{
	($command) = $ARGV[2] =~ /^([-.\w]+)$/;
}

error 'specified cluster does not exist' unless cluster_exists $version, $cluster;
%info = cluster_info ($version, $cluster);
validate_cluster_owner \%info;
error 'cluster is disabled' if $info{'start'} eq 'disabled';


my $socket = get_cluster_socketdir $version, $cluster;
my $owner = getpwuid $info{'owneruid'};
error 'could not get name of cluster owner' unless $owner;


change_ugid $info{'owneruid'}, $info{'ownergid'};
my $psql = get_program_path 'psql';

# command line parameters
my $sec_file = 0;
my $sec_parameters = 0;
my $sec_sql = 0;
my $sec_all = 0;
my $response_level = 'LOG';
my $naptime = PgCommon::get_conf_value($version, $cluster, 'postgresql.conf', 'ac_integrity_naptime');


if ($command eq 'refill')
{
	if (!$info{'running'})
	{
		print "Сan't send a request. Cluster is not running\n";
		exit 1;
	}

	GetOptions ('file' => \$sec_file,
				'parameters' => \$sec_parameters,
				'sql' => \$sec_sql,
	) or exit 1;
	$sec_all = 1 unless ($sec_file | $sec_parameters | $sec_sql);

	# Update [FILE] section
	if ($sec_file || $sec_all)
	{
		error 'could not get files data' unless 
		my @files_data = get_files_data($psql, $socket, $info{'port'}, ("pg_ident.conf",
																		"pg_hba.conf",
																		"pg_ctl.conf",
																		"start.conf",
																		"environment",
																		"pg_audit.conf",
																		"postgresql.conf"));

		update_integrity_conf(\@files_data, '\[FILE\]');
		print "The [FILE] section has been updated\n";

	}

	# Update [PARAMETERS] section
	if ($sec_parameters || $sec_all)
	{
		error 'could not get parameters data' unless 
		my @parameters_data = get_parameters_data($psql, $socket, $info{'port'});
		update_integrity_conf(\@parameters_data, '\[PARAMETERS\]');
		print "The [PARAMETERS] section has been updated\n";
	}

	# Update [SQL] section
	if ($sec_sql || $sec_all)
	{
		# Get list of databases
		error 'could not get databases list' unless 
		my $databases = qx($psql -h $socket -p $info{'port'} -d template1 -AXtc "SELECT datname FROM pg_database");
		my @proc_data;
		for my $datname ($databases =~ /(.+)/g) {
			next if $datname eq 'template0';
			@proc_data = (@proc_data, get_procedure_data($datname, $psql, $socket, $info{'port'}));
		}
		error 'could not get procedure data' unless @proc_data;
		update_integrity_conf(\@proc_data, '\[SQL\]');
		print "The [SQL] section has been updated\n";
	}
}elsif($command eq 'stop' || $command eq 'off')
{
	change_ugid $info{'owneruid'}, $info{'ownergid'};
	PgCommon::set_conf_value $version, $cluster, 'postgresql.conf', 'ac_integrity_check', 'NONE';
	if ($info{'running'}) {
# 	    my @argv = system ('pg_ctlcluster', $version, $cluster, 'restart');
	    system 'systemctl', 'restart', "postgresql\@$version-$cluster";
	    exit $? >> 8;
# 	    error "Could not restart cluster" if system @argv;
	}
	print "Integrity is turned off\n";
}elsif($command eq 'start' || $command eq 'on')
{
	change_ugid $info{'owneruid'}, $info{'ownergid'};
	GetOptions ('level=s' => \$response_level,
				'delay=s' => \$naptime,
	) or exit 1;
	PgCommon::set_conf_value $version, $cluster, 'postgresql.conf', 'ac_integrity_check', $response_level;
	PgCommon::set_conf_value $version, $cluster, 'postgresql.conf', 'ac_integrity_naptime', $naptime if $naptime;
	if ($info{'running'}) {
		error 'could not get postgresql.conf file hash' unless 
		my @new_postgresql_hash = get_files_data($psql, $socket, $info{'port'}, "postgresql.conf");
		update_integrity_conf(\@new_postgresql_hash, '^postgresql.conf :', 1);

# 	    my @argv = system ('pg_ctlcluster', $version, $cluster, 'restart');
	    system 'systemctl', 'restart', "postgresql\@$version-$cluster";
	    exit $? >> 8;
# 	    error "Could not restart cluster" if system @argv;
	}
	print "Integrity is turned on\n";
}elsif($command eq 'clear' || $command eq 'clean')
{
	GetOptions ('file' => \$sec_file,
				'parameters' => \$sec_parameters,
				'sql' => \$sec_sql,
	) or exit 1;
	$sec_all = 1 unless ($sec_file | $sec_parameters | $sec_sql);

	# Clear [SQL] section
	if ($sec_sql || $sec_all)
	{
		update_integrity_conf([], '\[SQL\]');
		print "The [SQL] section has been cleared\n";
	}

	# Clear [FILE] section
	if ($sec_file || $sec_all)
	{
		update_integrity_conf([], '\[FILE\]');
		print "The [FILE] section has been cleared\n";
	}

	# Clear [PARAMETERS] section
	if ($sec_parameters || $sec_all)
	{
		update_integrity_conf([],, '\[PARAMETERS\]');
		print "The [PARAMETERS] section has been cleared\n";
	}
}elsif($command eq 'check' )
{
	if (!$info{'running'})
	{
		print "Сan't send a request. Cluster is not running\n";
		exit 1;
	}
	my $not_ok_exit = 0;
	# Get list of databases
	error 'could not get databases list' unless 
	my $databases = qx($psql -h $socket -p $info{'port'} -d template1 -AXtc "SELECT datname FROM pg_database");
	my @result = [("Database", "Result")];
	for my $datname ($databases =~ /(.+)/g) {
		next if $datname eq 'template0';
		my $query = qx($psql -h $socket -p $info{'port'} -d $datname -AXtc "SELECT get_integrity_state()");
		$not_ok_exit = 1 unless ($query =~ /successfully/);
		chop($query);
		push @result, [$datname, $query];
	}

	print "\n";
	my $head = 1;
	my $fmtstring = format_width(\@result, 2);
	foreach my $line (@result) {
		if ($head) {
			printf "$fmtstring\n", @$line;
			$head = 0;
		} else {
			printf "\033[%dm$fmtstring\033[0m\n",(@$line[1] =~ /successfully/ ? 32 : 31),  @$line;
		}
	}
	exit $not_ok_exit;

}elsif($command eq 'state' )
{
	if (!$info{'running'})
	{
		print "Сan't send a request. Cluster is not running\n";
		exit 1;
	}
	my @dbs_result;
	my @itg_params;
	# Get list of databases
	error 'could not get databases list' unless 
	my $databases = qx($psql -h $socket -p $info{'port'} -d template1 -AXtc "SELECT datname FROM pg_database");

	push @dbs_result, ["Database", "Result", "Date", "Time"];
	for my $datname ($databases =~ /(.+)/g) {
		next if $datname eq 'template0';
		my $querry = qx($psql -h $socket -p $info{'port'} -d $datname -AXtc "SELECT COALESCE(last_integrity_check_result, \'-\'), COALESCE(DATE(last_integrity_check_time)::varchar(10),\'-\'), COALESCE(date_trunc(\'second\', (last_integrity_check_time)::time)::varchar(8), \'-\') FROM pg_stat_database WHERE datname=\'$datname\'");
		$querry =~ s/[\r\n]+$//;
		push @dbs_result, [$datname , split /\|/, $querry];
	}
	print "\n";
	print_state(\@dbs_result, format_width(\@dbs_result, 4));

	print "\n";

	push @itg_params, ["Auto", "Response", "Delay", "MaxProcesses"];

	error 'could not checker parameter' unless 
	my $itg_response = qx($psql -h $socket -p $info{'port'} -d template1 -AXtc "show ac_integrity_check");
	$itg_response=~ s/[\r\n]+$//;
	my $itg_auto = 'ON';
	if ($itg_response=~/NONE/)
	{
		$itg_response = 'LOG';
		$itg_auto = 'OFF';
	}
	error 'could not delay parameter' unless 
	my $itg_delay = qx($psql -h $socket -p $info{'port'} -d template1 -AXtc "show ac_integrity_naptime");
	$itg_delay=~ s/[\r\n]+$//;
	error 'could not max processes parameter' unless 
	my $itg_maxproc = qx($psql -h $socket -p $info{'port'} -d template1 -AXtc "show ac_integrity_processes");
	$itg_maxproc=~ s/[\r\n]+$//;
	push @itg_params, [($itg_auto, $itg_response, $itg_delay, $itg_maxproc)];

	print_state(\@itg_params, format_width(\@itg_params, 4));
}else
{
	error "Unknown command";
}

