Repository: Azure/WALinuxAgent
Pull request #696 (2017-05-03 20:13:44)

Removed unnecessary directories and files #693
Signed-off-by: Brendan Dixon <brendandixon@me.com>


Patch:
diff --git a/azurelinuxagent/common/version.py b/azurelinuxagent/common/version.py
index 3e30f9ef..35ba2dda 100644
--- a/azurelinuxagent/common/version.py
+++ b/azurelinuxagent/common/version.py
@@ -104,6 +104,7 @@ AGENT_PKG_GLOB = "{0}-*.zip".format(AGENT_NAME)
 
 AGENT_PATTERN = "{0}-(.*)".format(AGENT_NAME)
 AGENT_NAME_PATTERN = re.compile(AGENT_PATTERN)
+AGENT_PKG_PATTERN = re.compile(AGENT_PATTERN+"\.zip")
 AGENT_DIR_PATTERN = re.compile(".*/{0}".format(AGENT_PATTERN))
 
 EXT_HANDLER_PATTERN = b".*/WALinuxAgent-(\w.\w.\w[.\w]*)-.*-run-exthandlers"
@@ -127,6 +128,13 @@ def set_current_agent():
         version = AGENT_VERSION
     return agent, FlexibleVersion(version)
 
+def is_agent_package(path):
+    path = os.path.basename(path)
+    return not re.match(AGENT_PKG_PATTERN, path) is None
+
+def is_agent_path(path):
+    path = os.path.basename(path)
+    return not re.match(AGENT_NAME_PATTERN, path) is None
 
 CURRENT_AGENT, CURRENT_VERSION = set_current_agent()
 
diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py
index 6e0a5512..e5914fd9 100644
--- a/azurelinuxagent/ga/exthandlers.py
+++ b/azurelinuxagent/ga/exthandlers.py
@@ -20,6 +20,8 @@
 import glob
 import json
 import os
+import os.path
+import re
 import shutil
 import stat
 import subprocess
@@ -31,6 +33,7 @@ import azurelinuxagent.common.logger as logger
 import azurelinuxagent.common.utils.fileutil as fileutil
 import azurelinuxagent.common.utils.restutil as restutil
 import azurelinuxagent.common.utils.shellutil as shellutil
+import azurelinuxagent.common.version as version
 
 from azurelinuxagent.common.event import add_event, WALAEventOperation
 from azurelinuxagent.common.exception import ExtensionError, ProtocolError, HttpError
@@ -55,6 +58,13 @@ VALID_EXTENSION_STATUS = ['transitioning', 'error', 'success', 'warning']
 
 VALID_HANDLER_STATUS = ['Ready', 'NotReady', "Installing", "Unresponsive"]
 
+HANDLER_PATTERN = "^([^-]+)-(\d+(?:\.\d+)*)"
+HANDLER_NAME_PATTERN = re.compile(HANDLER_PATTERN+"$", re.IGNORECASE)
+HANDLER_PKG_EXT = ".zip"
+HANDLER_PKG_PATTERN = re.compile(HANDLER_PATTERN+"\\"+HANDLER_PKG_EXT+"$",
+                                re.IGNORECASE)
+
+
 def validate_has_key(obj, key, fullname):
     if key not in obj:
         raise ExtensionError("Missing: {0}".format(fullname))
@@ -163,6 +173,7 @@ def get_exthandlers_handler():
 class ExtHandlersHandler(object):
     def __init__(self):
         self.protocol_util = get_protocol_util()
+        self.protocol = None
         self.ext_handlers = None
         self.last_etag = None
         self.log_report = False
@@ -188,10 +199,74 @@ class ExtHandlersHandler(object):
         self.last_etag = etag
 
         self.report_ext_handlers_status()
+        self.cleanup_outdated_handlers()
 
     def run_status(self):
         self.report_ext_handlers_status()
         return
