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.

8 files changed, 459 insertions, 42 deletions

diff --git a/cloudinit/sources/DataSourceOVF.py b/cloudinit/sources/DataSourceOVF.py
index 6ac621f2..6e62f984 100644
--- a/cloudinit/sources/DataSourceOVF.py
+++ b/cloudinit/sources/DataSourceOVF.py
@@ -21,6 +21,8 @@ from cloudinit import util
 
 from cloudinit.sources.helpers.vmware.imc.config \
     import Config
+from cloudinit.sources.helpers.vmware.imc.config_custom_script \
+    import PreCustomScript, PostCustomScript
 from cloudinit.sources.helpers.vmware.imc.config_file \
     import ConfigFile
 from cloudinit.sources.helpers.vmware.imc.config_nic \
@@ -30,7 +32,7 @@ from cloudinit.sources.helpers.vmware.imc.config_passwd \
 from cloudinit.sources.helpers.vmware.imc.guestcust_error \
     import GuestCustErrorEnum
 from cloudinit.sources.helpers.vmware.imc.guestcust_event \
-    import GuestCustEventEnum
+    import GuestCustEventEnum as GuestCustEvent
 from cloudinit.sources.helpers.vmware.imc.guestcust_state \
     import GuestCustStateEnum
 from cloudinit.sources.helpers.vmware.imc.guestcust_util import (
@@ -127,17 +129,31 @@ class DataSourceOVF(sources.DataSource):
                 self._vmware_cust_conf = Config(cf)
                 (md, ud, cfg) = read_vmware_imc(self._vmware_cust_conf)
                 self._vmware_nics_to_enable = get_nics_to_enable(nicspath)
-                markerid = self._vmware_cust_conf.marker_id
-                markerexists = check_marker_exists(markerid)
+                imcdirpath = os.path.dirname(vmwareImcConfigFilePath)
+                product_marker = self._vmware_cust_conf.marker_id
+                hasmarkerfile = check_marker_exists(
+                    product_marker, os.path.join(self.paths.cloud_dir, 'data'))
+                special_customization = product_marker and not hasmarkerfile
+                customscript = self._vmware_cust_conf.custom_script_name
             except Exception as e:
-                LOG.debug("Error parsing the customization Config File")
-                LOG.exception(e)
-                set_customization_status(
-                    GuestCustStateEnum.GUESTCUST_STATE_RUNNING,
-                    GuestCustEventEnum.GUESTCUST_EVENT_CUSTOMIZE_FAILED)
-                raise e
-            finally:
-                util.del_dir(os.path.dirname(vmwareImcConfigFilePath))
+                _raise_error_status(
+                    "Error parsing the customization Config File",
+                    e,
+                    GuestCustEvent.GUESTCUST_EVENT_CUSTOMIZE_FAILED,
+                    vmwareImcConfigFilePath)
+
+            if special_customization:
+                if customscript:
+                    try:
+                        precust = PreCustomScript(customscript, imcdirpath)
+                        precust.execute()
+                    except Exception as e:
+                        _raise_error_status(
+                            "Error executing pre-customization script",
+                            e,
+                            GuestCustEvent.GUESTCUST_EVENT_CUSTOMIZE_FAILED,
+                            vmwareImcConfigFilePath)
+
             try:
                 LOG.debug("Preparing the Network configuration")
                 self._network_config = get_network_config_from_conf(
@@ -146,13 +162,13 @@ class DataSourceOVF(sources.DataSource):
                     True,
                     self.distro.osfamily)
             except Exception as e:
-                LOG.exception(e)
-                set_customization_status(
-                    GuestCustStateEnum.GUESTCUST_STATE_RUNNING,
-                    GuestCustEventEnum.GUESTCUST_EVENT_NETWORK_SETUP_FAILED)
-                raise e
+                _raise_error_status(
+                    "Error preparing Network Configuration",
+                    e,
+                    GuestCustEvent.GUESTCUST_EVENT_NETWORK_SETUP_FAILED,
+                    vmwareImcConfigFilePath)
 
-            if markerid and not markerexists:
+            if special_customization:
                 LOG.debug("Applying password customization")
                 pwdConfigurator = PasswordConfigurator()
                 adminpwd = self._vmware_cust_conf.admin_password
@@ -164,27 +180,41 @@ class DataSourceOVF(sources.DataSource):
                     else:
                         LOG.debug("Changing password is not needed")
                 except Exception as e:
-                    LOG.debug("Error applying Password Configuration: %s", e)
-                    set_customization_status(
-                        GuestCustStateEnum.GUESTCUST_STATE_RUNNING,
-                        GuestCustEventEnum.GUESTCUST_EVENT_CUSTOMIZE_FAILED)
-                    return False
-            if markerid:
-                LOG.debug("Handle marker creation")
+                    _raise_error_status(
+                        "Error applying Password Configuration",
+                        e,
+                        GuestCustEvent.GUESTCUST_EVENT_CUSTOMIZE_FAILED,
+                        vmwareImcConfigFilePath)
+
+                if customscript:
+                    try:
+                        postcust = PostCustomScript(customscript, imcdirpath)
+                        postcust.execute()
+                    except Exception as e:
+                        _raise_error_status(
+                            "Error executing post-customization script",
+                            e,
+                            GuestCustEvent.GUESTCUST_EVENT_CUSTOMIZE_FAILED,
+                            vmwareImcConfigFilePath)
+
+            if product_marker:
                 try:
-                    setup_marker_files(markerid)
+                    setup_marker_files(
+                        product_marker,
+                        os.path.join(self.paths.cloud_dir, 'data'))
                 except Exception as e:
-                    LOG.debug("Error creating marker files: %s", e)
-                    set_customization_status(
-                        GuestCustStateEnum.GUESTCUST_STATE_RUNNING,
-                        GuestCustEventEnum.GUESTCUST_EVENT_CUSTOMIZE_FAILED)
-                    return False
+                    _raise_error_status(
+                        "Error creating marker files",
+                        e,
+                        GuestCustEvent.GUESTCUST_EVENT_CUSTOMIZE_FAILED,
+                        vmwareImcConfigFilePath)
 
             self._vmware_cust_found = True
             found.append('vmware-tools')
 
             # TODO: Need to set the status to DONE only when the
             # customization is done successfully.
+            util.del_dir(os.path.dirname(vmwareImcConfigFilePath))
             enable_nics(self._vmware_nics_to_enable)
             set_customization_status(
                 GuestCustStateEnum.GUESTCUST_STATE_DONE,
@@ -539,31 +569,52 @@ def get_datasource_list(depends):
 
 
 # To check if marker file exists
-def check_marker_exists(markerid):
+def check_marker_exists(markerid, marker_dir):
     """
     Check the existence of a marker file.
     Presence of marker file determines whether a certain code path is to be
     executed. It is needed for partial guest customization in VMware.
+    @param markerid: is an unique string representing a particular product
+                     marker.
+    @param: marker_dir: The directory in which markers exist.
     """
     if not markerid:
         return False
-    markerfile = "/.markerfile-" + markerid
+    markerfile = os.path.join(marker_dir, ".markerfile-" + markerid + ".txt")
     if os.path.exists(markerfile):
         return True
     return False
 
 
 # Create a marker file
-def setup_marker_files(markerid):
+def setup_marker_files(markerid, marker_dir):
     """
     Create a new marker file.
     Marker files are unique to a full customization workflow in VMware
     environment.
+    @param markerid: is an unique string representing a particular product
+                     marker.
+    @param: marker_dir: The directory in which markers exist.
+
     """
-    if not markerid:
-        return
-    markerfile = "/.markerfile-" + markerid
-    util.del_file("/.markerfile-*.txt")
+    LOG.debug("Handle marker creation")
+    markerfile = os.path.join(marker_dir, ".markerfile-" + markerid + ".txt")
+    for fname in os.listdir(marker_dir):
+        if fname.startswith(".markerfile"):
+            util.del_file(os.path.join(marker_dir, fname))
     open(markerfile, 'w').close()
 
+
+def _raise_error_status(prefix, error, event, config_file):
+    """
+    Raise error and send customization status to the underlying VMware
+    Virtualization Platform. Also, cleanup the imc directory.
+    """
+    LOG.debug('%s: %s', prefix, error)
+    set_customization_status(
+        GuestCustStateEnum.GUESTCUST_STATE_RUNNING,
+        event)
+    util.del_dir(os.path.dirname(config_file))
+    raise error
+
 # vi: ts=4 expandtab
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:
diff --git a/tests/unittests/test_datasource/test_ovf.py b/tests/unittests/test_datasource/test_ovf.py
index 700da86c..fc4eb36e 100644
--- a/tests/unittests/test_datasource/test_ovf.py
+++ b/tests/unittests/test_datasource/test_ovf.py
@@ -5,11 +5,17 @@
 # This file is part of cloud-init. See LICENSE file for license information.
 
 import base64
-from collections import OrderedDict
+import os
 
-from cloudinit.tests import helpers as test_helpers
+from collections import OrderedDict
+from textwrap import dedent
 
+from cloudinit import util
+from cloudinit.tests.helpers import CiTestCase, wrap_and_call
+from cloudinit.helpers import Paths
 from cloudinit.sources import DataSourceOVF as dsovf
+from cloudinit.sources.helpers.vmware.imc.config_custom_script import (
+    CustomScriptNotFound)
 
 OVF_ENV_CONTENT = """<?xml version="1.0" encoding="UTF-8"?>
 <Environment xmlns="http://schemas.dmtf.org/ovf/environment/1"
@@ -42,7 +48,7 @@ def fill_properties(props, template=OVF_ENV_CONTENT):
     return template.format(properties=properties)
 
 
-class TestReadOvfEnv(test_helpers.TestCase):
+class TestReadOvfEnv(CiTestCase):
     def test_with_b64_userdata(self):
         user_data = "#!/bin/sh\necho hello world\n"
         user_data_b64 = base64.b64encode(user_data.encode()).decode()
@@ -72,7 +78,104 @@ class TestReadOvfEnv(test_helpers.TestCase):
         self.assertIsNone(ud)
 
 
-class TestTransportIso9660(test_helpers.CiTestCase):
+class TestMarkerFiles(CiTestCase):
+
+    def setUp(self):
+        super(TestMarkerFiles, self).setUp()
+        self.tdir = self.tmp_dir()
+
+    def test_false_when_markerid_none(self):
+        """Return False when markerid provided is None."""
+        self.assertFalse(
+            dsovf.check_marker_exists(markerid=None, marker_dir=self.tdir))
+
+    def test_markerid_file_exist(self):
+        """Return False when markerid file path does not exist,
+        True otherwise."""
+        self.assertFalse(
+            dsovf.check_marker_exists('123', self.tdir))
+
+        marker_file = self.tmp_path('.markerfile-123.txt', self.tdir)
+        util.write_file(marker_file, '')
+        self.assertTrue(
+            dsovf.check_marker_exists('123', self.tdir)
+        )
+
+    def test_marker_file_setup(self):
+        """Test creation of marker files."""
+        markerfilepath = self.tmp_path('.markerfile-hi.txt', self.tdir)
+        self.assertFalse(os.path.exists(markerfilepath))
+        dsovf.setup_marker_files(markerid='hi', marker_dir=self.tdir)
+        self.assertTrue(os.path.exists(markerfilepath))
+
+
+class TestDatasourceOVF(CiTestCase):
+
+    with_logs = True
+
+    def setUp(self):
+        super(TestDatasourceOVF, self).setUp()
+        self.datasource = dsovf.DataSourceOVF
+        self.tdir = self.tmp_dir()
+
+    def test_get_data_false_on_none_dmi_data(self):
+        """When dmi for system-product-name is None, get_data returns False."""
+        paths = Paths({'seed_dir': self.tdir})
+        ds = self.datasource(sys_cfg={}, distro={}, paths=paths)
+        retcode = wrap_and_call(
+            'cloudinit.sources.DataSourceOVF',
+            {'util.read_dmi_data': None},
+            ds.get_data)
+        self.assertFalse(retcode, 'Expected False return from ds.get_data')
+        self.assertIn(
+            'DEBUG: No system-product-name found', self.logs.getvalue())
+
+    def test_get_data_no_vmware_customization_disabled(self):
+        """When vmware customization is disabled via sys_cfg log a message."""
+        paths = Paths({'seed_dir': self.tdir})
+        ds = self.datasource(
+            sys_cfg={'disable_vmware_customization': True}, distro={},
+            paths=paths)
+        retcode = wrap_and_call(
+            'cloudinit.sources.DataSourceOVF',
+            {'util.read_dmi_data': 'vmware'},
+            ds.get_data)
+        self.assertFalse(retcode, 'Expected False return from ds.get_data')
+        self.assertIn(
+            'DEBUG: Customization for VMware platform is disabled.',
+            self.logs.getvalue())
+
+    def test_get_data_vmware_customization_disabled(self):
+        """When cloud-init workflow for vmware is enabled via sys_cfg log a
+        message.
+        """
+        paths = Paths({'seed_dir': self.tdir})
+        ds = self.datasource(
+            sys_cfg={'disable_vmware_customization': False}, distro={},
+            paths=paths)
+        conf_file = self.tmp_path('test-cust', self.tdir)
+        conf_content = dedent("""\
+            [CUSTOM-SCRIPT]
+            SCRIPT-NAME = test-script
+            [MISC]
+            MARKER-ID = 12345345
+            """)
+        util.write_file(conf_file, conf_content)
+        with self.assertRaises(CustomScriptNotFound) as context:
+            wrap_and_call(
+                'cloudinit.sources.DataSourceOVF',
+                {'util.read_dmi_data': 'vmware',
+                 'util.del_dir': True,
+                 'search_file': self.tdir,
+                 'wait_for_imc_cfg_file': conf_file,
+                 'get_nics_to_enable': ''},
+                ds.get_data)
+        customscript = self.tmp_path('test-script', self.tdir)
+        self.assertIn('Script %s not found!!' % customscript,
+                      str(context.exception))
+
+
+class TestTransportIso9660(CiTestCase):
 
     def setUp(self):
         super(TestTransportIso9660, self).setUp()
diff --git a/tests/unittests/test_vmware/__init__.py b/tests/unittests/test_vmware/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/tests/unittests/test_vmware/__init__.py
diff --git a/tests/unittests/test_vmware/test_custom_script.py b/tests/unittests/test_vmware/test_custom_script.py
new file mode 100644
index 00000000..2d9519b0
--- /dev/null
+++ b/tests/unittests/test_vmware/test_custom_script.py
@@ -0,0 +1,99 @@
+# Copyright (C) 2015 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.
+
+from cloudinit import util
+from cloudinit.sources.helpers.vmware.imc.config_custom_script import (
+    CustomScriptConstant,
+    CustomScriptNotFound,
+    PreCustomScript,
+    PostCustomScript,
+)
+from cloudinit.tests.helpers import CiTestCase, mock
+
+
+class TestVmwareCustomScript(CiTestCase):
+    def setUp(self):
+        self.tmpDir = self.tmp_dir()
+
+    def test_prepare_custom_script(self):
+        """
+        This test is designed to verify the behavior based on the presence of
+        custom script. Mainly needed for scenario where a custom script is
+        expected, but was not properly copied. "CustomScriptNotFound" exception
+        is raised in such cases.
+        """
+        # Custom script does not exist.
+        preCust = PreCustomScript("random-vmw-test", self.tmpDir)
+        self.assertEqual("random-vmw-test", preCust.scriptname)
+        self.assertEqual(self.tmpDir, preCust.directory)
+        self.assertEqual(self.tmp_path("random-vmw-test", self.tmpDir),
+                         preCust.scriptpath)
+        with self.assertRaises(CustomScriptNotFound):
+            preCust.prepare_script()
+
+        # Custom script exists.
+        custScript = self.tmp_path("test-cust", self.tmpDir)
+        util.write_file(custScript, "test-CR-strip/r/r")
+        postCust = PostCustomScript("test-cust", self.tmpDir)
+        self.assertEqual("test-cust", postCust.scriptname)
+        self.assertEqual(self.tmpDir, postCust.directory)
+        self.assertEqual(custScript, postCust.scriptpath)
+        self.assertFalse(postCust.postreboot)
+        postCust.prepare_script()
+        # Check if all carraige returns are stripped from script.
+        self.assertFalse("/r" in custScript)
+
+    def test_rc_local_exists(self):
+        """
+        This test is designed to verify the different scenarios associated
+        with the presence of rclocal.
+        """
+        # test when rc local does not exist
+        postCust = PostCustomScript("test-cust", self.tmpDir)
+        with mock.patch.object(CustomScriptConstant, "RC_LOCAL", "/no/path"):
+            rclocal = postCust.find_rc_local()
+            self.assertEqual("", rclocal)
+
+        # test when rc local exists
+        rclocalFile = self.tmp_path("vmware-rclocal", self.tmpDir)
+        util.write_file(rclocalFile, "# Run post-reboot guest customization",
+                        omode="w")
+        with mock.patch.object(CustomScriptConstant, "RC_LOCAL", rclocalFile):
+            rclocal = postCust.find_rc_local()
+            self.assertEqual(rclocalFile, rclocal)
+            self.assertTrue(postCust.has_previous_agent, rclocal)
+
+        # test when rc local is a symlink
+        rclocalLink = self.tmp_path("dummy-rclocal-link", self.tmpDir)
+        util.sym_link(rclocalFile, rclocalLink, True)
+        with mock.patch.object(CustomScriptConstant, "RC_LOCAL", rclocalLink):
+            rclocal = postCust.find_rc_local()
+            self.assertEqual(rclocalFile, rclocal)
+
+    def test_execute_post_cust(self):
+        """
+        This test is to identify if rclocal was properly populated to be
+        run after reboot.
+        """
+        customscript = self.tmp_path("vmware-post-cust-script", self.tmpDir)
+        rclocal = self.tmp_path("vmware-rclocal", self.tmpDir)
+        # Create a temporary rclocal file
+        open(customscript, "w")
+        util.write_file(rclocal, "tests\nexit 0", omode="w")
+        postCust = PostCustomScript("vmware-post-cust-script", self.tmpDir)
+        with mock.patch.object(CustomScriptConstant, "RC_LOCAL", rclocal):
+            # Test that guest customization agent is not installed initially.
+            self.assertFalse(postCust.postreboot)
+            self.assertIs(postCust.has_previous_agent(rclocal), False)
+            postCust.install_agent()
+
+            # Assert rclocal has been modified to have guest customization
+            # agent.
+            self.assertTrue(postCust.postreboot)
+            self.assertTrue(postCust.has_previous_agent, rclocal)
+
+# vi: ts=4 expandtab
diff --git a/tests/unittests/test_vmware_config_file.py b/tests/unittests/test_vmware_config_file.py
index 0f8cda95..036f6879 100644
--- a/tests/unittests/test_vmware_config_file.py
+++ b/tests/unittests/test_vmware_config_file.py
@@ -335,5 +335,12 @@ class TestVmwareConfigFile(CiTestCase):
         self.assertEqual('255.255.0.0', subnet.get('netmask'),
                          'Subnet netmask')
 
+    def test_custom_script(self):
+        cf = ConfigFile("tests/data/vmware/cust-dhcp-2nic.cfg")
+        conf = Config(cf)
+        self.assertIsNone(conf.custom_script_name)
+        cf._insertKey("CUSTOM-SCRIPT|SCRIPT-NAME", "test-script")
+        conf = Config(cf)
+        self.assertEqual("test-script", conf.custom_script_name)
 
 # vi: ts=4 expandtab