VMware: Support for user provided pre and post-customization scripts

In the VMware customization workflow, we have some options for the user
to upload scripts for additional customization. Based on user request,
those custom scripts can be either run before regular customization or
after. For post customization scripts, we decide whether to run the scripts
just after customization or post system reboot.

3 files changed, 158 insertions, 1 deletions

diff --git a/cloudinit/sources/helpers/vmware/imc/config.py b/cloudinit/sources/helpers/vmware/imc/config.py
index 49d441db..2eaeff34 100644
--- a/cloudinit/sources/helpers/vmware/imc/config.py
+++ b/cloudinit/sources/helpers/vmware/imc/config.py
@@ -100,4 +100,8 @@ class Config(object):
         """Returns marker id."""
         return self._configFile.get(Config.MARKERID, None)
 
+    @property
+    def custom_script_name(self):
+        """Return the name of custom (pre/post) script."""
+        return self._configFile.get(Config.CUSTOM_SCRIPT, None)
 # vi: ts=4 expandtab
diff --git a/cloudinit/sources/helpers/vmware/imc/config_custom_script.py b/cloudinit/sources/helpers/vmware/imc/config_custom_script.py
new file mode 100644
index 00000000..a7d4ad91
--- /dev/null
+++ b/cloudinit/sources/helpers/vmware/imc/config_custom_script.py
@@ -0,0 +1,153 @@
+# Copyright (C) 2017 Canonical Ltd.
+# Copyright (C) 2017 VMware Inc.
+#
+# Author: Maitreyee Saikia <msaikia@vmware.com>
+#
+# This file is part of cloud-init. See LICENSE file for license information.
+
+import logging
+import os
+import stat
+from textwrap import dedent
+
+from cloudinit import util
+
+LOG = logging.getLogger(__name__)
+
+
+class CustomScriptNotFound(Exception):
+    pass
+
+
+class CustomScriptConstant(object):
+    RC_LOCAL = "/etc/rc.local"
+    POST_CUST_TMP_DIR = "/root/.customization"
+    POST_CUST_RUN_SCRIPT_NAME = "post-customize-guest.sh"
+    POST_CUST_RUN_SCRIPT = os.path.join(POST_CUST_TMP_DIR,
+                                        POST_CUST_RUN_SCRIPT_NAME)
+    POST_REBOOT_PENDING_MARKER = "/.guest-customization-post-reboot-pending"
+
+
+class RunCustomScript(object):
+    def __init__(self, scriptname, directory):
+        self.scriptname = scriptname
+        self.directory = directory
+        self.scriptpath = os.path.join(directory, scriptname)
+
+    def prepare_script(self):
+        if not os.path.exists(self.scriptpath):
+            raise CustomScriptNotFound("Script %s not found!! "
+                                       "Cannot execute custom script!"
+                                       % self.scriptpath)
+        # Strip any CR characters from the decoded script
+        util.load_file(self.scriptpath).replace("\r", "")
+        st = os.stat(self.scriptpath)
+        os.chmod(self.scriptpath, st.st_mode | stat.S_IEXEC)
+
+
+class PreCustomScript(RunCustomScript):
+    def execute(self):
+        """Executing custom script with precustomization argument."""
+        LOG.debug("Executing pre-customization script")
+        self.prepare_script()
+        util.subp(["/bin/sh", self.scriptpath, "precustomization"])
+
+
+class PostCustomScript(RunCustomScript):
+    def __init__(self, scriptname, directory):
+        super(PostCustomScript, self).__init__(scriptname, directory)
+        # Determine when to run custom script. When postreboot is True,
+        # the user uploaded script will run as part of rc.local after
+        # the machine reboots. This is determined by presence of rclocal.
+        # When postreboot is False, script will run as part of cloud-init.
+        self.postreboot = False
+
+    def _install_post_reboot_agent(self, rclocal):
+        """
+        Install post-reboot agent for running custom script after reboot.
+        As part of this process, we are editing the rclocal file to run a
+        VMware script, which in turn is resposible for handling the user
+        script.
+        @param: path to rc local.
+        """
+        LOG.debug("Installing post-reboot customization from %s to %s",
+                  self.directory, rclocal)
+        if not self.has_previous_agent(rclocal):
+            LOG.info("Adding post-reboot customization agent to rc.local")
+            new_content = dedent("""
+                # Run post-reboot guest customization
+                /bin/sh %s
+                exit 0
+                """) % CustomScriptConstant.POST_CUST_RUN_SCRIPT
+            existing_rclocal = util.load_file(rclocal).replace('exit 0\n', '')
+            st = os.stat(rclocal)
+            # "x" flag should be set
+            mode = st.st_mode | stat.S_IEXEC
+            util.write_file(rclocal, existing_rclocal + new_content, mode)
+
+        else:
+            # We don't need to update rclocal file everytime a customization
+            # is requested. It just needs to be done for the first time.
+            LOG.info("Post-reboot guest customization agent is already "
+                     "registered in rc.local")
+        LOG.debug("Installing post-reboot customization agent finished: %s",
+                  self.postreboot)
+
+    def has_previous_agent(self, rclocal):
+        searchstring = "# Run post-reboot guest customization"
+        if searchstring in open(rclocal).read():
+            return True
+        return False
+
+    def find_rc_local(self):
+        """
+        Determine if rc local is present.
+        """
+        rclocal = ""
+        if os.path.exists(CustomScriptConstant.RC_LOCAL):
+            LOG.debug("rc.local detected.")
+            # resolving in case of symlink
+            rclocal = os.path.realpath(CustomScriptConstant.RC_LOCAL)
+            LOG.debug("rc.local resolved to %s", rclocal)
+        else:
+            LOG.warning("Can't find rc.local, post-customization "
+                        "will be run before reboot")
+        return rclocal
+
+    def install_agent(self):
+        rclocal = self.find_rc_local()
+        if rclocal:
+            self._install_post_reboot_agent(rclocal)
+            self.postreboot = True
+
+    def execute(self):
+        """
+        This method executes post-customization script before or after reboot
+        based on the presence of rc local.
+        """
+        self.prepare_script()
+        self.install_agent()
+        if not self.postreboot:
+            LOG.warning("Executing post-customization script inline")
+            util.subp(["/bin/sh", self.scriptpath, "postcustomization"])
+        else:
+            LOG.debug("Scheduling custom script to run post reboot")
+            if not os.path.isdir(CustomScriptConstant.POST_CUST_TMP_DIR):
+                os.mkdir(CustomScriptConstant.POST_CUST_TMP_DIR)
+            # Script "post-customize-guest.sh" and user uploaded script are
+            # are present in the same directory and needs to copied to a temp
+            # directory to be executed post reboot. User uploaded script is
+            # saved as customize.sh in the temp directory.
+            # post-customize-guest.sh excutes customize.sh after reboot.
+            LOG.debug("Copying post-customization script")
+            util.copy(self.scriptpath,
+                      CustomScriptConstant.POST_CUST_TMP_DIR + "/customize.sh")
+            LOG.debug("Copying script to run post-customization script")
+            util.copy(
+                os.path.join(self.directory,
+                             CustomScriptConstant.POST_CUST_RUN_SCRIPT_NAME),
+                CustomScriptConstant.POST_CUST_RUN_SCRIPT)
+            LOG.info("Creating post-reboot pending marker")
+            util.ensure_file(CustomScriptConstant.POST_REBOOT_PENDING_MARKER)
+
+# vi: ts=4 expandtab
diff --git a/cloudinit/sources/helpers/vmware/imc/config_nic.py b/cloudinit/sources/helpers/vmware/imc/config_nic.py
index 2fb07c59..2d8900e2 100644
--- a/cloudinit/sources/helpers/vmware/imc/config_nic.py
+++ b/cloudinit/sources/helpers/vmware/imc/config_nic.py
@@ -161,7 +161,7 @@ class NicConfigurator(object):
         if nic.primary and v4.gateways:
             self.ipv4PrimaryGateway = v4.gateways[0]
             subnet.update({'gateway': self.ipv4PrimaryGateway})
-            return [subnet]
+            return ([subnet], route_list)
 
         # Add routes if there is no primary nic
         if not self._primaryNic: