#!/usr/bin/env perl
use File::Basename;
use File::Path;
use File::Copy;
use File::Find;
use Getopt::Long;
Getopt::Long::Configure("bundling");
Getopt::Long::Configure("pass_through");

my $prinic; #TODO be flexible on node primary nic
my $othernics; #TODO be flexible on node primary nic
my $netdriver;
my @yumdirs;
my $arch = `uname -m`;
chomp($arch);
my %libhash;
my @filestoadd;
my $profile;
my $osver;
my $pathtofiles=dirname($0);
my $name = basename($0);
my $onlyinitrd=0;
if ($name =~ /geninitrd/) {
   $onlyinitrd=1;
}

GetOptions(
   #'a=s' => \$architecture,
   'p=s' => \$profile,
   'o=s' => \$osver,
   'n=s' => \$netdriver,
   'i=s' => \$prinic,
   'r=s' => \$othernics
);
unless ($osver and $profile and $netdriver and $prinic) {
   print 'Usage: genimage -i <nodebootif> -n <nodenetdrivers> [-r <otherifaces>] -o <OSVER> -p <PROFILE>'."\n";
   print "Examples:\n";
   print " genimage -i eth0 -n tg3 -o centos5.1 -p compute\n";
   print " genimage -i eth0 -r eth1,eth2 -n tg3,bnx2 -o centos5.1 -p compute\n";
   exit 1;
}
my @ndrivers;
foreach (split /,/,$netdriver) {
   unless (/\.ko$/) {
      s/$/.ko/;
   }
   if (/^$/) {
      next;
   }
   push @ndrivers,$_;
}

my $installroot = "/install";
my $srcdir = "$installroot/$osver/$arch";
@yumdirs=();
find(\&isyumdir, <$installroot/$osver/$arch/>);
unless (scalar(@yumdirs)) {
#unless ( -d $srcdir."/repodata" ) {
   print "Need $installroot/$osver/$arch/ available from a system that has ran copycds on $osver $arch";
   exit 1;
}
my $yumconfig;
open($yumconfig,">","/tmp/genimage.$$.yum.conf");
my $repnum=0;
foreach $srcdir (@yumdirs) {
   print $yumconfig "[$osver-$arch-$repnum]\nname=$osver-$arch-$repnum\nbaseurl=file://$srcdir\ngpgpcheck=0\n\n";
   $repnum += 1;
}
$repnum-=1;
close($yumconfig);
my $yumcmd = "yum -y -c /tmp/genimage.$$.yum.conf --installroot=$installroot/netboot/$osver/$arch/$profile/rootimg/ --disablerepo=* ";
foreach (0..$repnum) {
   $yumcmd .= "--enablerepo=$osver-$arch-$_ "
}
$yumcmd .= "install ";
mkpath("$installroot/netboot/$osver/$arch/$profile/rootimg/var/lib/yum");
my $pkglist;
if (-r "$pathtofiles/$profile.$osver.$arch.pkglist") {
   $pkglist = "$pathtofiles/$profile.$osver.$arch.pkglist";
} elsif (-r "$pathtofiles/$profile.$arch.pkglist") {
   $pkglist = "$pathtofiles/$profile.$arch.pkglist";
} elsif (-r "$pathtofiles/$profile.$osver.pkglist") {
   $pkglist = "$pathtofiles/$profile.$osver.pkglist";
} elsif (-r "$pathtofiles/$profile.pkglist") {
   $pkglist = "$pathtofiles/$profile.pkglist";
} else {
   print "Unable to find package list for $profile!";
   return 1;
}
open($yumconfig,"<","$pkglist");
while (<$yumconfig>) {
   chomp;
   $yumcmd .= $_ . " ";
}
$yumcmd =~ s/ $/\n/;
unless ($onlyinitrd) {
   my $rc = system($yumcmd);
   if ($rc) { 
      print "yum invocation failed\n";
      exit 1;
   }
   postscripts(); #run 'postscripts'
}
unlink "/tmp/genimage.$$.yum.conf";
mkinitrd();

sub getlibs {
	my $file = shift;
	my $liblist = `chroot $installroot/netboot/$osver/$arch/$profile/rootimg ldd $file`;
	my @libs = split/\n/,$liblist;
	my @return;
	foreach (@libs) {
      unless (/=>/) {
         (my $wjnk, my $lib,my $jnk) = split /\s+/,$_,3;
         $lib =~ s/^\///;
         $libhash{$lib}=1;
         next;
      }
		(my $temp1,my $temp2) = split />/,$_,2;
		(my $whitespace,$temp1,$temp2) = split /\s+/,$temp2,4;
		unless ($temp1 =~ /\//) {
			next;
		}
      $temp1 =~ s/^\///;
		$libhash{$temp1}=1;
   }
}

