#!/usr/bin/perl

# Modify the grub config file on the node to boot the specified kernel and initrd.
# This script is meant to be run on the node via xdsh -e.
# Currently requires that dns on the mn be configured and working to resolve the short node names.

use strict;
use Getopt::Long;
use Data::Dumper;
use Socket;

# Globals - these are set once and then only read.
my $HELP;
my $VERBOSE;
my $DRYRUN;
my $WAITTIME;
my $PROVMETHOD;
my $XCATNETBOOTTITLE = 'xCAT network boot kernel and initrd';

my $usage = sub {
   	my $exitcode = shift @_;
   	print "Usage: modifygrub [-?|-h|--help] [-v|--verbose] [--dryrun] [-w <waittime>] [-p <provmethod] <kernel-path> <initrd-path> <kernel-parms> <mn-ip>\n\n";
   	if (!$exitcode) {
   		print "Modify the grub config file on the node to boot the specified kernel and initrd.\n";
   	}
	exit $exitcode;
};


my $file = '/etc/os-release';
if (-f $file) { 
	#
	# SLES and RHEL also have /etc/os-release file, so actually need to open the file
	# and look for Ubuntu in the 'NAME=' lines before declaring Ubuntu
	#
	open my $info, $file or die "Could not open $file: $!";
	while( my $line = <$info>)  {
		if (index($line, 'NAME=') != -1) {
			if ($line =~ /Ubuntu/i) {
				die "This script does not support Ubuntu at this time.\n";
			}
		}
	}
	close $info;
}

# Process the cmd line args
Getopt::Long::Configure("bundling");
#Getopt::Long::Configure("pass_through");
Getopt::Long::Configure("no_pass_through");
if (!GetOptions('h|?|help'  => \$HELP, 'v|verbose' => \$VERBOSE, 'dryrun' => \$DRYRUN, 'w|waittime=s' => \$WAITTIME, 'p|provmethod=s' => \$PROVMETHOD)) { $usage->(1); }

if ($HELP) { $usage->(0); }
if (scalar(@ARGV) != 4) { $usage->(1); }
if (!defined($WAITTIME)) { $WAITTIME = 60; }	# seconds to wait after configuring the nic (to let the switch handle the state change)
my %args;
$args{kernelpath} = $ARGV[0];
$args{initrdpath} = $ARGV[1];
$args{kernelparms} = $ARGV[2];
$args{mnip} = $ARGV[3];

addKernelParms(\%args);		# replace and add some parms to args{kernelparms}
updateGrub(\%args);			# update the grub config with an entry filled with the info in args
	
exit(0);


