#!/usr/bin/perl -w

##########################################################################
# Copyright 2005 VMware, Inc.  All rights reserved. -- VMware Confidential
##########################################################################

#
# esxcfg-boot --
#
#       This script provides for interactions with the boot configuration
#       of ESX Server.  It allows the user to update or query the
#       configuration.
#

use lib "/usr/lib/vmware/esx-perl/perl5/site_perl/5.8.0";

use VMware::Boot::BootInfo qw(LoadChecksumManager UpdateBootConfig
                              CommitConfig CommitInitrd);
use VMware::Boot::BootManager;
use VMware::CmdTool;
use VMware::FileSys::StandardLocator qw(StdDir StdFile StdCommand);
use VMware::Log qw(:log);
use VMware::PCI::PCIInfo qw(FindPCIDeviceManager
                            GetPCIDeviceManagerFromConfig
                            SetPCIDeviceDataInConfig);
use VMware::PCI::Device;
use VMware::PCI::DeviceManager;
use VMware::PCI::DriverManager;
use strict;


########################################################################
#
# QueryBootConfig --
#
#       This function just gets whatever information is returned
#       by VMware::Boot::BootManager::QueryBootConfig and prints
#       it in a one-line, space-separated list.
#
# Results:
#       1 on success, 0 otherwise.
#
# Side effects:
#       May write to stdout.
#
########################################################################

sub QueryBootConfig
{
   my $config = shift; # IN: ESX configuration data, read-only.
   
   my $bootMgr = new VMware::Boot::BootManager($config->GetTree(['boot']));
   my $info = $bootMgr->QueryBootConfig();

   unless (defined $info) {
      LogError("Could not determine boot configuration.");
      return 0;
   }
   
   print join(' ', @$info), "\n";
   return 1;
}


########################################################################
#
# QueryVMKModules --
#
#       Writes out information for enabled modules in a format
#       suitable for passing to vmkload_mod.
#
# Results:
#       1 on success, 0 otherwise.
#
# Side effects:
#       Writes the module information, if any, to stdout.
#
########################################################################

sub QueryVMKModules
{
   system("esxcfg-module -q");
   return 1;
}


########################################################################
#
# CmdQuery --
#
#       Implements the "--query" command.
#
# Results:
#       1 on success, 0 otherwise.
#
# Side effects:
#       Loads esx.conf. May load grub.conf.
#
########################################################################

sub CmdQuery
{
   my $query = shift;   # IN: The query to run.
   
   my $config = VMware::CmdTool::LoadConfig();
   if ($query eq "boot") {
      return QueryBootConfig($config);
   } elsif ($query eq "vmkmod") {
      return QueryVMKModules($config);
   } else {
      LogError("Invalid query '$query' requested.");
      return 0;
   }
}


########################################################################
#
# UpdatePCIConfig --
#
#       Build the PCI configuration from existing config files
#       and/or probing the system with lspci.  
#       (Moot:) Assign all VMkernel-manageable device to the VMkernel.
#       Name devices and generate cosVSwitch config.
#       Update the config object with the results.
#       Ramdisk rebuilt for -p (updatepci) only. 
#
# Results:
#       True on success, false on failure.
#
# Side effects:
#       Updates contents of the configuration object.
#
########################################################################

sub UpdatePCIConfig
{
   my $config = shift;  # IN/OUT: The esx configuration object.
   my $cosVswitch = shift; # IN: If true, configure the COS vswitch.
   my $oldCPCI = shift; # IN: The previously configured CPCI line.

   LogInfo("Examining PCI Devices...");
   my $devMgr = FindPCIDeviceManager($config, 0, $oldCPCI);
   unless (defined $devMgr) {
      LogError("Could not determine PCI device configuration.");
      return undef;
   }

   LogInfo("Setting PCI Device Ownership...");
   foreach my $dev (@{$devMgr->GetDevices()}) {
      if ($dev->IsVMwareManageable()) {
         $dev->SetOwner(VMware::PCI::Device::VMKERNEL);
      }
   }

   #
   # Handle the vswitch before generating names as we will
   # want to name the COS vswitch specially first.
   #

   if ($cosVswitch) {
      ConfigureCOSVswitch($config, $devMgr);
   }
   LogInfo("Naming VMkernel-owned PCI Devices...");
   $devMgr->GenerateVMKDeviceNames();

   return SetPCIDeviceDataInConfig($config, $devMgr);
}


