#!/usr/bin/perl
#
# This script is a component of Warewulf,
# http://www.runlevelzero.net/greg/warewulf
#
#########################################################################
#
# Copyright (c) 2003, The Regents of the University of California,
# through
# Lawrence Berkeley National Laboratory (subject to receipt of any
# required approvals from the U.S. Dept. of Energy).  All rights
# reserved.
#
# 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.
#
# The GNU GPL Document can be found at:
# http://www.gnu.org/copyleft/gpl.html
#
#########################################################################
#
# Written and maintained by:
#       Greg Kurtzer, <gmkurtzer@lbl.gov>

use lib "/usr/lib/warewulf/", "/usr/lib64/warewulf/";
use Warewulf::PXE;
use Warewulf::Util;
use Warewulf::Config;
use Getopt::Long;
use File::Path;
use File::Basename;
use File::Copy;


$max_initrd_size = 15360;

my $usage = "USAGE: $0 [options] [kernel version]
  About:
   wwmkinitrd is responsiable for taking the template initial ram disk
   (found in /var/warewulf/wwinitrd) and turn it into the network
   bootable image used to boot the warewulf nodes. Any changes to
   /etc/warewulf/wwinitrd.config, or any of the files in the template FS
   should be followed with wwmkinitrd.

  Options:
    --kernel    Specify which kernel file that should be used.
                note: you still need to specify the kernel version!
    --moduledir If the modules are in an unstandard location
    --verbose   Run with extra output
    --help      Show this banner

  This tool is part of the Warewulf cluster distribution
     http://warewulf-cluster.org/
";

GetOptions(
   'moduledir:s'   => \$kernel_dir,
   'kernel:s'      => \$kernel,
   'help|h'        => \$help,
   'debug'         => \$debug,
   'verbose'       => \$verbose,
);

if ( $help ) {
   die $usage;
}

if ( $> != '0' ) {
   die "You need to run this as root!\n";
#   warn "You need to be root to make configuration changes!\n";
}

if ( $debug ) {
   $verbose = '1';
}

if ( $ARGV[0] ) {
   $kernel_version = $ARGV[0];
} else {
   chomp ($kernel_version = `uname -r`);
}

if ( ! -f "$kernel" and $kernel ) {
   die "ERROR: '--kernel=$kernel' doesn't exsist!\n";
}

if ( $kernel_dir ) {
   if ( ! -d "$kernel_dir" ) {
      die "Kernel module directory '$kernel_dir' doens't exsist!\n";
   }
   if ( ! -f "$kernel_dir/modules.dep" ) {
      die "Kernel module directory doesn't have modules.dep!\n";
   }
   if ( ! -f "$kernel_dir/modules.pcimap" ) {
      die "Kernel module directory doesn't have modules.pcimap!\n";
   }
   # Figure out what the root is:
   if ( $kernel_dir =~ /^(.*)\/lib\/modules\/.*/ ) {
      $module_root = $1;
   }
} else {
   if ( ! -d "/lib/modules/$kernel_version" ) {
      die "Kernel module directory '/lib/modules/$kernel_version' doens't exsist!\n";
   }
   $kernel_dir = "/lib/modules/$kernel_version";
   $module_root = "";
}


if ( ! -d "/tftpboot" ) {
   mkpath("/tftpboot");
}

sub module_deps {
   my ($module, @NULL) = @_;
   my $ret;
   my $eval;
   open(FILE, "<$kernel_dir/modules.dep");
   while(<FILE>) {
      chomp;
      $_ =~ /^(.+)(.)$/;
      if ( $2 eq '\\' ) {
         $eval .= "$1 ";
         next;
      } else {
         $eval .= "$_ ";
         ($mod, $deps) = split(/:\s*/, $eval);
         if ( "$module_root$mod" eq $module and $deps ) {
            $deps =~ s/([^\s]+)/$module_root$1 /g;
            $ret .= $deps;
            $debug and warn "Brining in module deps for $mod:\n   $deps\n";
         }
         $eval = ();
      }
   }
   return $ret;
}

%wwmkintird_conf = &read_conf("/etc/warewulf/wwinitrd.config");

$verbose and print "Building module lists...\n";

@modules = split(/\s+/, $wwmkintird_conf{KERNEL_MODULES});

foreach(@modules){
   $mods{$_} = '1';
}

open(MODS, "find $kernel_dir |");
while(<MODS>) {
   chomp;
   $module = basename("$_");
   if ( $module =~ /^(.+)(\.ko|\.o)(\.gz)?$/ ) {
      if ( $mods{$1} ) {
         $module_name{$1} = '1';
         $module_files{$_} = '1';
         $mod_deps = &module_deps($_);
         if ( $mod_deps ) {
            @deps = split(/\s+/, $mod_deps);
            foreach $dep (@deps){
               $module_files{$dep} = '1';
            }
         }
      }
   }
}

### Generate random location to build in
$random_string = &generate_random_string('16');

$verbose and print "building in /var/tmp/.wwinitrd-$random_string\n";

mkpath("/var/tmp/.wwinitrd-$random_string");

$verbose and print "creating loopback file system\n";
system("dd if=/dev/zero of=/var/tmp/.wwinitrd-$random_string.img bs=1024 count=$max_initrd_size >/dev/null 2>&1");
system("echo y | /sbin/mke2fs /var/tmp/.wwinitrd-$random_string.img >/dev/null 2>&1");
system("mount /var/tmp/.wwinitrd-$random_string.img /var/tmp/.wwinitrd-$random_string -o loop");


