#!/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; }
}

