#!/usr/bin/perl
## rrd_change.pl: Insert, append or rearrange data within RRD files
## Copyright (C) 2013 Robert Paschedag (robert dot paschedag at web dot de)
##
## This program 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 2
## of the License, or (at your option) any later version.
##
## 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.
##
## You should have received a copy of the GNU General Public License
## along with this program; if not, write to the Free Software
## Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#
# This script is inspired by the PNP4Nagios "rrd_convert.pl" script
#

use strict;
use warnings;
use File::Copy;
use Getopt::Long;

####################
####################
my $RRD_TOOL = '/usr/bin/rrdtool';
my $VERSION = '0.2';

my $rrd_ds_template = <<DSTEMPLATE;
	<ds>
		<name> ###NUM### </name>
		<type> ###TYPE### </type>
		<minimal_heartbeat> ###HEARTBEAT### </minimal_heartbeat>
		<min> NaN </min>
		<max> NaN </max>
		
		<!-- PDP Status -->
		<last_ds> UNKN </last_ds>
		<value> 0.0000000000e+00 </value>
		<unknown_sec> 0 </unknown_sec>
	</ds>

DSTEMPLATE

my $cdp_prep_template = <<CDPTEMPLATE;
	<ds>
	<primary_value> 0.0000000000e+00 </primary_value>
	<secondary_value> 0.0000000000e+00 </secondary_value>
	<value> NaN </value>
	<unknown_datapoints> 0 </unknown_datapoints>
	</ds>
CDPTEMPLATE

my $value_template = '<v> 0.0000000000e+00 </v>';

my ($opt_a, @opt_t, $opt_h, $opt_i, $opt_p, $opt_s, $rrdfile, $mode);

# insert_mode should be set here (in parse_options)
parse_options();

# define some file names
my $rrd_dump = $rrdfile . '.dump';
my $rrd_backup = $rrdfile . '.backup';
my $rrd_restore = $rrdfile. '.restore';

#get the last datasource within the rrd file
my $last_ds_num = get_last_ds_num($rrdfile);
my $ds_num;

# if we are in "insert mode", check the position the datasource should be inserted
if($mode eq 'i') {
	if($opt_p > $last_ds_num) {
		# set value to start one higher
		$ds_num = $last_ds_num + 1;
		# set "append mode"
		$mode = 'a';
		print "INFO: The positon to insert the datasource is higher than the maximum number of datasources within the RRD file\n",
			"Switching to 'append' mode.\n";
	} else {
		# we really insert.....set ds_num to start with
		$ds_num = $opt_p;
	}
} elsif($mode eq 'a') {
	# append
	$ds_num = $last_ds_num + 1;
} elsif($mode eq 's') {
	# check that the selected datasource numbers are present within the RRD file
	foreach(split(/,/, $opt_s)) {
		if($_ > $last_ds_num) {
			print "$0 error: One of the selected datasource numbers is higher than the number of datasources within the RRD file!\n";
			print "Aborting!\n";
			exit(1);
		}
	}
}

# check the values given for the new datasources (type and heartbeat)
# only needed in "insert" or "append" mode
my %NEW_DS = ();

if($mode eq 'a' || $mode eq 'i') {
	foreach (@opt_t) {
		if($_ !~ /^([A-Z]+):(\d+)$/) {
			print "$0 error: Wrong data format given: $_\n";
			print_help();
		}
		$NEW_DS{$ds_num}{TYPE} = $1;
		$NEW_DS{$ds_num}{HB} = $2;
	
		$ds_num++;
	}
}

# 1. dump the RRD
# 2. make a backup of the RRD
# 3. write new restore file with content changed
# 4. remove the dump file
# 5. restore the RRD with the restore file
# 6. remove the restore file

# 1.
rrd_dump($rrdfile, $rrd_dump);

# 2. 
rrd_backup($rrdfile, $rrd_backup);

my ($in_fd, $out_fd);
open($in_fd, '<', "$rrd_dump") || die "$0 error: Cannot open rrd dump file '$rrd_dump': $!\n";
open($out_fd, '>', "$rrd_restore") || die "$0 error: Cannot open rrd restore file for writing: $!\n";

# 3. 
if($mode eq 'i') {
	insert_ds(\%NEW_DS);
} elsif($mode eq 'a') {
	append_ds(\%NEW_DS);
} else {
	my @arr = split(/,/, $opt_s);
	select_ds(\@arr);
}

# close the files.
close($in_fd);
close($out_fd);

# 4.
unlink($rrd_dump);

# 5. 
rrd_restore($rrd_restore, $rrdfile);

# 6. 
unlink($rrd_restore);

print "Finished!\n";