# Add ip and net info to the kernel parms.  Modifies the kernelparms value of the args hash passed in.
sub addKernelParms {
	my $args = shift @_;

	# replace '!myipfn!' with the mn ip
	my $mnip = $args->{mnip};
	$args->{kernelparms} =~ s/!myipfn!/$mnip/g;

	# replace <nodename> with the nodename
	my $nodename = $ENV{NODE};			# this env var is set by xdsh
	$args->{kernelparms} =~ s/<nodename>/$nodename/g;
	# get node ip and add it to the kernel parms
	my ($nic, $ip, $netmask, $network, $broadcast, $gateway, $mac) = getNodeIpInfo($args);
	if (!$ip) { die "Error: could not find the NIC that would connect to the xCAT mgmt node's IP (".$args->{mnip}.").\n"; }

	#if (defined($PROVMETHOD) && $PROVMETHOD eq 'sysclone') {
	if ($args->{kernelpath} =~ m/genesis\.kernel\.x86_64/) {
		# genesis case, add additional parms for sysclone or nodeset shell
		my $bootif = $mac;
		$bootif =~ s/:/-/g;
		$bootif = "BOOTIF=01-$bootif";
		# DEVICE=eth0 IPADDR=10.0.0.99 NETMASK=255.255.255.0 NETWORK=10.0.0.0 BROADCAST=10.0.0.255 GATEWAY=10.0.0.1 GATEWAYDEV=eth0
		#todo:  should we also add ETHER_SLEEP=$WAITTIME textmode=1 dns=$mnip ?
		$args->{kernelparms} .= " $bootif IPADDR=$ip NETMASK=$netmask NETWORK=$network BROADCAST=$broadcast GATEWAY=$gateway HOSTNAME=$nodename DEVICE=$nic GATEWAYDEV=$nic";
	}
	else {			# scripted install (kickstart or autoyast)
		if ($args->{kernelparms} =~ m/autoyast=/) {			# sles
			# if site.managedaddressmode=static is set, it will put several of these kernel parms on there for us.  Do not duplicate in that case.
			if ($args->{kernelparms} !~ m/ hostip=/) { $args->{kernelparms} .= " hostip=$ip"; }
			if ($args->{kernelparms} !~ m/ netmask=/) { $args->{kernelparms} .= " netmask=$netmask"; }
			if ($args->{kernelparms} !~ m/ gateway=/) { $args->{kernelparms} .= " gateway=$gateway"; }
			if ($args->{kernelparms} !~ m/ hostname=/) { $args->{kernelparms} .= " hostname=$nodename"; }
			if ($args->{kernelparms} !~ m/ dns=/) { $args->{kernelparms} .= " dns=$mnip"; }
			# If they set installnic=mac (recommended), the netdevice will already be set to the mac.
			# If they explicitly set installnic=<nic>, then ksdevice will be set to that and we will not disturb it here.
			# Otherwise add netdevice set to the nic we calculated it should be.
			if ($args->{kernelparms} !~ m/ netdevice=/) { $args->{kernelparms} .= " netdevice=$nic"; }
			$args->{kernelparms} .= " netwait=$WAITTIME textmode=1";
	        # print Dumper($args->{kernelparms})
    	}
		elsif ($args->{kernelparms} =~ m/ks=/) {			# rhel/centos
			# See https://www.centos.org/docs/5/html/Installation_Guide-en-US/s1-kickstart2-startinginstall.html
			# and  http://fedoraproject.org/wiki/Anaconda/NetworkIssues for description of kickstart kernel parms
			# if site.managedaddressmode=static is set, it will put several of these kernel parms on there for us.  Do not duplicate in that case.
			if ($args->{kernelparms} !~ m/ ip=/) { $args->{kernelparms} .= " ip=$ip"; }
			if ($args->{kernelparms} !~ m/ netmask=/) { $args->{kernelparms} .= " netmask=$netmask"; }
			if ($args->{kernelparms} !~ m/ gateway=/) { $args->{kernelparms} .= " gateway=$gateway"; }
			if ($args->{kernelparms} !~ m/ hostname=/) { $args->{kernelparms} .= " hostname=$nodename"; }
			if ($args->{kernelparms} !~ m/ dns=/) { $args->{kernelparms} .= " dns=$mnip"; }
			# If they set installnic=mac (recommended), the ksdevice will already be set to the mac.
			# If they explicitly set installnic=<nic>, then ksdevice will be set to that and we will not disturb it here.
			# Otherwise ksdevice will be set to bootif, and we change it to the nic we calculated it should be.
			if ($args->{kernelparms} =~ m/ ksdevice=bootif/) { $args->{kernelparms} =~ s/ ksdevice=bootif/ ksdevice=$nic/; }
			elsif ($args->{kernelparms} !~ m/ ksdevice=/) { $args->{kernelparms} .= " ksdevice=$nic"; }
			$args->{kernelparms} .= " nicdelay=$WAITTIME linksleep=$WAITTIME textmode=1";
	        # print Dumper($args->{kernelparms})
    	}
    	else { die "Error: for scripted installs, only support SLES or RHEL/CentOS as the target OS.\n"; }
	}
}