########################################################################
#
# ConfigureCOSVSwitch --
#
#       Configure the initial Console OS networking information.
#       This should only be done during installation.
#
# Results:
#       True on success, false on failure (due to missing info, etc.).
#
# Side effects:
#       Config object activity.  Networking should work afterwards
#       if someone commits the config back to the file.
#
########################################################################

sub ConfigureCOSVswitch
{
   my $config = shift;  # IN/OUT: The esx configuration object.
   my $devMgr = shift;  # IN/OUT: The device manager representing our devs.

   my $cosNicName = 'vmnic0';
   my $netcfg = $config->GetTree(['net']);

   LogInfo("Configuring COS vswitch...");
   unless (defined $netcfg->{cosnic}) {
      LogError("cosnic not defined in esx.conf");
      return undef;
   }
   unless (defined $netcfg->{vmnetpg}) {
      LogInfo("vmnetpg not defined in esx.conf default to off");
      $netcfg->{vmnetpg} = 0;
   }
   unless (defined $netcfg->{defaultvlan}) {
      LogInfo("default vlan not defined in esx.conf default to 0");
      $netcfg->{defaultvlan} = 0;
   }

   my $vmnic0InUse = grep {
      defined($_) && $_ eq $cosNicName;
   } map {
      $_->GetVMKDeviceName();
   } @{$devMgr->GetDevices()};

   if ($vmnic0InUse) {
      #
      # TODO: Fix this for upgrade.
      #
      #       If we get here, we already have named (at least some) nics.
      #       Find the one that we want and use its name.  Which is probably
      #       vmnic0, but might not be.
      #

      LogInfo("vmnics already named, checking to see which one to use.");

      my $cosNic = $devMgr->GetDevice($netcfg->{cosnic});
      $cosNicName = $cosNic->GetVMKDeviceName();

      if (defined($cosNicName)) {
         LogInfo("Found $cosNicName attached to the COS vswitch.  Continuing.");

      } else {
         #
         # Punt.  The world is seriously askew at this point.
         # Either all devices should be named or none.
         #
         LogError("Devices partially named.  Unsure how to name COS nic.");
         return undef;
      }

   } else {
      $devMgr->SetVMKDeviceName($netcfg->{cosnic}, $cosNicName);
   }

   delete $netcfg->{cosnic};

   #
   # Now add all of the config info needed to set up the first vswitch.
   #

   $netcfg->{vswitch} = {} unless defined $netcfg->{vswitch};
   $netcfg->{vswitch}->{child} = [] unless defined $netcfg->{vswitch}->{child};
   if (defined $netcfg->{vswitch}->{child}->[0]) {
      LogError("vswitch 0 is already defined, but is needed " .
               "for Service Console networking.");
      return undef;
   }

   my $pg_children = [];
   my $pg_child = {};

   #
   # add the service console portgroup
   #

   $pg_child = {
      name => "Service Console",
      vlanId => $netcfg->{defaultvlan},
   };

   push @{$pg_children}, $pg_child;

   #
   # add the virtual machine portgroup if needed
   #

   if($netcfg->{vmnetpg}) {
      $pg_child = {
         name => "VM Network",
         vlanId => $netcfg->{defaultvlan},
      };

      push @{$pg_children}, $pg_child;
   }

   delete $netcfg->{vmnetpg};
   delete $netcfg->{defaultvlan};

   $netcfg->{vswitch}->{child}->[0] = {
      name => "vSwitch0",
      beacon => {
         enable => "false",
         threshold => 1,
         timeout => 1,
      },
      numPorts => 64,
      portgroup => {
         child => $pg_children,
      },
      teamPolicy => {
         maxActive => 1,
         uplinks => [
            {
               pnic => $cosNicName,
            },
         ],
      },
      uplinks => {
         child => [
            {
               pnic => $cosNicName,
            },
         ],
      },
      cdp => {
         status => "listen",
      }
   };
   $netcfg->{vswitch}->{pgSequence} = scalar(@{$pg_children});

   return $config->SetTree(['net'], $netcfg);
}


