#!/usr/bin/perl -wT
use English;
use Getopt::Std;
   
#
# XXX Paper and plastic IP addresses wired into the kernel choice.
#     Path to netdisk is wired in (should come from os_info table).
#     wd0 wired in. Should come from node_types table in DB
# 

#
# Load an image onto a disk. The image must be in the DB images table,
# which defines how/where to load, and what partitions are affected.
# The nodes and partitions tables are updated appropriately. 
# 
# usage: os_load [-s] <imageid> <node> [node ...]
#
sub usage()
{
    print STDOUT "Usage: os_load [-s] <imageid> <node> [node ...]\n".
	"Use the -s to setup reload only, but do not issue a reboot\n";
    exit(-1);
}
my  $optlist = "s";

#
# Configure variables
#
my $TB		= "@prefix@";

#
# Load the Testbed support stuff. 
#
push(@INC, "$TB/lib");
require libdb;

my $NETDISK     = "/tftpboot/netdisk";
my $PAPERADDR	= "boss.emulab.net";
my $PLASTICADDR	= "users.emulab.net";
my $nodereboot	= "$TB/bin/node_reboot";
my $dbg		= 0;
my @row;
my %imageid_row = ();
my @nodes       = ();
my $mereuser    = 0;
my $setuponly   = 0;
my $failures    = 0;

# un-taint path
$ENV{'PATH'} = '/bin:/usr/bin:/usr/local/bin';
delete @ENV{'IFS', 'CDPATH', 'ENV', 'BASH_ENV'};

$| = 1; #Turn off line buffering on output

#
# Parse command arguments. Once we return from getopts, all that should be
# left are the required arguments.
#
%options = ();
if (! getopts($optlist, \%options)) {
    usage();
}
if (@ARGV < 2) {
    usage();
}
if (defined($options{"s"})) {
    $setuponly = 1;
}
my $imageid = shift;

#
# Untaint args.
#
if ($imageid =~ /^([-\@\w.\+]+)$/) {
    $imageid = $1;
}
else {
    die("Bad data in $imageid.");
}

foreach my $node ( @ARGV ) {
    if ($node =~ /^([-\@\w]+)$/) {
	$node = $1;
    }
    else {
	die("Bad node name: $node.");
    }
    
    push(@nodes, $node);
}

#
# Figure out who called us. Root and admin types can do whatever they
# want. Normal users can only change nodes in experiments in their
# own projects.
#
if ($UID && !TBAdmin($UID)) {
    $mereuser = 1;
}

#
# Grab the imageid description from the DB. The permission check for this
# is that mere user can load an image from any project he/she is a member of,
# or any image that has a null pid field, since those are defined to be
# open to anyone. 
#
$db_result = DBQueryFatal("select * from images where imageid='$imageid'");
    
if ($db_result->numrows < 1) {
    die("No such imageid $imageid!");
}
%imageid_row = $db_result->fetchhash();

my $imagepid = 0;
if (defined($imageid_row{'pid'})) {
    $imagepid = $imageid_row{'pid'};
}
    
if ($mereuser && $imagepid && !ProjMember($imagepid)) {
    die("You do not have permission to load imageid $imageid!");
}

#
# Check to make sure that mere user is allowed to muck with nodes
#
if ($mereuser) {
    foreach my $node (@nodes) {
	if (! NodeAccessCheck(\$node)) {
	    die("You do not have permission to load images on $node\n");
	}
    }
}

my $loadpart  = $imageid_row{'loadpart'};
my $loadlen   = $imageid_row{'loadlength'};
my $imagepath = $imageid_row{'path'};
my $defosid   = $imageid_row{'default_osid'};

#
# 0 means load the entire disk.
#
my $diskpart = "";
if ($loadpart) {
    $diskpart = "wd0:s${loadpart}";
}
else {
    $diskpart = "wd0";
}

#
# For now, all testbed default images from from paper and all pid specific
# images come from plastic:/proj.
# 
my $cmdline = "";
if ($imagepid) {
    if (! ($imagepath =~ /^\/proj\//)) {
	die("Your image must reside in /proj\n");
    }
    $cmdline = "${PLASTICADDR}:$imagepath $diskpart";    
}
else {
    $cmdline = "${PAPERADDR}:$imagepath $diskpart";
}

#
# Loop for each node.
# 
foreach my $node (@nodes) {
    my $pc = $node;
	
    print STDOUT "Changing default OS for $pc to $defosid\n";
    DBQueryFatal("update nodes set ".
		 "def_boot_osid='$defosid',def_boot_path='' ".
		 "where node_id='$pc'");

    #
    # Assign partition table entries for each partition in the image.
    # This is complicated by the fact that an image that covers only
    # part of the slices, should only change the partition table entries
    # for the subset of slices that are written to disk. In reality, this
    # is silly since there is no telling what the disk is going to end
    # up looking like after a partial image is written, especially if its
    # a user generated image. Not sure how to handle this yet. For now
    # lets just say that a user defined images essentially wipe the disk
    # except for the stuff they write. 
    #
    for ($i = 1; $i <= 4; $i++) {
	my $partname = "part${i}_osid";
	
	if (defined($imageid_row{$partname})) {
	    my $osid = $imageid_row{$partname};

	    DBQueryFatal("replace into partitions ".
			 "(partition, osid, node_id) ".
			 "values('$i', '$osid', '$pc')");
	}
	else {
	    DBQueryFatal("delete from partitions ".
			 "where node_id='$pc' and partition='$i'");
	}
    }
    
    print STDOUT "Setting up reload for $pc\n";
    DBQueryFatal("update nodes set ".
		 "next_boot_path='$NETDISK',".
		 "next_boot_cmd_line='$cmdline' ".
		 "where node_id='$pc'");
}

#
# Fire off a reboot.
#
if (! $setuponly) {
    system("$nodereboot @nodes");
    $failures = $? >> 8;
}

print STDOUT "OS Reload Done!\n";
exit $failures;