sub mkinitrd {
   mkpath("/tmp/xcatinitrd.$$/bin");
   rename(<$installroot/netboot/$osver/$arch/$profile/rootimg/boot/vmlinuz*>,"$installroot/netboot/$osver/$arch/$profile/kernel");
	symlink("bin","/tmp/xcatinitrd.$$/sbin");
	mkpath("/tmp/xcatinitrd.$$/usr/bin");
	mkpath("/tmp/xcatinitrd.$$/usr/sbin");
	mkpath("/tmp/xcatinitrd.$$/usr/lib");
	mkpath("/tmp/xcatinitrd.$$/usr/lib64");
	mkpath("/tmp/xcatinitrd.$$/lib/firmware");
	mkpath("/tmp/xcatinitrd.$$/lib64/firmware");
	mkpath("/tmp/xcatinitrd.$$/proc");
	mkpath("/tmp/xcatinitrd.$$/sys");
	mkpath("/tmp/xcatinitrd.$$/dev/mapper");
	mkpath("/tmp/xcatinitrd.$$/sysroot");
	mkpath("/tmp/xcatinitrd.$$/etc/ld.so.conf.d");
	mkpath("/tmp/xcatinitrd.$$/var/lib/dhclient");
	my $inifile;
	open($inifile,">","/tmp/xcatinitrd.$$/init");
	print $inifile "#!/sbin/busybox.anaconda sh\n";
	print $inifile "busybox.anaconda mount -t proc /proc /proc\n";
	print $inifile "busybox.anaconda --install\n";
	print $inifile "mount -t sysfs /sys /sys\n";
   print $inifile "mount -o mode=0755 -t tmpfs /dev /dev\n";
   print $inifile "mkdir /dev/pts\n";
   print $inifile "mount -t devpts -o gid=5,mode=620 /dev/pts /dev/pts\n";
   print $inifile "mkdir /dev/shm\n";
   print $inifile "mkdir /dev/mapper\n";
   print $inifile "mknod /dev/null c 1 3\n";
   print $inifile "mknod /dev/zero c 1 5\n";
   print $inifile "mknod /dev/systty c 4 0\n";
   print $inifile "mknod /dev/tty c 5 0\n";
   print $inifile "mknod /dev/console c 5 1\n";
   print $inifile "mknod /dev/ptmx c 5 2\n";
   print $inifile "mknod /dev/rtc c 10 135\n";
   print $inifile "mknod /dev/tty0 c 4 0\n";
   print $inifile "mknod /dev/tty1 c 4 1\n";
   print $inifile "mknod /dev/tty2 c 4 2\n";
   print $inifile "mknod /dev/tty3 c 4 3\n";
   print $inifile "mknod /dev/tty4 c 4 4\n";
   print $inifile "mknod /dev/tty5 c 4 5\n";
   print $inifile "mknod /dev/tty6 c 4 6\n";
   print $inifile "mknod /dev/tty7 c 4 7\n";
   print $inifile "mknod /dev/tty8 c 4 8\n";
   print $inifile "mknod /dev/tty9 c 4 9\n";
   print $inifile "mknod /dev/tty10 c 4 10\n";
   print $inifile "mknod /dev/tty11 c 4 11\n";
   print $inifile "mknod /dev/tty12 c 4 12\n";
   print $inifile "mknod /dev/ttyS0 c 4 64\n";
   print $inifile "mknod /dev/ttyS1 c 4 65\n";
   print $inifile "mknod /dev/ttyS2 c 4 66\n";
   print $inifile "mknod /dev/ttyS3 c 4 67\n";
   foreach (@ndrivers) {
      print $inifile "insmod /lib/$_\n";
   }
   print $inifile "netstart\n";
   print $inifile "cd /\n";
   print $inifile "for i in `cat /proc/cmdline`; do\n";
   print $inifile "   KEY=`echo \$i |awk -F= '{print \$1}'`\n";
   print $inifile "   if [ \"\$KEY\" == 'imgurl' ]; then\n";
   print $inifile "      VALUE=`echo \$i |awk -F= '{print \$2}'`\n";
   print $inifile "      if [ \"http\" == \"`echo \$VALUE|awk -F: '{print \$1}'`\" ]; then\n";
   print $inifile "        wget \$VALUE\n";
   print $inifile "     fi\n";
   print $inifile "   fi\n";
   print $inifile "done\n";
   print $inifile "if [ -r /rootimg.sfs ]; then\n";
   print $inifile "  echo Setting up squashfs with ram overlay\n";
   print $inifile "  mknod /dev/loop0 b 7 0\n";
   print $inifile "  mkdir -p /ro\n";
   print $inifile "  mkdir -p /rw\n";
   print $inifile "  mount -t squashfs /rootimg.sfs /ro\n";
   print $inifile "  mount -t tmpfs rw /rw\n";
   print $inifile "  mount -t aufs -o dirs=/rw:/ro mergedroot  /sysroot\n";
   print $inifile "elif [ -r /rootimg.gz ]; then\n";
   print $inifile "  mount -t tmpfs rootfs /sysroot\n";
   print $inifile "  cd /sysroot\n";
   print $inifile "  echo -n \"Extracting root filesystem:\"\n";
   print $inifile "  if [ -x /bin/cpio ]; then\n";
   print $inifile "  zcat /rootimg.gz |/bin/cpio -idum\n";
   print $inifile "  else\n";
   print $inifile "  zcat /rootimg.gz |cpio -idum\n";
   print $inifile "  fi\n";
   print $inifile "  echo Done\n";
   print $inifile "else\n";
   print $inifile "  echo -n Failed to download image, panicing in 5...\n";
   print $inifile "  for i in 4 3 2 1 0; do\n";
   print $inifile "    /bin/sleep 1\n";
   print $inifile "    echo -n \$i...\n";
   print $inifile "  done\n";
   print $inifile "  echo\n";
   print $inifile "  exit\n";
   print $inifile "fi\n";
   print $inifile "cd /\n";
   print $inifile "cp /var/lib/dhclient/dhclient.leases /sysroot/dev/.dhclient-eth0.leases\n";
   print $inifile "mknod /sysroot/dev/console c 5 1\n";
   print $inifile "exec switch_root -c /dev/console /sysroot /sbin/init\n";
   close($inifile);
   open($inifile,">"."/tmp/xcatinitrd.$$/bin/netstart");
   print $inifile "#!/sbin/nash\n";
   print $inifile "network --device eth0 --bootproto dhcp\n";
   close($inifile);
	chmod(0755,"/tmp/xcatinitrd.$$/init");
	chmod(0755,"/tmp/xcatinitrd.$$/bin/netstart");
   @filestoadd=();
   foreach (@ndrivers) {
      if (-f "$pathtofiles/$_") {
         push @filestoadd,[$_,"lib/$_"];
      }
   }
   foreach ("bin/cpio","sbin/nash","sbin/busybox.anaconda","sbin/rmmod") {
      getlibs($_);
      push @filestoadd,$_;
   }
   push @filestoadd,keys %libhash;
   find(\&isnetdriver, <$installroot/netboot/$osver/$arch/$profile/rootimg/lib/modules/*>);
   foreach (@filestoadd) {
      if (ref($_)) {
         my $srcpath = "$installroot/netboot/$osver/$arch/$profile/rootimg/".$_->[0];
         if (-f "$pathtofiles/".$_->[0]) {
            $srcpath="$pathtofiles/".$_->[0];
         }
         copy($srcpath,"/tmp/xcatinitrd.$$/".$_->[1]);
         chmod 0755,"/tmp/xcatinitrd.$$/".$_->[1];
      } else {
         my $srcpath = "$installroot/netboot/$osver/$arch/$profile/rootimg/$_";
         if (-f "$pathtofiles/$_") {
            $srcpatch = "$pathtofiles/$_";
         }
         copy("$srcpath","/tmp/xcatinitrd.$$/$_");
         chmod 0755,"/tmp/xcatinitrd.$$/".$_;
      }
   }

   #copy("$installroot/netboot/$osver/$arch/$profile/rootimg/lib/modules/*d","/tmp/xcatinitrd.$$/$_");
   system("cd /tmp/xcatinitrd.$$;find .|cpio -H newc -o|gzip -9 -c - > $installroot/netboot/$osver/$arch/$profile/initrd.gz");
   system("rm -rf /tmp/xcatinitrd.$$");

}

sub isyumdir {
   if ($File::Find::name =~ /\/repodata$/) {
      my $location = $File::Find::name;
      $location =~ s/\/repodata$//;
      push @yumdirs,$location;
   }
}

sub isnetdriver {
   foreach (@ndrivers) {
      if ($File::Find::name =~ /$_/) {
         my $filetoadd = $File::Find::name;
         $filetoadd =~ s!$installroot/netboot/$osver/$arch/$profile/rootimg/!!;
         push @filestoadd,[$filetoadd,"lib/$_"];
      }
   }
}
      	






sub postscripts { # TODO: customized postscripts
   generic_post();
   if (-d "$installroot/postscripts/hostkeys") {
      for my $key (<$installroot/postscripts/hostkeys/*key>) {
         copy ($key,"$installroot/netboot/$osver/$arch/$profile/rootimg/etc/ssh/");
      }
      chmod 0600,</$installroot/netboot/$osver/$arch/$profile/rootimg/etc/ssh/*key>;
   }
   if (-d "/$installroot/postscripts/.ssh") {
      mkpath("/$installroot/netboot/$osver/$arch/$profile/rootimg/root/.ssh");
      chmod(0700,"/$installroot/netboot/$osver/$arch/$profile/rootimg/root/.ssh/");
      for my $file (</$installroot/postscripts/.ssh/*>) {
         copy ($file,"/$installroot/netboot/$osver/$arch/$profile/rootimg/root/.ssh/");
      }
      chmod(0600,</$installroot/netboot/$osver/$arch/$profile/rootimg/root/.ssh/*>);
   }
}

sub generic_post { #This function is meant to leave the image in a state approximating a normal install
   my $cfgfile;
   unlink("$installroot/netboot/$osver/$arch/$profile/rootimg/dev/null");
   system("mknod $installroot/netboot/$osver/$arch/$profile/rootimg/dev/null c 1 3");
   open($cfgfile,">","$installroot/netboot/$osver/$arch/$profile/rootimg/etc/fstab");
   print $cfgfile "devpts  /dev/pts devpts   gid=5,mode=620 0 0\n";
   print $cfgfile "tmpfs   /dev/shm tmpfs    defaults       0 0\n";
   print $cfgfile "proc    /proc    proc     defaults       0 0\n";
   print $cfgfile "sysfs   /sys     sysfs    defaults       0 0\n";
   close($cfgfile);
   open($cfgfile,">","$installroot/netboot/$osver/$arch/$profile/rootimg/etc/sysconfig/network");
   print $cfgfile "NETWORKING=yes\n";
   close($cfgfile);
   open($cfgfile,">","$installroot/netboot/$osver/$arch/$profile/rootimg/etc/sysconfig/network-scripts/ifcfg-$prinic");
   print $cfgfile "ONBOOT=yes\nBOOTPROTO=dhcp\nDEVICE=$prinic\n";
   close($cfgfile);
   foreach (split /,/,$othernics) {
      if (/^$/) { next; }
      open($cfgfile,">","$installroot/netboot/$osver/$arch/$profile/rootimg/etc/sysconfig/network-scripts/ifcfg-$_");
      print $cfgfile "ONBOOT=yes\nBOOTPROTO=dhcp\nDEVICE=$_\n";
      close($cfgfile);
   }
   open($cfgfile,">","$installroot/netboot/$osver/$arch/$profile/rootimg/etc/rc3.d/S60gettyset");
   print $cfgfile "#!/bin/bash\n"; 
   print $cfgfile "for i in `cat /proc/cmdline`; do\n"; 
   print $cfgfile '  KEY=`echo $i|cut -d= -f 1`'."\n";
   print $cfgfile "  if [ \"\$KEY\" == \"console\" ]; then\n";
   print $cfgfile "    VALUE=`echo \$i | cut -d= -f 2`\n"; 
   print $cfgfile "     COTTY=`echo \$VALUE|cut -d, -f 1`\n";
   print $cfgfile "     COSPEED=`echo \$VALUE|cut -d, -f 2|cut -dn -f 1`\n";
   print $cfgfile "     if echo \$VALUE | grep n8r; then\n";
   print $cfgfile "        FLOWFLAG=\"-h\"\n";
   print $cfgfile "     fi\n";
   print $cfgfile "     echo xco:2345:respawn:/sbin/agetty \$FLOWFLAG \$COTTY \$COSPEED xterm >> /etc/inittab\n";
   print $cfgfile "     init q\n";
   print $cfgfile "  fi\n";
   print $cfgfile "done\n";
   close($cfgfile);
   chmod(0755,"$installroot/netboot/$osver/$arch/$profile/rootimg/etc/rc3.d/S60gettyset");
   #link("$installroot/netboot/$osver/$arch/$profile/rootimg/sbin/init","$installroot/netboot/$osver/$arch/$profile/rootimg/init");
   
   rename(<$installroot/netboot/$osver/$arch/$profile/rootimg/boot/vmlinuz*>,"$installroot/netboot/$osver/$arch/$profile/kernel");
}