# print warning information if we inserted....
if($mode eq 'i' || $mode eq 's') {
	print "\nATTENTION! DON'T FORGET TO CHANGE CUSTOM PNP4NAGIOS GRAPHING TEMPLATES USING *THIS* RRD FILE!!!\n";
}

exit(0);

#########################
#########################
sub rrd_backup {
	my ($rrd_file, $rrd_backup) = @_;
	move($rrd_file, $rrd_backup);
}

sub remove_file {
	my $dump = shift;
	
	unlink($dump);
}

sub rrd_dump {
	my ($rrd, $xml) = @_;
	system("$RRD_TOOL dump $rrd > $xml");
	
	if($?) {
		print "RRD Error: Cannot dump RRD file: $!\n";
		exit(1);
	}
}
	
sub rrd_restore {
	my ($xml, $rrd) = @_;
	
	system("$RRD_TOOL restore -f $xml $rrd");
	
	if($?) {
		print "RRD Error: Could not restore RRD: $!\n";
		exit(1);
	}
}
	
sub get_ds_string {
	#my ($num, $type, $heartbeat) = @_;
	my $ref = shift;
	
	my $out;
	
	# sort !!!!!!!
	foreach (sort(keys(%$ref))) {
		my $tmpout;
		($tmpout = $rrd_ds_template) =~ s/###NUM###/$_/;
		($tmpout = $tmpout) =~ s/###TYPE###/$ref->{$_}{TYPE}/;
		($tmpout = $tmpout) =~ s/###HEARTBEAT###/$ref->{$_}{HB}/;
		
		$out .= $tmpout;
	}
	
	return sprintf("%s", $out);
}

sub get_cdp_prep_ds_string {
	my $num_ds = shift;
	my $out;
	
	for (my $i = 0; $i < $num_ds; $i++) {
		$out .= $cdp_prep_template;
	}
	
	return sprintf("%s", $out);
}
	
sub get_new_values_string {
	my $num_vals = shift;
	
	my $out;
	
	for (my $i = 0; $i < $num_vals; $i++) {
		$out .= $value_template;
	}
	
	return sprintf("%s", $out);
}

sub get_last_ds_num {
	my $rrd = shift;
	
	my @DATA;
	
	@DATA = `$RRD_TOOL info $rrd`;
	
	if($?) {
		print "RRD Error: Cannot get info of RRD: $!\n";
		exit(1);
	}
	
	my $last_ds = 1;
	
	foreach (@DATA) {
		if(/^ds\[(\d+)\]\.type/) {
			$last_ds = $1;
		}
	}
	return $last_ds;
}

sub print_help {
	
	my $msg = <<EOF;
	
$0, version $VERSION
$0 <-a|-i|-s <arg>> <-f <file>> [-p <pos>] [-t <data>]

    -h, --help              this help page
	
    -a, --append            append datasources to the RRD (options '-i', '-a' and '-s' are exclusive!)
	
    -i, --insert            insert "one" datasource at specified position
                            (see option '-p'). If 'p' is not given, position '1' will be used!
						 
    -p, --pos <pos>         the position, on which the new datasource should be "inserted"
                            (must be used with option '-i'). Minumum value is 1!
						 
    -s, --select <p1[,px]>  comma separated list of datasources that should be "extracted"
                            and "inserted" into the "new" RRD file. Datasources *not* listed
                            will be removed (there must be no spaces between the values).

                            Example: RRD File with 10 datasources (1 to 10). We only want to keep
                                     the datasources 3, 6 and 9. All other will be lost
								  
                                     $0 -s 3,6,9 -f <RRDFILE>
								  
                                     The new RRD file will get the datasources with numbers 1 to 3, where
                                     1 is the old number 3
                                     2 is the old number 6 and
                                     3 is the old number 9
								  
                            You see, ordering is important here!
						 
    -f, --file <file>       the RRD file to process
	
    -t, --type <data>       the data type and heartbeat value
                            In "append" mode, this option can be specified multiple times.
                            In "insert" mode, only one datasource may be given.
						    In "select" mode, this has no effect.
                            Separate the TYPE and HEARTBEAT by a colon (:)
                            (e.g. -t GAUGE:300 -t COUNTER:600)
                            Use "rrdtool info" to get the current values of the datasouces
					  
EOF
	print $msg;
	exit(1);
}

