Replace unnecessary shell out logics with python based ones (#232)

azurelinuxagent/common/osutil/default.py

@@ -16,6 +16,7 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
+import multiprocessing
 import os
 import re
 import shutil
@@ -761,19 +762,11 @@ class DefaultOSUtil(object):
         return base64.b64decode(data)
 
     def get_total_mem(self):
-        cmd = "grep MemTotal /proc/meminfo |awk '{print $2}'"
-        ret = shellutil.run_get_output(cmd)
-        if ret[0] == 0:
-            return int(ret[1])/1024
-        else:
-            raise OSUtilError("Failed to get total memory: {0}".format(ret[1]))
+        # Get total memory in bytes and divide by 1024**2 to get the valu in MB.
+        return os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / (1024**2)
 
     def get_processor_cores(self):
-        ret = shellutil.run_get_output("grep 'processor.*:' /proc/cpuinfo |wc -l")
-        if ret[0] == 0:
-            return int(ret[1])
-        else:
-            raise OSUtilError("Failed to get processor cores")
+        return multiprocessing.cpu_count()
 
     def set_admin_access_to_ip(self, dest_ip):
         #This allows root to access dest_ip

azurelinuxagent/common/protocol/metadata.py

@@ -16,6 +16,7 @@
 #
 # Requires Python 2.4+ and Openssl 1.0+
 
+import base64
 import json
 import os
 import shutil
@@ -294,8 +295,11 @@ class Certificates(object):
         p7b_file = os.path.join(conf.get_lib_dir(), P7B_FILE_NAME)
 
         # Wrapping the certificate lines.
-        b64_cmd = "echo {0} | base64 -d > {1}"
-        shellutil.run(b64_cmd.format(data, p7b_file))
+        # decode and save the result into p7b_file
+        fileStream = open(p7b_file, 'w')
+        fileStream.write(textutil.b64decode(data))
+        fileStream.close()
+
         ssl_cmd = "openssl pkcs7 -text -in {0} -inform der | grep -v '^-----' "
         ret, data = shellutil.run_get_output(ssl_cmd.format(p7b_file))
 

azurelinuxagent/common/utils/fileutil.py

@@ -21,11 +21,11 @@
 File operation util functions
 """
 
+import glob
 import os
 import re
 import shutil
 import pwd
-import tempfile
 import azurelinuxagent.common.logger as logger
 from azurelinuxagent.common.future import ustr
 import azurelinuxagent.common.utils.textutil as textutil
@@ -111,9 +111,11 @@ def chmod(path, mode):
         os.chmod(path, mode)
 
 def rm_files(*args):
-    for path in args:
-        if os.path.isfile(path):
-            os.remove(path)
+    for paths in args:
+        #Find all possible file paths
+        for path in glob.glob(paths):
+            if os.path.isfile(path):
+                os.remove(path)
 
 def rm_dirs(*args):
     """
@@ -169,3 +171,12 @@ def findstr_in_file(file_path, pattern_str):
 
     return None
 
+def get_all_files(root_path):
+    """
+    Find all files under the given root path
+    """
+    result = []
+    for root, dirs, files in os.walk(root_path):
+        result.extend([os.path.join(root, file) for file in files])
+
+    return result

azurelinuxagent/common/utils/textutil.py

@@ -278,6 +278,11 @@ def b64encode(s):
         return base64.b64encode(bytes(s, 'utf-8')).decode('utf-8')
     return base64.b64encode(s)
 
+def b64decode(s):
+    from azurelinuxagent.common.version import PY_VERSION_MAJOR
+    if PY_VERSION_MAJOR > 2:
+        return base64.b64decode(s).decode('utf-8')
+    return base64.b64decode(s)
 
 def safe_shlex_split(s):
     import shlex

azurelinuxagent/ga/exthandlers.py

@@ -21,6 +21,7 @@ import glob
 import json
 import os
 import shutil
+import stat
 import subprocess
 import time
 import zipfile
@@ -518,8 +519,10 @@ class ExtHandlerInstance(object):
         except IOError as e:
             raise ExtensionError(u"Failed to write and unzip plugin", e)
 
-        chmod = "find {0} -type f | xargs chmod u+x".format(self.get_base_dir())
-        shellutil.run(chmod)
+        #Add user execute permission to all files under the base dir
+        for file in fileutil.get_all_files(self.get_base_dir()):
+            fileutil.chmod(file, os.stat(file).st_mode | stat.S_IXUSR)
+
         self.report_event(message="Download succeeded")
 
         self.logger.info("Initialize extension directory")

azurelinuxagent/pa/deprovision/default.py

@@ -63,8 +63,8 @@ class DeprovisionHandler(object):
 
     def regen_ssh_host_key(self, warnings, actions):
         warnings.append("WARNING! All SSH host key pairs will be deleted.")
-        actions.append(DeprovisionAction(shellutil.run,
-                                         ['rm -f /etc/ssh/ssh_host_*key*']))
+        actions.append(DeprovisionAction(fileutil.rm_files,
+                                         ['/etc/ssh/ssh_host_*key*']))
 
     def stop_agent_service(self, warnings, actions):
         warnings.append("WARNING! The waagent service will be stopped.")

azurelinuxagent/pa/provision/default.py

@@ -81,7 +81,7 @@ class ProvisionHandler(object):
     def reg_ssh_host_key(self):
         keypair_type = conf.get_ssh_host_keypair_type()
         if conf.get_regenerate_ssh_host_key():
-            shellutil.run("rm -f /etc/ssh/ssh_host_*key*")
+            fileutil.rm_files("/etc/ssh/ssh_host_*key*")
             keygen_cmd = "ssh-keygen -N '' -t {0} -f /etc/ssh/ssh_host_{1}_key"
             shellutil.run(keygen_cmd.format(keypair_type, keypair_type))
         thumbprint = self.get_ssh_host_key_thumbprint(keypair_type)

tests/common/osutil/test_default.py

@@ -142,5 +142,27 @@ class TestOSUtil(AgentTestCase):
                 self.assertTrue(endpoint is not None)
                 self.assertEqual(endpoint, "second")
 
+    def test_get_total_mem(self):
+        """
+        Validate the returned value matches to the one retrieved by invoking shell command
+        """
+        cmd = "grep MemTotal /proc/meminfo |awk '{print $2}'"
+        ret = shellutil.run_get_output(cmd)
+        if ret[0] == 0:
+            self.assertEqual(int(ret[1]) / 1024, get_osutil().get_total_mem())
+        else:
+            self.fail("Cannot retrieve total memory using shell command.")
+
+    def test_get_processor_cores(self):
+        """
+        Validate the returned value matches to the one retrieved by invoking shell command
+        """
+        cmd = "grep 'processor.*:' /proc/cpuinfo |wc -l"
+        ret = shellutil.run_get_output(cmd)
+        if ret[0] == 0:
+            self.assertEqual(int(ret[1]), get_osutil().get_processor_cores())
+        else:
+            self.fail("Cannot retrieve number of process cores using shell command.")
+
 if __name__ == '__main__':
     unittest.main()

tests/ga/test_extension.py

@@ -162,7 +162,7 @@ class TestExtension(AgentTestCase):
         self.assertNotEquals(0, len(vm_status.vmAgent.extensionHandlers))
         handler_status = vm_status.vmAgent.extensionHandlers[0]
         self.assertEquals(expected_status, handler_status.status)
-        self.assertEquals("OSTCExtensions.ExampleHandlerLinux", 
+        self.assertEquals("OSTCExtensions.ExampleHandlerLinux",
                           handler_status.name)
         self.assertEquals(version, handler_status.version)
         self.assertEquals(expected_ext_count, len(handler_status.extensions))

tests/tools.py

@@ -56,6 +56,7 @@ class AgentTestCase(unittest.TestCase):
     def setUp(self):
         prefix = "{0}_".format(self.__class__.__name__)
         self.tmp_dir = tempfile.mkdtemp(prefix=prefix)
+        self.test_file = 'test_file'
         conf.get_autoupdate_enabled = Mock(return_value=True)
         conf.get_lib_dir = Mock(return_value=self.tmp_dir)
         ext_log_dir = os.path.join(self.tmp_dir, "azure")

tests/utils/test_file_util.py

@@ -24,8 +24,9 @@ from azurelinuxagent.common.future import ustr
 import azurelinuxagent.common.utils.fileutil as fileutil
 
 class TestFileOperations(AgentTestCase):
+
     def test_read_write_file(self):
-        test_file=os.path.join(self.tmp_dir, 'test_file')
+        test_file=os.path.join(self.tmp_dir, self.test_file)
         content = ustr(uuid.uuid4())
         fileutil.write_file(test_file, content)
 
@@ -34,7 +35,7 @@ class TestFileOperations(AgentTestCase):
         os.remove(test_file)
     
     def test_rw_utf8_file(self):
-        test_file=os.path.join(self.tmp_dir, 'test_file')
+        test_file=os.path.join(self.tmp_dir, self.test_file)
         content = u"\u6211"
         fileutil.write_file(test_file, content, encoding="utf-8")
 
@@ -43,14 +44,14 @@ class TestFileOperations(AgentTestCase):
         os.remove(test_file)
 
     def test_remove_bom(self):
-        test_file=os.path.join(self.tmp_dir, 'test_file')
+        test_file=os.path.join(self.tmp_dir, self.test_file)
         data = b'\xef\xbb\xbfhehe'
         fileutil.write_file(test_file, data, asbin=True)
         data = fileutil.read_file(test_file, remove_bom=True)
         self.assertNotEquals(0xbb, ord(data[0]))
    
     def test_append_file(self):
-        test_file=os.path.join(self.tmp_dir, 'test_file')
+        test_file=os.path.join(self.tmp_dir, self.test_file)
         content = ustr(uuid.uuid4())
         fileutil.append_file(test_file, content)
 
@@ -68,5 +69,53 @@ class TestFileOperations(AgentTestCase):
         filename = fileutil.base_name(filepath)
         self.assertEquals('abc', filename)
 
+    def test_remove_files(self):
+        import random
+        import string
+        import glob
+        random_word = lambda : ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5))
+
+        #Create 10 test files
+        test_file = os.path.join(self.tmp_dir, self.test_file)
+        test_file2 = os.path.join(self.tmp_dir, 'another_file')
+        test_files = [test_file + random_word() for _ in range(5)] + \
+                     [test_file2 + random_word() for _ in range(5)]
+        for file in test_files:
+            open(file, 'a').close()
+
+        #Remove files using fileutil.rm_files
+        test_file_pattern = test_file + '*'
+        test_file_pattern2 = test_file2 + '*'
+        fileutil.rm_files(test_file_pattern, test_file_pattern2)
+
+        self.assertEqual(0, len(glob.glob(os.path.join(self.tmp_dir, test_file_pattern))))
+        self.assertEqual(0, len(glob.glob(os.path.join(self.tmp_dir, test_file_pattern2))))
+
+    def test_get_all_files(self):
+        import random
+        import string
+        random_word = lambda: ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5))
+
+        # Create 10 test files at the root dir and 10 other in the sub dir
+        test_file = os.path.join(self.tmp_dir, self.test_file)
+        test_file2 = os.path.join(self.tmp_dir, 'another_file')
+        expected_files = [test_file + random_word() for _ in range(5)] + \
+                     [test_file2 + random_word() for _ in range(5)]
+
+        test_subdir = os.path.join(self.tmp_dir, 'test_dir')
+        os.mkdir(test_subdir)
+        test_file_in_subdir = os.path.join(test_subdir, self.test_file)
+        test_file_in_subdir2 = os.path.join(test_subdir, 'another_file')
+        expected_files.extend([test_file_in_subdir + random_word() for _ in range(5)] + \
+                     [test_file_in_subdir2 + random_word() for _ in range(5)])
+
+        for file in expected_files:
+            open(file, 'a').close()
+
+        # Get All files using fileutil.get_all_files
+        actual_files = fileutil.get_all_files(self.tmp_dir)
+
+        self.assertEqual(set(expected_files), set(actual_files))
+
 if __name__ == '__main__':
     unittest.main()