+
+    def cleanup_outdated_handlers(self):
+        handlers = []
+        pkgs = []
+
+        # Build a collection of uninstalled handlers and orphaned packages
+        # Note:
+        # -- An orphaned package is one without a corresponding handler
+        #    directory
+        for item in os.listdir(conf.get_lib_dir()):
+            path = os.path.join(conf.get_lib_dir(), item)
+
+            if version.is_agent_package(path) or version.is_agent_path(path):
+                continue
+
+            if os.path.isdir(path):
+                if re.match(HANDLER_NAME_PATTERN, item) is None:
+                    continue
+                try:
+                    eh = ExtHandler()
+
+                    separator = item.rfind('-')
+
+                    eh.name = item[0:separator]
+                    eh.properties.version = str(FlexibleVersion(item[separator+1:]))
+
+                    handler = ExtHandlerInstance(eh, self.protocol)
+                except Exception as e:
+                    continue
+                if handler.get_handler_state() != ExtHandlerState.NotInstalled:
+                    continue
+                handlers.append(handler)
+
+            elif os.path.isfile(path) and \
+                    not os.path.isdir(path[0:-len(HANDLER_PKG_EXT)]):
+                if not re.match(HANDLER_PKG_PATTERN, item):
+                    continue
+                pkgs.append(path)
+
+        # Then, remove the orphaned packages
+        for pkg in pkgs:
+            try:
+                os.remove(pkg)
+                logger.verbose("Removed orphaned extension package "
+                            "{0}".format(pkg))
+            except Exception as e:
+                logger.warn("Failed to remove orphaned package: {0}".format(
+                                pkg))
+
+        # Finally, remove the directories and packages of the
+        # uninstalled handlers
+        for handler in handlers:
+            handler.rm_ext_handler_dir()
+            pkg = os.path.join(conf.get_lib_dir(),
+                    handler.get_full_name() + HANDLER_PKG_EXT)
+            if os.path.isfile(pkg):
+                try:
+                    os.remove(pkg)
+                    logger.verbose("Removed extension package "
+                                "{0}".format(pkg))
+                except Exception as e:
+                    logger.warn("Failed to remove extension package: "
+                                "{0}".format(pkg))
    
     def handle_ext_handlers(self, etag=None):
         if self.ext_handlers.extHandlers is None or \
@@ -635,6 +710,7 @@ class ExtHandlerInstance(object):
         except IOError as e:
             message = "Failed to remove extension handler directory: {0}".format(e)
             self.report_event(message=message, is_success=False)
+            self.logger.warn(message)
 
     def update(self):
         self.set_operation(WALAEventOperation.Update)


Test patch:
diff --git a/tests/data/ext/sample_ext.zip b/tests/data/ext/sample_ext-1.2.0.zip
similarity index 100%
rename from tests/data/ext/sample_ext.zip
rename to tests/data/ext/sample_ext-1.2.0.zip
diff --git a/tests/data/ext/sample_ext/HandlerManifest.json b/tests/data/ext/sample_ext-1.2.0/HandlerManifest.json
similarity index 100%
rename from tests/data/ext/sample_ext/HandlerManifest.json
rename to tests/data/ext/sample_ext-1.2.0/HandlerManifest.json
diff --git a/tests/data/ext/sample_ext/sample.py b/tests/data/ext/sample_ext-1.2.0/sample.py
similarity index 100%
rename from tests/data/ext/sample_ext/sample.py
rename to tests/data/ext/sample_ext-1.2.0/sample.py
diff --git a/tests/ga/test_extension.py b/tests/ga/test_extension.py
index 1df0a043..2a60ea36 100644
--- a/tests/ga/test_extension.py
+++ b/tests/ga/test_extension.py
@@ -15,6 +15,14 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
+import glob
+import os
+import os.path
+import shutil
+import tempfile
+import zipfile
+
+import azurelinuxagent.common.conf as conf
 import azurelinuxagent.common.utils.fileutil as fileutil
 
 from tests.protocol.mockwiredata import *
@@ -30,6 +38,103 @@ from azurelinuxagent.common.protocol.restapi import ExtHandlerStatus, \
 from azurelinuxagent.ga.exthandlers import *
 from azurelinuxagent.common.protocol.wire import WireProtocol
 