# get this nodes nic, ip, netmask, gateway, and mac.  Returns them in a 5 element array.
sub getNodeIpInfo {
	my $args = shift @_;
	my ($ipprefix) = $args->{mnip}=~m/^(\d+)\./;		#todo: this is a hack, just using the 1st octet of the mn ip addr
	verbose("using IP prefix $ipprefix");

	# parse ip addr show output, looking for ipprefix, to determine nic, ip, mac
	my @output = runcmd("/sbin/ip addr show");
	my ($nic, $mac, $ipandmask);
	foreach my $line (@output) {
		my ($nictmp, $mactmp, $iptmp);
		if (($nictmp) = $line=~m/^\d+:\s+(\S+): /) { $nic = $nictmp; }		# new stanza, remember it
		if (($mactmp) = $line=~m|^\s+link/ether\s+(\S+) |) { $mac = $mactmp; }		# got mac, remember it
		if (($iptmp) = $line=~m/^\s+inet\s+($ipprefix\S+) /) { $ipandmask = $iptmp; last; }		# got ip, we are done
	}
	if (!defined($ipandmask)) { die "Error: can't find a NIC with a prefix $ipprefix that communicates with".$args->{mnip}.".\n"; }
	my ($ip, $netmask, $network, $broadcast) = convertIpAndMask($ipandmask);

	# if the nic is a bonded nic (common on sl), then find the 1st real nic that is up that is part of it.
	# also find that real nics real mac
	my $realnic;
	if ($nic =~ /^bond/) {
		my @nics = grep(m/\s+master\s+$nic\s+/, @output);
		if (!scalar(@nics)) { die "Error: can't find the NICs that are part of $nic.\n"; }
		foreach my $line (@nics) {
			my ($nictmp, $state) = $line=~m/^\d+:\s+(\S+): .* state\s+(\S+)/;
			if (defined($nictmp) && defined($state) && $state eq 'UP') { $realnic = $nictmp; last; }		# got ip, we are done
		}
		if (!defined($realnic)) { die "Error: can't find a physical NIC that is up and part of $nic.\n"; }

		# now get the real mac of this real nic (when 2 nics are bonded, ip addr show displays one of the nics
		# macs for both nics and the bond).  So we have to depend on /proc/net/bonding/$bond instead.
		my @bondout = runcmd("cat /proc/net/bonding/$nic");
		my $foundnic;
		foreach my $line (@bondout) {
			my $mactmp;
			if ($line=~m/^Slave Interface:\s+$realnic/) { $foundnic = 1; }		# found the stanza for this nic, remember it
			if ($foundnic && (($mactmp) = $line=~m/^Permanent HW addr:\s+(\S+)/)) { $mac = $mactmp; last; }
		}
	}
	else { $realnic = $nic; }

	# centos/redhat seems to name the nic in a different order than sles on some svrs.
	# sles seems to name them in the same order as 'ip addr show' displays them, centos does not.
	# so if we are on centos right now, we need to count down to determine the number that sles
	# will give the nic that we have selected, because it is the sles naming that we care about,
	# because that is the initrd that will be running in the scripted install case.
	# This works similarly (at least in some case) when rhel is the target OS.
	# The preferred way is for the user to set installnic=mac, then we do not need to run this code.
	# For the sysclone case, genesis doxcat uses the mac to find the nic.
	if (isRedhat() && $args->{kernelparms} !~ m/ ksdevice=\S*:/ && $args->{kernelparms} !~ m/ netdevice=\S*:/) {
		my @nics = grep(m/^\d+:\s+eth/, @output);
		my $i = 0;
		foreach my $line (@nics) {
			my ($nictmp) = $line=~m/^\d+:\s+(\S+):/;
			if (defined($nictmp) && $nictmp eq $realnic) { $realnic = "eth$i"; last; }		# got ip, we are done
			$i++;
		}
		print "Determined that SLES/RHEL will call the install NIC $realnic (it has mac $mac)\n";
	}

	# finally, find the gateway
	my $gateway;
	my @output = runcmd("/sbin/ip route");
	# we are looking for a line like: 10.0.0.0/8 via 10.54.51.1 dev bond0
	my @networks = grep(m/ via .* $nic\s*$/, @output);
	if (scalar(@networks)) { ($gateway) = $networks[0]=~m/ via\s+(\S+)/; }
	else {
		# use the mn ip as a fall back
		$gateway = $args->{mnip};
		verbose("using xCAT mgmt node IP as the fall back gateway.");
	}

	verbose("IP info: realnic=$realnic, ip=$ip, netmask=$netmask, gateway=$gateway, mac=$mac");
	return ($realnic, $ip, $netmask, $network, $broadcast, $gateway, $mac);
}