system("cp -rap /var/warewulf/wwinitrd/* /var/tmp/.wwinitrd-$random_string");
$verbose and print "Migrating kernel modules...";
foreach (keys %module_files) {
   $verbose and print ".";
   $dirname = dirname("$_");
   $dirname =~ s/^.*(\/lib\/modules\/.+)$/$1/;
   $basename = basename("$_");
   mkpath("/var/tmp/.wwinitrd-$random_string/$dirname");
   copy("$_", "/var/tmp/.wwinitrd-$random_string/$dirname/$basename");
   ($basename_ungz = $basename) =~ s/\.gz//;
   if ($basename ne $basename_ungz) {
       system("gunzip /var/tmp/.wwinitrd-$random_string/$dirname/$basename");
   }
   system("strip -pg /var/tmp/.wwinitrd-$random_string/$dirname/$basename_ungz");
   if ($basename ne $basename_ungz) {
       system("cd /var/tmp/.wwinitrd-$random_string/$dirname/; gzip $basename_ungz");
   }
}
$verbose and print "\n";
copy("$kernel_dir/modules.pcimap", "/var/tmp/.wwinitrd-$random_string/lib/modules/$kernel_version/modules.pcimap");
system("sed -e 's/\.gz//g' $kernel_dir/modules.dep > /var/tmp/.wwinitrd-$random_string/$kernel_dir/modules.dep");
system("find /var/tmp/.wwinitrd-$random_string/$kernel_dir/  -name *.gz -exec gunzip {} \\; 2>/dev/null");

if ( -f "/sbin/insmod.static" ) {
   $verbose and print "Using the hosts's static insmod\n";
   unlink("/var/tmp/.wwinitrd-$random_string/sbin/insmod");
   system("cp -ap /sbin/insmod.static /var/tmp/.wwinitrd-$random_string/sbin/insmod");
}

@load_modules = split(/\s+/, $wwmkintird_conf{LOAD_MODULES});
foreach $module ( @load_modules ) {
   if ( $module =~ /^(.+)=(.+)$/ ) {
      $line .= "$module ";
   } else {
      if ( $module_name{$module} ) {
         if ( $line ) {
            push (@load_mods, $line);
         }
         $line = "$module ";
         $load_mod .= "$module ";
      }
   }
}
if ( $line ) {
   push (@load_mods, $line);
}

$verbose and print "Building device tree\n";
mkpath("/var/tmp/.wwinitrd-$random_string/dev/pts");
@DEVS = split(/\s+/, $wwmkintird_conf{DEVS});
foreach $dev (@DEVS) {
   $debug and warn "   /dev/$dev\n";
   system("/bin/cp -ra /dev/$dev /var/tmp/.wwinitrd-$random_string/dev/$dev 2>/dev/null");
}

$verbose and print "Loading the following modules automatically at node boot:\n";
$verbose and print "   $load_mod\n";

open(LOAD_MODULES, "> /var/tmp/.wwinitrd-$random_string/etc/load_modules");
foreach (@load_mods) {
   print LOAD_MODULES "$_\n";
}
close LOAD_MODULES;

$verbose and print "Unmounting loop device\n";
system("umount /var/tmp/.wwinitrd-$random_string.img");
$verbose and print "Compressing initial ram disk\n";
system("gzip -9f /var/tmp/.wwinitrd-$random_string.img");
system("/bin/mv /var/tmp/.wwinitrd-$random_string.img.gz /tftpboot/wwinitrd-$kernel_version.img");

#$verbose and print "Creating wwinitrd archive\n";
#chdir("/var/tmp/.wwinitrd-$random_string/");
#system("find . | cpio --quiet -o -H newc > /tftpboot/wwinitrd-$kernel_version.img");
#$error += $?;
#$verbose and print "Compressing wwinitrd archive\n";
#system("gzip -9f /tftpboot/wwinitrd-$kernel_version.img");
#$error += $?;

if ( $kernel ) {
   copy("$kernel", "/tftpboot/");
} elsif ( -f "/boot/vmlinuz-$kernel_version" ) {
   copy("/boot/vmlinuz-$kernel_version", "/tftpboot/kernel-$kernel_version");
} elsif ( -f "/boot/vmlinux-$kernel_version" ) {
   copy("/boot/vmlinux-$kernel_version", "/tftpboot/kernel-$kernel_version");
} else {
   warn "Could not find kernel with version '$kernel_version' in /boot!\n";
   warn "You should copy it to: /tftpboot/kernel-$kernel-version\n";
}


if ( $error ) {
   die "Left /var/tmp/.wwinitrd-$random_string around for debugging\n";
} else {
   system("rm -rf /var/tmp/.wwinitrd-$random_string");
}

unlink("/tftpboot/kernel");
unlink("/tftpboot/wwinitrd.img");
link("/tftpboot/kernel-$kernel_version", "/tftpboot/kernel");
move("/tftpboot/wwinitrd-$kernel_version.img.gz" , "/tftpboot/wwinitrd-$kernel_version.img");
link("/tftpboot/wwinitrd-$kernel_version.img", "/tftpboot/wwinitrd.img");
$verbose and print "WROTE: /tftpboot/kernel-$kernel_version\n";
$verbose and print "WROTE: /tftpboot/kernel\n";
$verbose and print "WROTE: /tftpboot/wwinitrd-$kernel_version.img\n";
$verbose and print "WROTE: /tftpboot/wwinitrd.img\n";