+class TestExtensionCleanup(AgentTestCase):
+    def setUp(self):
+        AgentTestCase.setUp(self)
+        self.ext_handlers = ExtHandlersHandler()
+        self.lib_dir = tempfile.mkdtemp()
+
+    def _install_handlers(self, start=0, count=1,
+                        handler_state=ExtHandlerState.Installed):
+        src = os.path.join(data_dir, "ext", "sample_ext-1.2.0.zip")
+        version = FlexibleVersion("1.2.0")
+        version += start - version.patch
+
+        for i in range(start, start+count):
+            eh = ExtHandler()
+            eh.name = "sample_ext"
+            eh.properties.version = str(version)
+            handler = ExtHandlerInstance(eh, "unused")
+
+            dst = os.path.join(self.lib_dir,
+                    handler.get_full_name()+HANDLER_PKG_EXT)
+            shutil.copy(src, dst)
+
+            if not handler_state is None:
+                zipfile.ZipFile(dst).extractall(handler.get_base_dir())
+                handler.set_handler_state(handler_state)
+
+            version += 1
+    
+    def _count_packages(self):
+        return len(glob.glob(os.path.join(self.lib_dir, "*.zip")))
+
+    def _count_installed(self):
+        paths = os.listdir(self.lib_dir)
+        paths = [os.path.join(self.lib_dir, p) for p in paths]
+        return len([p for p in paths
+                        if os.path.isdir(p) and self._is_installed(p)])
+
+    def _count_uninstalled(self):
+        paths = os.listdir(self.lib_dir)
+        paths = [os.path.join(self.lib_dir, p) for p in paths]
+        return len([p for p in paths
+                        if os.path.isdir(p) and not self._is_installed(p)])
+
+    def _is_installed(self, path):
+        path = os.path.join(path, 'config', 'HandlerState')
+        return fileutil.read_file(path) != "NotInstalled"
+
+    @patch("azurelinuxagent.common.conf.get_lib_dir")
+    def test_cleanup_leaves_installed_extensions(self, mock_conf):
+        mock_conf.return_value = self.lib_dir
+
+        self._install_handlers(start=0, count=5, handler_state=ExtHandlerState.Installed)
+        self._install_handlers(start=5, count=5, handler_state=ExtHandlerState.Enabled)
+
+        self.assertEqual(self._count_packages(), 10)
+        self.assertEqual(self._count_installed(), 10)
+
+        self.ext_handlers.cleanup_outdated_handlers()
+
+        self.assertEqual(self._count_packages(), 10)
+        self.assertEqual(self._count_installed(), 10)
+        self.assertEqual(self._count_uninstalled(), 0)
+
+    @patch("azurelinuxagent.common.conf.get_lib_dir")
+    def test_cleanup_removes_uninstalled_extensions(self, mock_conf):
+        mock_conf.return_value = self.lib_dir
+
+        self._install_handlers(start=0, count=5, handler_state=ExtHandlerState.Installed)
+        self._install_handlers(start=5, count=5, handler_state=ExtHandlerState.NotInstalled)
+
+        self.assertEqual(self._count_packages(), 10)
+        self.assertEqual(self._count_installed(), 5)
+        self.assertEqual(self._count_uninstalled(), 5)
+
+        self.ext_handlers.cleanup_outdated_handlers()
+
+        self.assertEqual(self._count_packages(), 5)
+        self.assertEqual(self._count_installed(), 5)
+        self.assertEqual(self._count_uninstalled(), 0)
+
+    @patch("azurelinuxagent.common.conf.get_lib_dir")
+    def test_cleanup_removes_orphaned_packages(self, mock_conf):
+        mock_conf.return_value = self.lib_dir
+
+        self._install_handlers(start=0, count=5, handler_state=ExtHandlerState.Installed)
+        self._install_handlers(start=5, count=5, handler_state=None)
+
+        self.assertEqual(self._count_packages(), 10)
+        self.assertEqual(self._count_installed(), 5)
+        self.assertEqual(self._count_uninstalled(), 0)
+
+        self.ext_handlers.cleanup_outdated_handlers()
+
+        self.assertEqual(self._count_packages(), 5)
+        self.assertEqual(self._count_installed(), 5)
+        self.assertEqual(self._count_uninstalled(), 0)
+
 class TestHandlerStateMigration(AgentTestCase):
     def setUp(self):
         AgentTestCase.setUp(self)
diff --git a/tests/protocol/mockwiredata.py b/tests/protocol/mockwiredata.py
index c789de5c..4e45623c 100644
--- a/tests/protocol/mockwiredata.py
+++ b/tests/protocol/mockwiredata.py
@@ -30,7 +30,7 @@ DATA_FILE = {
         "ga_manifest" : "wire/ga_manifest.xml",
         "trans_prv": "wire/trans_prv",
         "trans_cert": "wire/trans_cert",
-        "test_ext": "ext/sample_ext.zip"
+        "test_ext": "ext/sample_ext-1.2.0.zip"
 }
 
 DATA_FILE_NO_EXT = DATA_FILE.copy()