#!/usr/bin/perl -wT
#
# Copyright (c) 2001- 2016 University of Utah and the Flux Group.
# 
# {{{EMULAB-LICENSE
# 
# This file is part of the Emulab network testbed software.
# 
# This file is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at
# your option) any later version.
# 
# This file 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 Affero General Public
# License for more details.
# 
# You should have received a copy of the GNU Affero General Public License
# along with this file.  If not, see <http://www.gnu.org/licenses/>.
# 
# }}}
#

#
# This utility manages setting, removing and displaying taint states for
# any Emulab objects that support them (e.g., nodes and OS descriptors).
#

use strict;
use English;
use Getopt::Std;

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

#
# Testbed Support libraries
#
use lib "@prefix@/lib";
use libdb;
use EmulabConstants;
use User;

# Protos
sub usage();
sub fatal($);
sub ParseArgs(@);
sub DoAction($$$@);
sub GetObject($$);
sub PrintAllObjects($);

#
# Top level constants
#
my @ALLTAINTS = TB_TAINTSTATE_ALL();
my @TAINT_OBJECT_TYPES = ("node","os");
my @TAINT_ACTIONS = ("add","remove","set","clear","show","showall");

#
# Global variables
#

#
# Get, set, and display taint states for supported objects
#
sub usage()
{
    print STDERR "Usage: managetaint <action> <objtype> [<OID> <taint> ...]\n";
    print STDERR "  -h       This message\n";
    print STDERR " Actions:\n";
    print STDERR "  add      Apply one or more taint modes to an object.\n";
    print STDERR "  remove   Remove one or more taint modes from an object.\n";
    print STDERR "  set      Set a taint mode on specified object.\n";
    print STDERR "  clear    Clear all taint states from an object.\n";
    print STDERR "  show     Show all taint modes set for an object.\n";
    print STDERR "  showall  Display all taint modes across all objects.\n";
    print STDERR " Object Types (and associated object identifiers):\n";
    print STDERR "  Node     Testbed Node\n";
    print STDERR "   OID     node_id\n";
    print STDERR "  OS       OS Descriptor\n";
    print STDERR "   OID     <project>/<osname>";
    print STDERR " Taint States: ". join(",",@ALLTAINTS) ."\n";
    exit(-1);
}
my $optlist  = "h";
my $debug    = 0;

#
# Please do not run as root. Hard to track what has happened.
#
if ($UID == 0) {
    die("*** $0:\n".
	"    Please do not run this as root!\n");
}

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

#
# Untaint the path
# 
$ENV{'PATH'} = "/bin:/sbin:/usr/bin:";

#
# Verify user, must be admin.
#
my $this_user = User->ThisUser();
if (! defined($this_user)) {
    fatal("You ($UID) do not exist!");
}
# Probably will need to relax this so that regular users can add taint states.
if (!$this_user->IsAdmin()) {
    fatal("You are not a testbed administrator!");
}

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

my ($action, $objtype, $oid, @in_states) = ParseArgs(@ARGV);

fatal("Argument parsing failed!") 
    if (!defined($objtype));

DoAction($action, $objtype, $oid, @in_states) or
    fatal("Failed to perform $action on $objtype object $oid");

exit 0;

#
# Input argument parsing and sanity checking.
#
sub ParseArgs(@) {
    my @UARGV = ();

    # Untaint arguments
    foreach my $targ (@_) {
	if ($targ =~ /^[-\w_\/]+$/) {
	    push @UARGV, lc($targ);
	} else {
	    fatal("Malformed argument provided on command line: $targ");
	}
    }

    my ($action, $objtype, $oid, @in_states) = @UARGV;

    # Argument sanity checks.
    if (!grep {$_ eq $action} @TAINT_ACTIONS) {
	fatal("Invalid action: $action");
    }

    if (!grep {$_ eq $objtype} @TAINT_OBJECT_TYPES) {
	fatal("Invalid taint object type: $objtype");
    }

    if (@in_states) {
	foreach my $tstate (@in_states) {
	    if (!grep {$_ eq $tstate} @ALLTAINTS) {
		fatal("Unknown taint state: $tstate");
	    }
	}
    }

    if (!defined($oid) && $action ne "showall") {
	fatal("You must supply an OID for this operation");
    }

    if (($action eq "set" || $action eq "remove" || $action eq "add") 
	&& !scalar(@in_states)) {
	fatal("You must supply one or more taint states when using the ".
	      "'add', 'remove', or 'set' actions");
    }

    return ($action, $objtype, $oid, @in_states);
}