########################################################################
#
# Usage --
#
#       Returns the usage statement as a string.
#
# Results:
#       The usage statement.
#
# Side effects:
#       None.
#
########################################################################

sub Usage
{
   #
   # Currently, verbosity and config override options are
   # not documented as they are mostly for debugging.
   #
   # The cos-vswitch option is also undocumented as it is
   # an install-time-only option.
   #

   return <<EOD;
esxcfg-boot -h --help
            -q --query boot|vmkmod
            -p --update-pci
            -b --update-boot
            -d --rootdev UUID=<uuid>
            -a --kernelappend <kernel append>
            -r --refresh-initrd
            -g --regenerate-grub

   Queries cannot be combined with each other or other options.
   Passing -p or -d enables -b even if it is not passed explicitly.
   -b implies -g plus a new initrd creation.
   -b and -r are incompatible, but -g and -r can be combined.
EOD

}


########################################################################
#
# Init --
#
#       Initialize the script and handle all command-line processing.
#
# Results:
#       A hash holding the processed command line data, or undef
#       on failure.
#
# Side effects:
#       Calls VMware::CmdTool::InitTool.
# 
########################################################################

sub Init 
{
   #
   # Define command line options and set up defaults.
   #
   # In general, flags default to 0 and arguments with
   # string parameters default to undef.
   #
   my $cmdHash = {
      "query" => undef,       # Query system configuration.
      "updatePCI" => 0,       # Examine and update PCI configuration.
      "updateBoot" => 0,      # Examine and update boot loader configuration.
      "refreshInitrd" => 0,   # Update config files in the initrd.
      "condRefresh" => 0,     # Conditional refresh of the initrd.
      "needRefresh" => 0,     # Say whether a refresh is needed.
      "rootDev" => undef,     # Set the root device.
      "kernelAppend" => undef,# Set the kernel append.
      "cosVswitch" => 0,      # Set up the COS vswitch on install.
      "regenGrub" => 0,       # Regenerate grub.conf
      "installGrub" => 0,     # Install GRUB to the boot target
      "newrd" => 0,           # Create new stateless style initrd
      "simple" => 0,          # Create simple.map file
      "schedRebuild" => 0,    # Schedule a ramdisk rebuild at the next opportunity. 
      "nameDevices" => 0,     # Find and name all PCI devs on the system.
   };

   my $opts = {
      "query|q=s" => \$cmdHash->{query},
      "rootdev|d=s" => \$cmdHash->{rootDev},
      "kernelappend|a=s" => \$cmdHash->{kernelAppend},
      "update-pcicfg|p" => \$cmdHash->{updatePCI},
      "update-bootcfg|b" => \$cmdHash->{updateBoot},
      "refresh-initrd|r" => \$cmdHash->{refreshInitrd},
      "cos-vswitch|s" => \$cmdHash->{cosVswitch},
      "cond-refresh|o" => \$cmdHash->{condRefresh},
      "need-refresh|k" => \$cmdHash->{needRefresh},
      "regenerate-grub|g" => \$cmdHash->{regenGrub},
      "install-grub|n" => \$cmdHash->{installGrub},
      "newrd" => \$cmdHash->{newrd},
      "simple" => \$cmdHash->{simple},
      "sched-rdbuild" => \$cmdHash->{schedRebuild},
      "name-devices" => \$cmdHash->{nameDevices},     
   };
   unless (VMware::CmdTool::InitTool("esxcfg_boot_cfg", "esxcfg_boot_log",
                                     $opts, Usage())) {
      return undef;
   }

   if ($cmdHash->{simple}) {
      GenerateSimpleMap("/etc/vmware/vmware-devices.map", "/etc/vmware/simple.map") ||
         die "Failed to generate /etc/vmware/simple.map\n";
      exit 0;
   }

   if (defined $cmdHash->{rootDev}) {
      my $chars = qr/[\da-fA-F]/;
      my $four = qr/$chars{4}/;
      my $eight = qr/$chars{8}/;
      my $twelve = qr/$chars{12}/;

      if ($cmdHash->{rootDev} =~
              /^(UUID=)($eight-$four-$four-$four-$twelve)$/) {
         my $prefix = $1;
         my $uuid = $2;
         $uuid =~ tr/ABCDEF/abcdef/;
         $cmdHash->{rootDev} = "$prefix$uuid";

      } else {
         print "Root device must be in UUID=<uuid> form.\n";
         exit -1;
      }
   }

   #
   # Handle dependencies and incompatible options.
   # TODO: cosVswitch should not require a full updatePCI but
   #       it will do for now.  To fix, need to get the right
   #       DeviceManager object without calling the utility
   #       function that does the current/expected comparison.
   #       Fixing this is not a high priority.
   #

   $cmdHash->{updatePCI} = 1 if $cmdHash->{cosVswitch};
   $cmdHash->{nameDevices} = 1 if $cmdHash->{updatePCI};
   $cmdHash->{updateBoot} = 1 if $cmdHash->{updatePCI} ||
                                 defined $cmdHash->{rootDev} ||
                                 defined $cmdHash->{kernelAppend};
   $cmdHash->{regenGrub} = 1 if $cmdHash->{updateBoot};

   unless ($cmdHash->{updatePCI} or
           $cmdHash->{updateBoot} or
           $cmdHash->{refreshInitrd} or
           $cmdHash->{condRefresh} or
           $cmdHash->{needRefresh} or
           $cmdHash->{regenGrub} or
           $cmdHash->{installGrub} or
           $cmdHash->{cosVswitch} or
           $cmdHash->{schedRebuild} or
           $cmdHash->{nameDevices} or
           defined $cmdHash->{query}) {
      die Usage();
   }
   if (($cmdHash->{updateBoot} or $cmdHash->{updatePCI}) and
       ($cmdHash->{refreshInitrd} and ! $cmdHash->{nameDevices})) {
      die Usage();
   }
   if (defined($cmdHash->{query}) &&
       ($cmdHash->{updateBoot} or
        $cmdHash->{refreshInitrd} or
        $cmdHash->{regenGrub} or
        $cmdHash->{installGrub} or
        $cmdHash->{schedRebuild} or
        $cmdHash->{nameDevices} or
        $cmdHash->{cosVswitch})) {
      die Usage();
   }

   if ($cmdHash->{needRefresh} && 
       (defined($cmdHash->{query}) or 
        $cmdHash->{regenGrub} or
        $cmdHash->{condRefresh} or
        $cmdHash->{refreshInitrd} or
        $cmdHash->{schedRebuild} or
        $cmdHash->{nameDevices} or
        $cmdHash->{installGrub})) {
      die Usage();
   }
   return $cmdHash;
}