sub parse_options {
	
	GetOptions(
		"h|help" => \$opt_h,
		"i|insert" => \$opt_i,
		"a|append" => \$opt_a,
		"s|select=s" => \$opt_s,
		"p|pos=i" => \$opt_p,
		"t|type=s" => \@opt_t,
		"f|file=s" => \$rrdfile
	);

	# print help
	defined($opt_h) && print_help();
	
	# print help if neither option '-i' nor '-a' nor '-s' given
	if(!defined($opt_a) && !defined($opt_i) && !defined($opt_s)) {
		print "$0 error: You must specify either option '-i' or option '-a' or option '-s'!\n";
		print_help();
	}

	# also print help if more than one option ('-i', '-a', or '-s') given
	if( (defined($opt_a) && defined($opt_i)) ||
		(defined($opt_a) && defined($opt_s)) ||
		(defined($opt_i) && defined($opt_s)) ) {
		print "$0 error: Options '-i', '-a' and '-s' are exclusive!\n";
		print_help();
	}
	
	# print help if option '-i' set but not option '-p'
	if(defined($opt_i) && !defined($opt_p)) {
		print "$0 error: You must specifiy option '-p' when trying to insert a datasource!\n";
		print_help();
	}
	
	# check if $opt_p is greater than 0
	if(defined($opt_p) && $opt_p < 1) {
		print "$0 error: Minimum number for option '-p' is 1!\n";
		print_help();
	}
	
	# set the mode
	$mode = 'i' if($opt_i);
	$mode = 'a' if($opt_a);
	$mode = 's' if($opt_s);
	
	if(!defined($rrdfile)) {
		print "$0 error: An RRD file must be specified with option '-f'!\n";
		print_help();
	}
	
	# print help if no data types have been given (option '-t')
	if(($mode eq 'a' || $mode eq 'i') && !@opt_t) {
		print "$0 error: option '-t' must be given when you insert or append datasources!\n";
		print_help();
	}

	# print help if option '-i' is defined AND the are more than 1 datasources given
	if(($mode eq 'i') && (scalar(@opt_t) > 1)) {
		print "$0 error: Only one datasource may be given in insert mode (option '-i')!\n";
		print_help();
	}
	
	# option select
	if($mode eq 's') {
		# first check per regex
		if($opt_s !~ /^\d+(,\d+)*$/) {
			print "$0 error: List of datasources to option '-s' must be comma separated!\n";
			print_help();
		}
		
		# stupid sanity check that every number given only once
		my %tmp_hash = ();
		foreach(split(/,/, $opt_s)) {
			if(exists($tmp_hash{$_})) {
				print "$0 error: Every datasource number to option '-s' may only be given once!\n";
				print_help();
			}
			$tmp_hash{$_} = 1;
		}
		
		# if we are here, then everything is ok
		%tmp_hash = ();
	}
}

sub append_ds {
	
	my $newds_ref = shift;
	
	while(<$in_fd>) {
		if(/Round Robin Archives/) {
			# insert the new datasources
			printf $out_fd "%s", get_ds_string($newds_ref);
			print $out_fd $_;
			next;
		}
	
		if(/<\/cdp_prep>/) {
			# insert cdp prep data sources
			printf $out_fd "%s", get_cdp_prep_ds_string(scalar(keys(%{$newds_ref})));
			print $out_fd $_;
			next;
		}
	
		if(/(.*<row>)/) {
			my $rowstart = $1;
			if(/(<v>.*)<\/row>$/) {
				my $old_vals = $1;
				my $new_vals_string = get_new_values_string(scalar(keys(%{$newds_ref})));
			
				print $out_fd $rowstart . $old_vals . $new_vals_string . '</row>' . "\n";
			}
			next;
		}
		print $out_fd $_;
	}
}

sub insert_ds {
	my $newds_ref = shift;
	
	my $num_ds_block = 0;
	my $in_archive_block = 0;
	my $in_cdprep_block = 0;
	
	my @tmparr = keys(%{$newds_ref});
	my $num_ds_insert = $tmparr[0];
	
	# start reading dump file

	while(<$in_fd>) {
	
		# change the ds number of all higher ds blocks
		if(m{<name>\s+(\d+)\s+</name>}) {
			my $num = $1;
			if($num >= $num_ds_insert) {
				$num++;
				print $out_fd '<name> ' . $num . ' </name>', "\n";
				next;
			}
		}
	
		if(/Round Robin Archives/) {
			$in_archive_block = 1;
		}
	
		if(/<cdp_prep>/) {
			$in_cdprep_block = 1;
			# start new count
			$num_ds_block = 0;
		}
	
		if(/<\/cdp_prep>/) {
			$in_cdprep_block = 0;
		}
	
		if(/<ds>/) {
			$num_ds_block++;
		
			# check, if we are in a cdprep block.
			if($in_cdprep_block) {
				# check, if this is the right ds block to insert
				if($num_ds_block == $num_ds_insert) {
					printf $out_fd "%s", get_cdp_prep_ds_string(1);
				}
			} else {
				if($num_ds_block == $num_ds_insert) {
					printf $out_fd "%s", get_ds_string($newds_ref);
				}
			}
		}
	
		if(/(.*<row>)/) {
			my $rowstart = $1;
			my @vals = m{(<v>[^<].*?</v>)}gc;
			my $tmpout;
			for(my $i = 0; $i < scalar(@vals); $i++) {
				if(($i+1) == $num_ds_insert) {
					$tmpout .= get_new_values_string(1);
				}
				$tmpout .= $vals[$i];
			}
			
			# print database record
			print $out_fd $rowstart . $tmpout . '</row>', "\n";
			
			# next line!
			next;
		}
		
		# now print the current line
		print $out_fd $_;
	}
}