#
# Perform $action on the object of type $objtype identified by $oid.
# @in_states are taint states to be applied / removed (based on $action).
#
sub DoAction($$$@) {
    my ($action, $objtype, $oid, @in_states) = @_;

    # Any objects returned are expected to implement the 'libTaintStates'
    # interface.
    my $obj;
    if ($action ne "showall") {
	$obj = GetObject($objtype,$oid) or
	    fatal("Could not lookup object with type '$objtype' ".
		  "and identity '$oid'");
    }

  ACTION: for ($action) {
      /^add$/ && do {
	  foreach my $tstate (@in_states) {
	      $obj->AddTaintState($tstate) == 0 or
		  fatal("Failed to apply '$tstate' taint to '$oid'");

	  }	  
	  last ACTION;
      };

      /^remove$/ && do {
	  foreach my $tstate (@in_states) {
	      $obj->RemoveTaintState($tstate) == 0 or
		  fatal("Failed to remove '$tstate' taint from '$oid'");

	  }
	  last ACTION;
      };

      /^set$/ && do {
	  $obj->SetTaintStates(@in_states) == 0 or
	      fatal("Failed to set taint states for '$oid'");
	  last ACTION;
      };

      /^clear$/ && do {
	  $obj->SetTaintStates(()) == 0 or
	      fatal("Failed to clear taint states from '$oid'");
	  last ACTION;
      };

      /^show$/ && do {
	  my @taint_states = $obj->GetTaintStates();
	  if (@taint_states) {
	      my $pstr = join(", ", @taint_states);
	      print "Taint states on '$oid': $pstr\n";
	  } else {
	      print "No taint states set on '$oid'\n";
	  }
	  last ACTION;
      };

      /^showall$/ && do {
	  PrintAllObjects($objtype);
	  last ACTION;
      };

      # Default
      fatal("Unknown action: $action");
  }

    return 1;
}

#
# Return the Emulab DB abstraction object based on the object type
# specified and the lookup id provided.  Returns undef if not found.
#
sub GetObject($$) {
    my ($objtype, $oid) = @_;
    
    my $retobj;

  TYPE: for ($objtype) {
      /^node$/ && do {
	  require Node;
	  $retobj = Node->Lookup($oid);
	  last TYPE;
      };

      /^os$/ && do {
	  require OSImage;
	  my ($ospid, $osname) = split("/", $oid);
	  fatal("OS identifier must be specified as '<project_id>/<os_name>'") 
	      if (!defined($osname));
	  $retobj = OSImage->Lookup($ospid, $osname);
	  last TYPE;
      };

      # Default
      fatal("Unknown object type: $objtype");
  }

    return $retobj;
}

#
# Grab all objects of the given type that have taint states.  Return them
# as an array of arrays.  The elements of the contained arrays are the
# object identifier and the list of taint states for the associated object.
#
sub PrintAllObjects($) {
    my ($objtype) = @_;

  TYPE: for ($objtype) {
      /^node$/ && do {
	  my $query_result =
	      DBQueryWarn("select node_id, taint_states from nodes ".
			  "where taint_states is not null ".
			  "and taint_states != ''");

	  fatal("Database lookup error.")
	      if (!$query_result);

	  print "Node_id\tTaint_States\n";
	  print "----------------------------------------------------------\n";

	  while (my ($nodeid, $tstates) = $query_result->fetchrow_array()) {
	      print "$nodeid\t$tstates\n";
	  }

	  last TYPE;
      };

      /^os$/ && do {
	  my $query_result =
	      DBQueryWarn("select pid, osname, taint_states from os_info_versions ".
			  "where taint_states is not null ".
			  "and taint_states != ''");

	  fatal("Database lookup error.")
	      if (!$query_result);

	  print "OS(pid/name)\tTaint_States\n";
	  print "----------------------------------------------------------\n";

	  while (my ($pid, $osname, $tstates) = 
		 $query_result->fetchrow_array()) {
	      print "$pid/$osname\t$tstates\n";
	  }

	  last TYPE;
      };

      # Default
      fatal("Unknown object type: $objtype");
  }

    return;
}

sub fatal($)
{
    my ($mesg) = $_[0];

    die("*** $0:\n".
	"    $mesg\n");
}