########################################################################
#
# GenerateSimpleMap
#
#       Generates a simple device -> mapping files suitable for parsing 
#       with a shell script.  This should eventually be moved into
#       esxcfg-pciids.
#
# Results:
#       1 on success
#
# Side effects:
#       None.
#
########################################################################
sub GenerateSimpleMap 
{
   my $mapFile = shift;
   my $outputFile = shift;

   if (!open(F, "<$mapFile")) {
      LogError "Failed to open $mapFile: $!\n";
      return 0;
   }
   my @file = <F>;
   close(F);

   if (!open(F, ">$outputFile")) { 
      LogError "Failed to open $outputFile: $!\n";
      return 0;
   }

   my @supportedDevices = grep(s/^device,//, @file);

   foreach my $device (@supportedDevices) {
      chomp($device);
      my ($vend, $dvid, $type, $name, $driver, $sub) = split(/,/, $device);
      my ($subv, $subd) = split(/;/, $sub || "");

      if (defined $subv) {
         $subv =~ s/subsys_vendorid=//;
      } else {
         $subv = 0;
      }
      if (defined $subd) {
         $subd =~ s/subsys_deviceid=//;
      } else {
         $subd = 0;
      }
      printf F "%0.4x:%0.4x %0.4x:%0.4x $driver\n", hex($vend), hex($dvid), hex($subv), hex($subd);
   }
   close(F);
   return 1;
}

########################################################################
#
# main --
#
# Results:
#       Exit code 0 on success, non-zero otherwise.
#
# Side effects:
#       Everything.
#
########################################################################

sub main
{
   my $cmdHash = Init();
   exit -1 unless $cmdHash;

   my $flagfile = StdFile('rdrebuild_flag', 'esx');
   my $status = 1;
   LogDebug("$0 initialized");

   #
   # Handle basic setup and the simple case of a read-only query.
   #

   if (defined($cmdHash->{query})) {
      $status = CmdQuery($cmdHash->{query});
      VMware::CmdTool::Finish($status, 1);
   }
   
   # Setting the rebuild flag is easy. 
   if ($cmdHash->{schedRebuild}) {
      if (!open(F, ">$flagfile")) {
         LogError("Failed to created rebuild flag file.");
         exit -1;
      }
      close(F);
      exit 0;
   }

   #
   # Load esx.conf. We will hold a lock on it throughout the entire
   # update process.
   #
   # XXX: When FileSys::Lock supports readlocks, we can replace this
   #      with a more fine-grained locking scheme.
   #

   my $config = VMware::CmdTool::LoadConfig(1);
   my $checksumMgr = LoadChecksumManager();
   unless (defined $checksumMgr) {
      LogError("Could not load checksum manager.");
      $config->UnlockFile();
      VMware::CmdTool::Finish(0, 1);
   }

   if ($cmdHash->{newrd}) {
      my $bootDisk = $config->Get("/boot/virtdisk");
      if (!defined($bootDisk)) {
         LogWarn("No bootdisk defined in esx.conf.");
      } elsif (!-e $bootDisk) {
         LogWarn("Couldn't find bootdisk $bootDisk!");
      }
      $config->Set("/cos/stateless", "1");
   }

   my $updateConfig = ($cmdHash->{nameDevices} || $cmdHash->{regenGrub}); 
   my $updateInitrd = ($cmdHash->{updateBoot} || $cmdHash->{refreshInitrd});

   #
   # Check refresh: True if esx.conf has changed or rebuild flag is set. 
   #
   if ($cmdHash->{needRefresh}) {
      if ($checksumMgr->Compare() || -e $flagfile){
         print "yes\n"; 
      } else {
         print "no\n"; 
      }
      $config->UnlockFile();
      VMware::CmdTool::Finish(1, 1);
   }

   #
   # Conditional refresh: If esx.conf has changed or rebuild flag is set, 
   # regenerate the initrd. 
   #
   if ($cmdHash->{condRefresh}) {
      if ($checksumMgr->Compare() || -e $flagfile){
         $updateInitrd = 1;
      } else {
         $config->UnlockFile();
         VMware::CmdTool::Finish(1, 1);
      }
   }

   my $bootMgr = new VMware::Boot::BootManager($config->GetTree(['boot']));

   if ($status && $cmdHash->{nameDevices} && !$cmdHash->{newrd}) {
      my $bootInfo = $bootMgr->QueryBootConfig();
      my $oldCPCI = (defined $bootInfo) ? $$bootInfo[1] : undef;
      $status = UpdatePCIConfig($config,
                                $cmdHash->{cosVswitch},
                                $oldCPCI);
   }
   if ($status && $cmdHash->{regenGrub}) {
      $status = UpdateBootConfig($config, $bootMgr,
                                 $cmdHash->{rootDev},
                                 $cmdHash->{kernelAppend},
                                 $cmdHash->{newrd}, 0);
   }
   if ($status && $updateConfig) {
      $status = CommitConfig($config, $bootMgr);
   }

   #
   # We must make the initrd last because it needs to
   # contain the committed configuration changes.
   #

   if ($status && $updateInitrd) {
      GenerateSimpleMap("/etc/vmware/vmware-devices.map", "/etc/vmware/simple.map") ||
         die "Failed to generate /etc/vmware/simple.map\n";
      my $fullrebuild = unlink($flagfile);
      $status = CommitInitrd($config, $bootMgr, $checksumMgr, undef, $fullrebuild, $cmdHash->{newrd});
   }

   if ($status && $cmdHash->{installGrub}) {
      $status = $bootMgr->InstallBootloader();
   }

   $config->UnlockFile();
   VMware::CmdTool::Finish($status, 1);
}


main();