# Convert an ip/mask in slash notation (like 10.1.1.1/26) to separate ip, netmask, network, and broadcast values,
# like: 10.1.1.1, 255.255.255.192, 10.1.1.0, 10.1.1.63
sub convertIpAndMask {
	my $ipandmask = shift @_;
	my ($ip, $masknum) = split('/', $ipandmask);

	# build the netmask
	my $nmbin = oct("0b" . '1' x $masknum . '0' x (32-$masknum));	# create a str like '1111100', then convert to binary
	my @nmarr=unpack('C4',pack('N',$nmbin));		# separate into the 4 octets
	my $netmask=join('.',@nmarr);			# put them together into the normal looking netmask

	# create binary form of ip
	my @iparr=split(/\./,$ip);
	my ( $ipbin ) = unpack('N', pack('C4',@iparr ) );

	# Calculate network address by logical AND operation of ip & netmask and convert network address to IP address format
	my $netbin = ( $ipbin & $nmbin );
	my @netarr=unpack('C4', pack('N',$netbin ) );
	my $network=join(".",@netarr);

	# Calculate broadcast address by inverting the netmask and adding it to the network address
	my $bcbin = ( $ipbin & $nmbin ) + ( ~ $nmbin );
	my @bcarr=unpack('C4', pack('N',$bcbin ) ) ;
	my $broadcast=join(".",@bcarr);

	return ($ip, $netmask, $network, $broadcast);
}


# not used - resolve the hostname to an ip addr
sub getipaddr {
	my $hostname = shift @_;
	my $packed_ip;
    $packed_ip = inet_aton($hostname);
    if (!$packed_ip) { return undef; }
    return inet_ntoa($packed_ip);
}


# Update the grub config file with a new stanza for booting our kernel and initrd
sub updateGrub {
	my $args = shift @_;

	# how we specify the path for the kernel and initrd is different on redhat and suse
	my $fileprefix;
	if (isRedhat()) { $fileprefix = '/'; }
	elsif (isSuse()) { $fileprefix = '/boot/'; }
	else { die "Error: currently only support red hat or suse distros.\n"; }

	# open the grub file and see if it is in there or if we have to add it
	my $grubfile = findGrubPath();
	verbose("reading $grubfile");
	open(FILE, $grubfile) || die "Error: can not open config file $grubfile for reading: $!\n";
	my @lines = <FILE>;
	close FILE;

	# this is the entry we want in the grub file
	my @rootlines = grep(/^\s+root\s+/, @lines);		# copy one of the existing root lines
	if (!scalar(@rootlines)) { die "Error: can't find an existing line for 'root' in the grub config file\n"; }
	my ($rootline) = $rootlines[0] =~ m/^\s*(.*?)\s*$/;
	my @entry = (
		"title $XCATNETBOOTTITLE\n",
		"\t$rootline\n",
		"\tkernel " . $fileprefix . $args->{kernelpath} . ' ' . $args->{kernelparms} . "\n",
		"\tinitrd " . $fileprefix . $args->{initrdpath} . "\n",
		);
	if ($DRYRUN) {
		print "Dry run:  would add this stanza to $grubfile:\n";
		foreach my $l (@entry) { print $l; }
		return;
	}

	my $needtowritefile = 1;
	if (grep(/^title\s+$XCATNETBOOTTITLE/, @lines)) { $needtowritefile = updateGrubEntry(\@lines, \@entry); }		# there is already an entry in there
	else { addGrubEntry (\@lines, \@entry); }

	# write the file with the new/updated xcat entry
	if ($needtowritefile) {
		verbose("updating $grubfile");
		open(FILE, '>', $grubfile) || die "Error: can not open config file $grubfile for writing: $!\n";
		print FILE @lines;
		close FILE;
	}
	else { print "Info: $grubfile did not need modifying. It was already up to date.\n"; }
}