sub select_ds {
	my $ds_arr_ref = shift;
	
	# example: 1,4,6,8,9
	# mapping from old to new number
	my %old_to_new = ();
	
	# the rrd ds numbers start at 1
	my $index = 0;
	foreach(@{$ds_arr_ref}) {
		$index++;
		if($_ > $last_ds_num) {
			print "$0 error: You cannot select a datasource number that is not within the rrd file!\n";
			print "Abort!\n";
			exit(1);
		}
		
		# add a mapping
		$old_to_new{$_} = $index;
	}
	
	# %old_to_new looks like this now
	# 1 -> 1
	# 4 -> 2
	# 6 -> 3
	# 8 -> 4
	# 9 -> 5
	
	# these are buffers for the data we read
	my %ds_buffer = ();
	my %cdp_prep_buffer = ();
	
	my $num_ds_block = 0;
	my $in_archive_block = 0;
	my $in_cdprep_block = 0;
	my $in_ds_block = 0;
	
	# start reading the xml file
	while(<$in_fd>) {
		
		if(/Round Robin Archives/) {
			# we are just at the end of the datasources definitions
			# print the definitions in new order (from the %ds_buffer)
			# the "new" names should have been changed
			foreach(@{$ds_arr_ref}) {
				print $out_fd $ds_buffer{$_};
			}
			
			# set flag
			$in_archive_block = 1;
			
			# clear the ds_buffer. we don't need it anymore
			%ds_buffer = ();
		}
	
		if(/<cdp_prep>/) {
			# set flag
			$in_cdprep_block = 1;
			# start new count
			$num_ds_block = 0;
		}
	
		if(/<\/cdp_prep>/) {
			# at the end of a cdp_prep block
			# print "ds" blocks in new order (from %cdp_prep_buffer)
			foreach(@{$ds_arr_ref}) {
				print $out_fd $cdp_prep_buffer{$_};
			}			
			
			# clear the buffer for the next block
			%cdp_prep_buffer = ();
			
			# set flag
			$in_cdprep_block = 0;
		}
		
		if(/<ds>/) {
			$num_ds_block++;
			
			# add the line to the buffer
			if($in_cdprep_block) {
				$cdp_prep_buffer{$num_ds_block} .= $_;
			} else {
				$ds_buffer{$num_ds_block} .= $_;
			}
			$in_ds_block = 1;
			
			# read next line!
			next;
		}
		
		if(/<\/ds>/) {
			# add the line to the buffer
			if($in_cdprep_block) {
				$cdp_prep_buffer{$num_ds_block} .= $_;
			} else {
				$ds_buffer{$num_ds_block} .= $_;
			}
			
			$in_ds_block = 0;
			# read the next line!
			next;
		}
		
		# change the "name" of the datasource to the new number
		# this also implies that we must add it to the %ds_buffer
		if(/<name>\s+(\d+)\s+<\/name>/) {
			if(exists($old_to_new{$1})) {
				$ds_buffer{$num_ds_block} .= '<name> ' . $old_to_new{$1} . ' </name>' . "\n";
			} else {
				$ds_buffer{$num_ds_block} .= $_;
			}
			
			# read next line
			next;
		}
		
		# if we are in a "ds" block, add the line to the buffer
		if($in_ds_block) {
			if($in_cdprep_block) {
				$cdp_prep_buffer{$num_ds_block} .= $_;
			} else {
				$ds_buffer{$num_ds_block} .= $_;
			}
			
			# read next line
			next;
		}
		
		# change "database"
		if(/(.*<row>)/) {
			my $rowstart = $1;
			my @vals = m{(<v>[^<].*?</v>)}gc;
			
			my $tmpout;
			# rearrange the values
			foreach(@{$ds_arr_ref}) {
				# datasources start counting from 1, array index from 0
				$tmpout .= $vals[$_ - 1];
			}
			
			# print the resulting line
			print $out_fd $rowstart . $tmpout . '</row>', "\n";
			
			# read next line
			next;
		}
			
		# this looks like a "normal" line we can print directly
		# everything else should have been check before
		print $out_fd $_;
	}
}