# add our entry as the 1st one in the grub file
sub addGrubEntry {
	my ($lines, $entry) = @_;
	# find the index of the 1st stanza (it starts with 'title')
	my $i;
	for ($i=0; $i<scalar(@$lines); $i++) {
		if ($lines->[$i] =~ m/^title\s+/) { verbose('adding xcat entry before:'.$lines->[$i]); last; }		# found it
	}
	
	# splice the entry right before the i-th line (which may also be 1 past the end)
	splice(@$lines, $i, 0, @$entry);
}


# check the xcat entry in the grub file and see if it needs to be updated.  Return 1 if it does.
sub updateGrubEntry {
	my ($lines, $entry) = @_;
	#print Dumper($lines), Dumper($entry);
	# find the index of the xcat stanza
	my $i;
	for ($i=0; $i<scalar(@$lines); $i++) {
		if ($lines->[$i] =~ m/^title\s+$XCATNETBOOTTITLE/) { last; }		# found it
	}

	# compare the next few lines with the corresponding line in @$entries and replace if different
	my $replaced = 0;
	for (my $j=0; $j<scalar(@$entry); $j++) {
		#print "comparing:\n  ", $lines->[$i+$j], "\n  ", $entry->[$j], "\n";
		if ($lines->[$i+$j] ne $entry->[$j]) {		# this line was different
			$lines->[$i+$j] = $entry->[$j];
			$replaced = 1;
		}
	}
	
	return $replaced;
}


# depending on the distro, find the correct grub file and return its path
sub findGrubPath {
	# for rhel/centos it is /boot/grub/grub.conf, for sles it is /boot/grub/menu.lst
	my @paths = qw(/boot/grub/grub.conf /boot/grub/menu.lst);
	foreach my $p (@paths) {
		if (-f $p) { return $p; }
	}
	
	die "Error: Can't find grub config file.\n";

	#todo: support ubuntu: you add an executable file in /etc/grub.d named 06_xcatnetboot that prints out the
	#		entry to add.  Then run grub-mkconfig.
}


# Pring msg only if -v was specified
sub verbose { if ($VERBOSE) { print shift, "\n"; } }

# Check the distro we are running on
sub isSuse { return (-e '/etc/SuSE-release'); }
sub isRedhat { return (-e '/etc/redhat-release' || -e '/etc/centos-release' || -e '/etc/fedora-release'); }		# add chk for fedora



# Run a command.  If called in the context of return an array, it will capture the output
# of the cmd and return it.  Otherwise, it will display the output to stdout.
# If the cmd has a non-zero rc, this function will die with a msg.
sub runcmd
{
    my ($cmd) = @_;
    my $rc;

    $cmd .= ' 2>&1' ;
    verbose($cmd);

   	my @output;
   	if (wantarray) {
		@output = `$cmd`;
		$rc = $?;
	}
	else {
		system($cmd);
		$rc = $?;
	}

    if ($rc) {
        $rc = $rc >> 8;
        if ($rc > 0) { die "Error: rc $rc return from cmd: $cmd\n"; }
        else { die "Error: system error returned from cmd: $cmd\n"; }
    }
    elsif (wantarray) { return @output; }
}